[compiler] Infer phi types, extend mutable ranges to account for Store effects

compiler/packages/babel-plugin-react-compiler/src/Inference/InferMutableRanges.ts

@@ -50,8 +50,38 @@ export function inferMutableRanges(ir: HIRFunction): void {
   // Re-infer mutable ranges for all values
   inferMutableLifetimes(ir, true);
 
-  // Re-infer mutable ranges for aliases
-  inferMutableRangesForAlias(ir, aliases);
+  /**
+   * The second inferMutableLifetimes() call updates mutable ranges
+   * of values to account for Store effects. Now we need to update
+   * all aliases of such values to extend their ranges as well. Note
+   * that the store only mutates the the directly aliased value and
+   * not any of its inner captured references. For example:
+   *
+   * ```
+   * let y;
+   * if (cond) {
+   *   y = [];
+   * } else {
+   *   y = [{}];
+   * }
+   * y.push(z);
+   * ```
+   *
+   * The Store effect from the `y.push` modifies the values that `y`
+   * directly aliases - the two arrays from the if/else branches -
+   * but does not modify values that `y` "contains" such as the
+   * object literal or `z`.
+   */
+  prevAliases = aliases.canonicalize();
+  while (true) {
+    inferMutableRangesForAlias(ir, aliases);
+    inferAliasForPhis(ir, aliases);
+    const nextAliases = aliases.canonicalize();
+    if (areEqualMaps(prevAliases, nextAliases)) {
+      break;
+    }
+    prevAliases = nextAliases;
+  }
 }
 
 function areEqualMaps<T>(a: Map<T, T>, b: Map<T, T>): boolean {

compiler/packages/babel-plugin-react-compiler/src/TypeInference/InferTypes.ts

@@ -483,30 +483,135 @@ class Unifier {
     }
 
     if (type.kind === 'Phi') {
-      const operands = new Set(type.operands.map(i => this.get(i).kind));
-
-      CompilerError.invariant(operands.size > 0, {
+      CompilerError.invariant(type.operands.length > 0, {
         reason: 'there should be at least one operand',
         description: null,
         loc: null,
         suggestions: null,
       });
-      const kind = operands.values().next().value;
 
-      // there's only one unique type and it's not a type var
-      if (operands.size === 1 && kind !== 'Type') {
-        this.unify(v, type.operands[0]);
+      let candidateType: Type | null = null;
+      for (const operand of type.operands) {
+        const resolved = this.get(operand);
+        if (candidateType === null) {
+          candidateType = resolved;
+        } else if (!typeEquals(resolved, candidateType)) {
+          candidateType = null;
+          break;
+        } // else same type, continue
+      }
+
+      if (candidateType !== null) {
+        this.unify(v, candidateType);
         return;
       }
     }
 
     if (this.occursCheck(v, type)) {
+      const resolvedType = this.tryResolveType(v, type);
+      if (resolvedType !== null) {
+        this.substitutions.set(v.id, resolvedType);
+        return;
+      }
       throw new Error('cycle detected');
     }
 
     this.substitutions.set(v.id, type);
   }
 
+  tryResolveType(v: TypeVar, type: Type): Type | null {
+    switch (type.kind) {
+      case 'Phi': {
+        /**
+         * Resolve the type of the phi by recursively removing `v` as an operand.
+         * For example we can end up with types like this:
+         *
+         * v = Phi [
+         *   T1
+         *   T2
+         *   Phi [
+         *     T3
+         *     Phi [
+         *       T4
+         *       v <-- cycle!
+         *     ]
+         *   ]
+         * ]
+         *
+         * By recursively removing `v`, we end up with:
+         *
+         * v = Phi [
+         *   T1
+         *   T2
+         *   Phi [
+         *     T3
+         *     Phi [
+         *       T4
+         *     ]
+         *   ]
+         * ]
+         *
+         * Which avoids the cycle
+         */
+        const operands = [];
+        for (const operand of type.operands) {
+          if (operand.kind === 'Type' && operand.id === v.id) {
+            continue;
+          }
+          const resolved = this.tryResolveType(v, operand);
+          if (resolved === null) {
+            return null;
+          }
+          operands.push(resolved);
+        }
+        return {kind: 'Phi', operands};
+      }
+      case 'Type': {
+        const substitution = this.get(type);
+        if (substitution !== type) {
+          const resolved = this.tryResolveType(v, substitution);
+          if (resolved !== null) {
+            this.substitutions.set(type.id, resolved);
+          }
+          return resolved;
+        }
+        return type;
+      }
+      case 'Property': {
+        const objectType = this.tryResolveType(v, this.get(type.objectType));
+        if (objectType === null) {
+          return null;
+        }
+        return {
+          kind: 'Property',
+          objectName: type.objectName,
+          objectType,
+          propertyName: type.propertyName,
+        };
+      }
+      case 'Function': {
+        const returnType = this.tryResolveType(v, this.get(type.return));
+        if (returnType === null) {
+          return null;
+        }
+        return {
+          kind: 'Function',
+          return: returnType,
+          shapeId: type.shapeId,
+        };
+      }
+      case 'ObjectMethod':
+      case 'Object':
+      case 'Primitive':
+      case 'Poly': {
+        return type;
+      }
+      default: {
+        assertExhaustive(type, `Unexpected type kind '${(type as any).kind}'`);
+      }
+    }
+  }
+
   occursCheck(v: TypeVar, type: Type): boolean {
     if (typeEquals(v, type)) return true;
 

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/phi-type-inference-array-push-consecutive-phis.expect.md

@@ -0,0 +1,110 @@
+
+## Input
+
+```javascript
+import {makeArray} from 'shared-runtime';
+
+function Component(props) {
+  const x = {};
+  let y;
+  if (props.cond) {
+    if (props.cond2) {
+      y = [props.value];
+    } else {
+      y = [props.value2];
+    }
+  } else {
+    y = [];
+  }
+  // This should be inferred as `<store> y` s.t. `x` can still
+  // be independently memoized. *But* this also must properly
+  // extend the mutable range of the array literals in the
+  // if/else branches
+  y.push(x);
+
+  return [x, y];
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{cond: true, cond2: true, value: 42}],
+  sequentialRenders: [
+    {cond: true, cond2: true, value: 3.14},
+    {cond: true, cond2: true, value: 42},
+    {cond: true, cond2: true, value: 3.14},
+    {cond: true, cond2: false, value2: 3.14},
+    {cond: true, cond2: false, value2: 42},
+    {cond: true, cond2: false, value2: 3.14},
+    {cond: false},
+    {cond: false},
+  ],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime";
+import { makeArray } from "shared-runtime";
+
+function Component(props) {
+  const $ = _c(3);
+  let t0;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    t0 = {};
+    $[0] = t0;
+  } else {
+    t0 = $[0];
+  }
+  const x = t0;
+  let t1;
+  if ($[1] !== props) {
+    let y;
+    if (props.cond) {
+      if (props.cond2) {
+        y = [props.value];
+      } else {
+        y = [props.value2];
+      }
+    } else {
+      y = [];
+    }
+
+    y.push(x);
+
+    t1 = [x, y];
+    $[1] = props;
+    $[2] = t1;
+  } else {
+    t1 = $[2];
+  }
+  return t1;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ cond: true, cond2: true, value: 42 }],
+  sequentialRenders: [
+    { cond: true, cond2: true, value: 3.14 },
+    { cond: true, cond2: true, value: 42 },
+    { cond: true, cond2: true, value: 3.14 },
+    { cond: true, cond2: false, value2: 3.14 },
+    { cond: true, cond2: false, value2: 42 },
+    { cond: true, cond2: false, value2: 3.14 },
+    { cond: false },
+    { cond: false },
+  ],
+};
+
+```
+
+### Eval output
+(kind: ok) [{},[3.14,"[[ cyclic ref *1 ]]"]]
+[{},[42,"[[ cyclic ref *1 ]]"]]
+[{},[3.14,"[[ cyclic ref *1 ]]"]]
+[{},[3.14,"[[ cyclic ref *1 ]]"]]
+[{},[42,"[[ cyclic ref *1 ]]"]]
+[{},[3.14,"[[ cyclic ref *1 ]]"]]
+[{},["[[ cyclic ref *1 ]]"]]
+[{},["[[ cyclic ref *1 ]]"]]

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/phi-type-inference-array-push-consecutive-phis.js

@@ -0,0 +1,37 @@
+import {makeArray} from 'shared-runtime';
+
+function Component(props) {
+  const x = {};
+  let y;
+  if (props.cond) {
+    if (props.cond2) {
+      y = [props.value];
+    } else {
+      y = [props.value2];
+    }
+  } else {
+    y = [];
+  }
+  // This should be inferred as `<store> y` s.t. `x` can still
+  // be independently memoized. *But* this also must properly
+  // extend the mutable range of the array literals in the
+  // if/else branches
+  y.push(x);
+
+  return [x, y];
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{cond: true, cond2: true, value: 42}],
+  sequentialRenders: [
+    {cond: true, cond2: true, value: 3.14},
+    {cond: true, cond2: true, value: 42},
+    {cond: true, cond2: true, value: 3.14},
+    {cond: true, cond2: false, value2: 3.14},
+    {cond: true, cond2: false, value2: 42},
+    {cond: true, cond2: false, value2: 3.14},
+    {cond: false},
+    {cond: false},
+  ],
+};

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/phi-type-inference-array-push.expect.md

@@ -36,10 +36,17 @@ export const FIXTURE_ENTRYPOINT = {
 ```javascript
 import { c as _c } from "react/compiler-runtime";
 function Component(props) {
-  const $ = _c(2);
+  const $ = _c(3);
   let t0;
-  if ($[0] !== props) {
-    const x = {};
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    t0 = {};
+    $[0] = t0;
+  } else {
+    t0 = $[0];
+  }
+  const x = t0;
+  let t1;
+  if ($[1] !== props) {
     let y;
     if (props.cond) {
       y = [props.value];
@@ -49,13 +56,13 @@ function Component(props) {
 
     y.push(x);
 
-    t0 = [x, y];
-    $[0] = props;
-    $[1] = t0;
+    t1 = [x, y];
+    $[1] = props;
+    $[2] = t1;
   } else {
-    t0 = $[1];
+    t1 = $[2];
   }
-  return t0;
+  return t1;
 }
 
 export const FIXTURE_ENTRYPOINT = {