[compiler] New inference repros/fixes

Adds a mechanism where if we can't determine a precise set of aliasing (data flow) effects for a function signature, to fall back to treating function expressions as general-purpose mutable values. Concretely:

* `InferMutationAliasingEffects` has an optimization where if you're `Apply`-ing a specific, known function expression, that we use the inferred signature for that function. Now we only use this mode _if_ we could infer a precise signature.
* `InferFunctionExpressionAliasingSignature` now bails out if there is a phi with a backedge, since we weren't handling this case.

See the repro for this case: i confirmed that it was previously incorrect.

Also adds a repro for a case of an invalid array length, which exposes the fact that InferFunctionExpressionAliasingSignatures is accumulating aliases w/o reducing them in any way. A complex enough function expression can lead to the accumulated array of effects overflowing the max array size. Fix in the next PR.

Author: josephsavona

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

@@ -20,11 +20,10 @@ import {inferReactiveScopeVariables} from '../ReactiveScopes';
 import {rewriteInstructionKindsBasedOnReassignment} from '../SSA';
 import {inferMutableRanges} from './InferMutableRanges';
 import inferReferenceEffects from './InferReferenceEffects';
-import {assertExhaustive, retainWhere} from '../Utils/utils';
+import {assertExhaustive} from '../Utils/utils';
 import {inferMutationAliasingEffects} from './InferMutationAliasingEffects';
 import {inferFunctionExpressionAliasingEffectsSignature} from './InferFunctionExpressionAliasingEffectsSignature';
 import {inferMutationAliasingRanges} from './InferMutationAliasingRanges';
-import {hashEffect} from './AliasingEffects';
 
 export default function analyseFunctions(func: HIRFunction): void {
   for (const [_, block] of func.body.blocks) {
@@ -69,89 +68,103 @@ function lowerWithMutationAliasing(fn: HIRFunction): void {
   analyseFunctions(fn);
   inferMutationAliasingEffects(fn, {isFunctionExpression: true});
   deadCodeElimination(fn);
-  inferMutationAliasingRanges(fn, {isFunctionExpression: true});
+  const mutationEffects = inferMutationAliasingRanges(fn, {
+    isFunctionExpression: true,
+  }).unwrap();
   rewriteInstructionKindsBasedOnReassignment(fn);
   inferReactiveScopeVariables(fn);
-  const effects = inferFunctionExpressionAliasingEffectsSignature(fn);
-  fn.env.logger?.debugLogIRs?.({
-    kind: 'hir',
-    name: 'AnalyseFunction (inner)',
-    value: fn,
-  });
-  if (effects != null) {
-    fn.aliasingEffects ??= [];
-    fn.aliasingEffects?.push(...effects);
-  }
-  if (fn.aliasingEffects != null) {
-    const seen = new Set<string>();
-    retainWhere(fn.aliasingEffects, effect => {
-      const hash = hashEffect(effect);
-      if (seen.has(hash)) {
-        return false;
-      }
-      seen.add(hash);
-      return true;
-    });
-  }
+  const aliasingEffects = inferFunctionExpressionAliasingEffectsSignature(fn);
+  if (aliasingEffects == null) {
+    /**
+     * If the data flow within the function expression is too complex to resolve,
+     * we null out the function expressions' alisingEffects. This ensures that
+     * InferMutationAliasingEffects at the outer level will fall back to a more
+     * conservative of the function as a plain mutable value — see the handling of
+     * 'Apply' in applyEffect(). As part of this we also set all context variables
+     * to Capture, which will downgrade if they turn out to be primitives, frozen,
+     * globals, etc.
+     */
+    fn.aliasingEffects = null;
+    for (const operand of fn.context) {
+      operand.effect = Effect.Capture;
+    }
+  } else {
+    /**
+     * We were able to determine a precise set of mutation and aliasing (data flow)
+     * effects for the function. Annotate which context variables are Capture based
+     * on which are mutated and/or captured into some other value. Capture is used
+     * as a signal to InferMutationAliasingEffects (at the outer function level) that
+     * the function itself has to be treated as mutable. See 'CreateFunction' handling
+     * in applyEffect().
+     */
+    const functionEffects = [...mutationEffects, ...aliasingEffects];
+    fn.aliasingEffects = functionEffects;
 
-  /**
-   * Phase 2: populate the Effect of each context variable to use in inferring
-   * the outer function. For example, InferMutationAliasingEffects uses context variable
-   * effects to decide if the function may be mutable or not.
-   */
-  const capturedOrMutated = new Set<IdentifierId>();
-  for (const effect of effects ?? []) {
-    switch (effect.kind) {
-      case 'Assign':
-      case 'Alias':
-      case 'Capture':
-      case 'CreateFrom': {
-        capturedOrMutated.add(effect.from.identifier.id);
-        break;
-      }
-      case 'Apply': {
-        CompilerError.invariant(false, {
-          reason: `[AnalyzeFunctions] Expected Apply effects to be replaced with more precise effects`,
-          loc: effect.function.loc,
-        });
-      }
-      case 'Mutate':
-      case 'MutateConditionally':
-      case 'MutateTransitive':
-      case 'MutateTransitiveConditionally': {
-        capturedOrMutated.add(effect.value.identifier.id);
-        break;
-      }
-      case 'Impure':
-      case 'Render':
-      case 'MutateFrozen':
-      case 'MutateGlobal':
-      case 'CreateFunction':
-      case 'Create':
-      case 'Freeze':
-      case 'ImmutableCapture': {
-        // no-op
-        break;
-      }
-      default: {
-        assertExhaustive(
-          effect,
-          `Unexpected effect kind ${(effect as any).kind}`,
-        );
+    /**
+     * Phase 2: populate the Effect of each context variable to use in inferring
+     * the outer function. For example, InferMutationAliasingEffects uses context variable
+     * effects to decide if the function may be mutable or not.
+     */
+    const capturedOrMutated = new Set<IdentifierId>();
+    for (const effect of functionEffects) {
+      switch (effect.kind) {
+        case 'Assign':
+        case 'Alias':
+        case 'Capture':
+        case 'CreateFrom': {
+          capturedOrMutated.add(effect.from.identifier.id);
+          break;
+        }
+        case 'Apply': {
+          CompilerError.invariant(false, {
+            reason: `[AnalyzeFunctions] Expected Apply effects to be replaced with more precise effects`,
+            loc: effect.function.loc,
+          });
+        }
+        case 'Mutate':
+        case 'MutateConditionally':
+        case 'MutateTransitive':
+        case 'MutateTransitiveConditionally': {
+          capturedOrMutated.add(effect.value.identifier.id);
+          break;
+        }
+        case 'Impure':
+        case 'Render':
+        case 'MutateFrozen':
+        case 'MutateGlobal':
+        case 'CreateFunction':
+        case 'Create':
+        case 'Freeze':
+        case 'ImmutableCapture': {
+          // no-op
+          break;
+        }
+        default: {
+          assertExhaustive(
+            effect,
+            `Unexpected effect kind ${(effect as any).kind}`,
+          );
+        }
       }
     }
-  }
 
-  for (const operand of fn.context) {
-    if (
-      capturedOrMutated.has(operand.identifier.id) ||
-      operand.effect === Effect.Capture
-    ) {
-      operand.effect = Effect.Capture;
-    } else {
-      operand.effect = Effect.Read;
+    for (const operand of fn.context) {
+      if (
+        capturedOrMutated.has(operand.identifier.id) ||
+        operand.effect === Effect.Capture
+      ) {
+        operand.effect = Effect.Capture;
+      } else {
+        operand.effect = Effect.Read;
+      }
     }
   }
+
+  fn.env.logger?.debugLogIRs?.({
+    kind: 'hir',
+    name: 'AnalyseFunction (inner)',
+    value: fn,
+  });
 }
 
 function lower(func: HIRFunction): void {

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

@@ -5,9 +5,16 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-import {HIRFunction, IdentifierId, Place, ValueKind, ValueReason} from '../HIR';
-import {getOrInsertDefault} from '../Utils/utils';
-import {AliasingEffect} from './AliasingEffects';
+import {
+  BlockId,
+  HIRFunction,
+  IdentifierId,
+  Place,
+  ValueKind,
+  ValueReason,
+} from '../HIR';
+import {getOrInsertDefault, retainWhere} from '../Utils/utils';
+import {AliasingEffect, hashEffect} from './AliasingEffects';
 
 /**
  * This function tracks data flow within an inner function expression in order to
@@ -77,11 +84,14 @@ export function inferFunctionExpressionAliasingEffectsSignature(
     );
   }
 
-  // todo: fixpoint
+  const seenBlocks = new Set<BlockId>();
   for (const block of fn.body.blocks.values()) {
     for (const phi of block.phis) {
       const operands: Array<AliasedIdentifier> = [];
-      for (const operand of phi.operands.values()) {
+      for (const [pred, operand] of phi.operands) {
+        if (!seenBlocks.has(pred)) {
+          return null;
+        }
         const inputs = lookup(operand, 'Alias');
         if (inputs != null) {
           operands.push(...inputs);
@@ -91,6 +101,7 @@ export function inferFunctionExpressionAliasingEffectsSignature(
         dataFlow.set(phi.place.identifier.id, operands);
       }
     }
+    seenBlocks.add(block.id);
     for (const instr of block.instructions) {
       if (instr.effects == null) continue;
       for (const effect of instr.effects) {
@@ -182,15 +193,18 @@ export function inferFunctionExpressionAliasingEffectsSignature(
     });
   }
 
+  const seen = new Set<string>();
+  retainWhere(effects, effect => {
+    const hash = hashEffect(effect);
+    if (seen.has(hash)) {
+      return false;
+    }
+    seen.add(hash);
+    return true;
+  });
   return effects;
 }
 
-export enum MutationKind {
-  None = 0,
-  Conditional = 1,
-  Definite = 2,
-}
-
 type AliasingKind = 'Alias' | 'Capture' | 'CreateFrom' | 'Assign';
 function joinEffects(
   effect1: AliasingKind,

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

@@ -822,7 +822,8 @@ function applyEffect(
       const functionValues = state.values(effect.function);
       if (
         functionValues.length === 1 &&
-        functionValues[0].kind === 'FunctionExpression'
+        functionValues[0].kind === 'FunctionExpression' &&
+        functionValues[0].loweredFunc.func.aliasingEffects != null
       ) {
         /*
          * We're calling a locally declared function, we already know it's effects!

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

@@ -22,8 +22,8 @@ import {
   eachTerminalOperand,
 } from '../HIR/visitors';
 import {assertExhaustive, getOrInsertWith} from '../Utils/utils';
-import {MutationKind} from './InferFunctionExpressionAliasingEffectsSignature';
-import {Result} from '../Utils/Result';
+import {Err, Ok, Result} from '../Utils/Result';
+import {AliasingEffect} from './AliasingEffects';
 
 /**
  * Infers mutable ranges for all values in the program, using previously inferred
@@ -49,7 +49,7 @@ import {Result} from '../Utils/Result';
 export function inferMutationAliasingRanges(
   fn: HIRFunction,
   {isFunctionExpression}: {isFunctionExpression: boolean},
-): Result<void, CompilerError> {
+): Result<Array<AliasingEffect>, CompilerError> {
   /**
    * Part 1: Infer mutable ranges for values. We build an abstract model of
    * values, the alias/capture edges between them, and the set of mutations.
@@ -215,7 +215,7 @@ export function inferMutationAliasingRanges(
   for (const render of renders) {
     state.render(render.index, render.place.identifier, errors);
   }
-  fn.aliasingEffects ??= [];
+  const functionEffects: Array<AliasingEffect> = [];
   for (const param of [...fn.context, ...fn.params]) {
     const place = param.kind === 'Identifier' ? param : param.place;
     const node = state.nodes.get(place.identifier);
@@ -226,13 +226,13 @@ export function inferMutationAliasingRanges(
     if (node.local != null) {
       if (node.local.kind === MutationKind.Conditional) {
         mutated = true;
-        fn.aliasingEffects.push({
+        functionEffects.push({
           kind: 'MutateConditionally',
           value: {...place, loc: node.local.loc},
         });
       } else if (node.local.kind === MutationKind.Definite) {
         mutated = true;
-        fn.aliasingEffects.push({
+        functionEffects.push({
           kind: 'Mutate',
           value: {...place, loc: node.local.loc},
         });
@@ -241,13 +241,13 @@ export function inferMutationAliasingRanges(
     if (node.transitive != null) {
       if (node.transitive.kind === MutationKind.Conditional) {
         mutated = true;
-        fn.aliasingEffects.push({
+        functionEffects.push({
           kind: 'MutateTransitiveConditionally',
           value: {...place, loc: node.transitive.loc},
         });
       } else if (node.transitive.kind === MutationKind.Definite) {
         mutated = true;
-        fn.aliasingEffects.push({
+        functionEffects.push({
           kind: 'MutateTransitive',
           value: {...place, loc: node.transitive.loc},
         });
@@ -436,7 +436,9 @@ export function inferMutationAliasingRanges(
     }
   }
 
-  return errors.asResult();
+  return errors.hasErrors() && !isFunctionExpression
+    ? Err(errors)
+    : Ok(functionEffects);
 }
 
 function appendFunctionErrors(errors: CompilerError, fn: HIRFunction): void {
@@ -452,6 +454,12 @@ function appendFunctionErrors(errors: CompilerError, fn: HIRFunction): void {
   }
 }
 
+export enum MutationKind {
+  None = 0,
+  Conditional = 1,
+  Definite = 2,
+}
+
 type Node = {
   id: Identifier;
   createdFrom: Map<Identifier, number>;

compiler/packages/babel-plugin-react-compiler/src/Inference/MUTABILITY_ALIASING_MODEL.md

@@ -514,9 +514,9 @@ Intuition: these effects are inverses of each other (capturing into an object, e
 Capture then CreatFrom is equivalent to Alias: we have to assume that the result _is_ the original value and that a local mutation of the result could mutate the original.
 
 ```js
-const y = [x]; // capture
-const z = y[0]; // createfrom
-mutate(z); // this clearly can mutate x, so the result must be one of Assign/Alias/CreateFrom 
+const b = [a]; // capture
+const c = b[0]; // createfrom
+mutate(c); // this clearly can mutate a, so the result must be one of Assign/Alias/CreateFrom 
 ```
 
 We use Alias as the return type because the mutability kind of the result is not derived from the source value (there's a fresh object in between due to the capture), so the full set of effects in practice would be a Create+Alias.
@@ -528,17 +528,17 @@ CreateFrom c <- b
 Alias c <- a
 ```
 
-Meanwhile the opposite direction preservers the capture, because the result is not the same as the source:
+Meanwhile the opposite direction preserves the capture, because the result is not the same as the source:
 
 ```js
-const y = x[0]; // createfrom
-const z = [y]; // capture
-mutate(z); // does not mutate x, so the result must be Capture
+const b = a[0]; // createfrom
+const c = [b]; // capture
+mutate(c); // does not mutate a, so the result must be Capture
 ```
 
 ```
-Capture b <- a
-CreateFrom c <- b
+CreateFrom b <- a
+Capture c <- b
 =>
-Capture b <- a
+Capture c <- a
 ```