[compiler] Infer alias effects for function expressions

compiler/packages/babel-plugin-react-compiler/src/HIR/HIR.ts

@@ -300,6 +300,10 @@ export type FunctionEffect =
       places: ReadonlySet<Place>;
       effect: Effect;
       loc: SourceLocation;
+    }
+  | {
+      kind: 'CaptureEffect';
+      places: ReadonlySet<Place>;
     };
 
 /*

compiler/packages/babel-plugin-react-compiler/src/HIR/PrintHIR.ts

@@ -546,12 +546,23 @@ export function printInstructionValue(instrValue: ReactiveValue): string {
       const effects =
         instrValue.loweredFunc.func.effects
           ?.map(effect => {
-            if (effect.kind === 'ContextMutation') {
-              return `ContextMutation places=[${[...effect.places]
-                .map(place => printPlace(place))
-                .join(', ')}] effect=${effect.effect}`;
-            } else {
-              return `GlobalMutation`;
+            switch (effect.kind) {
+              case 'ContextMutation': {
+                return `ContextMutation places=[${[...effect.places]
+                  .map(place => printPlace(place))
+                  .join(', ')}] effect=${effect.effect}`;
+              }
+              case 'GlobalMutation': {
+                return 'GlobalMutation';
+              }
+              case 'ReactMutation': {
+                return 'ReactMutation';
+              }
+              case 'CaptureEffect': {
+                return `CaptureEffect places=[${[...effect.places]
+                  .map(place => printPlace(place))
+                  .join(', ')}]`;
+              }
             }
           })
           .join(', ') ?? '';

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

@@ -11,6 +11,7 @@ import {
   HIRFunction,
   Identifier,
   LoweredFunction,
+  Place,
   isRefOrRefValue,
   makeInstructionId,
 } from '../HIR';
@@ -19,15 +20,23 @@ import {inferReactiveScopeVariables} from '../ReactiveScopes';
 import {rewriteInstructionKindsBasedOnReassignment} from '../SSA';
 import {inferMutableRanges} from './InferMutableRanges';
 import inferReferenceEffects from './InferReferenceEffects';
+import DisjointSet from '../Utils/DisjointSet';
+import {
+  eachInstructionLValue,
+  eachInstructionValueOperand,
+} from '../HIR/visitors';
+import prettyFormat from 'pretty-format';
+import {printIdentifier} from '../HIR/PrintHIR';
+import {Iterable_some} from '../Utils/utils';
 
 export default function analyseFunctions(func: HIRFunction): void {
   for (const [_, block] of func.body.blocks) {
     for (const instr of block.instructions) {
       switch (instr.value.kind) {
         case 'ObjectMethod':
         case 'FunctionExpression': {
-          lower(instr.value.loweredFunc.func);
-          infer(instr.value.loweredFunc);
+          const aliases = lower(instr.value.loweredFunc.func);
+          infer(instr.value.loweredFunc, aliases);
 
           /**
            * Reset mutable range for outer inferReferenceEffects
@@ -44,21 +53,82 @@ export default function analyseFunctions(func: HIRFunction): void {
   }
 }
 
-function lower(func: HIRFunction): void {
+function lower(func: HIRFunction): DisjointSet<Identifier> {
   analyseFunctions(func);
   inferReferenceEffects(func, {isFunctionExpression: true});
   deadCodeElimination(func);
-  inferMutableRanges(func);
+  const aliases = inferMutableRanges(func);
   rewriteInstructionKindsBasedOnReassignment(func);
   inferReactiveScopeVariables(func);
   func.env.logger?.debugLogIRs?.({
     kind: 'hir',
     name: 'AnalyseFunction (inner)',
     value: func,
   });
+  inferAliasesForCapturing(func, aliases);
+  return aliases;
+}
+
+export function debugAliases(aliases: DisjointSet<Identifier>): void {
+  console.log(
+    prettyFormat(
+      aliases
+        .buildSets()
+        .map(set => [...set].map(ident => printIdentifier(ident))),
+    ),
+  );
+}
+
+/**
+ * The alias sets returned by InferMutableRanges() accounts only for aliases that
+ * are known to mutate together. Notably this skips cases where a value is captured
+ * into some other value, but neither is subsequently mutated. An example is pushing
+ * a mutable value onto an array, where neither the array or value are subsequently
+ * mutated.
+ *
+ * This function extends the aliases sets to account for such capturing, so that we
+ * can detect cases where one of the values in a set is mutated later (in an outer function)
+ * we can correctly infer them as mutating together.
+ */
+function inferAliasesForCapturing(
+  fn: HIRFunction,
+  aliases: DisjointSet<Identifier>,
+): void {
+  for (const block of fn.body.blocks.values()) {
+    for (const instr of block.instructions) {
+      const {lvalue, value} = instr;
+      const hasStore =
+        lvalue.effect === Effect.Store ||
+        Iterable_some(
+          eachInstructionValueOperand(value),
+          operand => operand.effect === Effect.Store,
+        );
+      if (!hasStore) {
+        continue;
+      }
+      const operands: Array<Identifier> = [];
+      for (const lvalue of eachInstructionLValue(instr)) {
+        operands.push(lvalue.identifier);
+      }
+      for (const operand of eachInstructionValueOperand(instr.value)) {
+        if (
+          operand.effect === Effect.Store ||
+          operand.effect === Effect.Capture
+        ) {
+          operands.push(operand.identifier);
+        }
+      }
+      if (operands.length > 1) {
+        aliases.union(operands);
+      }
+    }
+  }
 }
 
-function infer(loweredFunc: LoweredFunction): void {
+function infer(
+  loweredFunc: LoweredFunction,
+  aliases: DisjointSet<Identifier>,
+): void {
   for (const operand of loweredFunc.func.context) {
     const identifier = operand.identifier;
     CompilerError.invariant(operand.effect === Effect.Unknown, {
@@ -85,6 +155,23 @@ function infer(loweredFunc: LoweredFunction): void {
       operand.effect = Effect.Read;
     }
   }
+  const contextIdentifiers = new Map(
+    loweredFunc.func.context.map(place => [place.identifier, place]),
+  );
+  for (const set of aliases.buildSets()) {
+    const contextOperands: Set<Place> = new Set(
+      [...set]
+        .map(identifier => contextIdentifiers.get(identifier))
+        .filter(place => place != null) as Array<Place>,
+    );
+    if (contextOperands.size !== 0) {
+      loweredFunc.func.effects ??= [];
+      loweredFunc.func.effects?.push({
+        kind: 'CaptureEffect',
+        places: contextOperands,
+      });
+    }
+  }
 }
 
 function isMutatedOrReassigned(id: Identifier): boolean {

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

@@ -60,6 +60,10 @@ function inferInstr(
       alias = instrValue.value;
       break;
     }
+    case 'IteratorNext': {
+      alias = instrValue.collection;
+      break;
+    }
     default:
       return;
   }

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

@@ -0,0 +1,33 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+import {HIRFunction, Identifier} from '../HIR/HIR';
+import DisjointSet from '../Utils/DisjointSet';
+
+export function inferAliasForFunctionCaptureEffects(
+  func: HIRFunction,
+  aliases: DisjointSet<Identifier>,
+): void {
+  for (const [_, block] of func.body.blocks) {
+    for (const instr of block.instructions) {
+      const {value} = instr;
+      if (
+        value.kind !== 'FunctionExpression' &&
+        value.kind !== 'ObjectMethod'
+      ) {
+        continue;
+      }
+      const loweredFunction = value.loweredFunc.func;
+      for (const effect of loweredFunction.effects ?? []) {
+        if (effect.kind !== 'CaptureEffect') {
+          continue;
+        }
+        aliases.union([...effect.places].map(place => place.identifier));
+      }
+    }
+  }
+}

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

@@ -95,45 +95,58 @@ function inheritFunctionEffects(
 
   return effects
     .flatMap(effect => {
-      if (effect.kind === 'GlobalMutation' || effect.kind === 'ReactMutation') {
-        return [effect];
-      } else {
-        const effects: Array<FunctionEffect | null> = [];
-        CompilerError.invariant(effect.kind === 'ContextMutation', {
-          reason: 'Expected ContextMutation',
-          loc: null,
-        });
-        /**
-         * Contextual effects need to be replayed against the current inference
-         * state, which may know more about the value to which the effect applied.
-         * The main cases are:
-         * 1. The mutated context value is _still_ a context value in the current scope,
-         *    so we have to continue propagating the original context mutation.
-         * 2. The mutated context value is a mutable value in the current scope,
-         *    so the context mutation was fine and we can skip propagating the effect.
-         * 3. The mutated context value  is an immutable value in the current scope,
-         *    resulting in a non-ContextMutation FunctionEffect. We propagate that new,
-         *    more detailed effect to the current function context.
-         */
-        for (const place of effect.places) {
-          if (state.isDefined(place)) {
-            const replayedEffect = inferOperandEffect(state, {
-              ...place,
-              loc: effect.loc,
-              effect: effect.effect,
-            });
-            if (replayedEffect != null) {
-              if (replayedEffect.kind === 'ContextMutation') {
-                // Case 1, still a context value so propagate the original effect
-                effects.push(effect);
-              } else {
-                // Case 3, immutable value so propagate the more precise effect
-                effects.push(replayedEffect);
-              }
-            } // else case 2, local mutable value so this effect was fine
+      switch (effect.kind) {
+        case 'GlobalMutation':
+        case 'ReactMutation': {
+          return [effect];
+        }
+        case 'ContextMutation': {
+          const effects: Array<FunctionEffect | null> = [];
+          CompilerError.invariant(effect.kind === 'ContextMutation', {
+            reason: 'Expected ContextMutation',
+            loc: null,
+          });
+          /**
+           * Contextual effects need to be replayed against the current inference
+           * state, which may know more about the value to which the effect applied.
+           * The main cases are:
+           * 1. The mutated context value is _still_ a context value in the current scope,
+           *    so we have to continue propagating the original context mutation.
+           * 2. The mutated context value is a mutable value in the current scope,
+           *    so the context mutation was fine and we can skip propagating the effect.
+           * 3. The mutated context value  is an immutable value in the current scope,
+           *    resulting in a non-ContextMutation FunctionEffect. We propagate that new,
+           *    more detailed effect to the current function context.
+           */
+          for (const place of effect.places) {
+            if (state.isDefined(place)) {
+              const replayedEffect = inferOperandEffect(state, {
+                ...place,
+                loc: effect.loc,
+                effect: effect.effect,
+              });
+              if (replayedEffect != null) {
+                if (replayedEffect.kind === 'ContextMutation') {
+                  // Case 1, still a context value so propagate the original effect
+                  effects.push(effect);
+                } else {
+                  // Case 3, immutable value so propagate the more precise effect
+                  effects.push(replayedEffect);
+                }
+              } // else case 2, local mutable value so this effect was fine
+            }
           }
+          return effects;
+        }
+        case 'CaptureEffect': {
+          return [];
+        }
+        default: {
+          assertExhaustive(
+            effect,
+            `Unexpected effect kind '${(effect as any).kind}'`,
+          );
         }
-        return effects;
       }
     })
     .filter((effect): effect is FunctionEffect => effect != null);
@@ -298,26 +311,31 @@ export function inferTerminalFunctionEffects(
 export function transformFunctionEffectErrors(
   functionEffects: Array<FunctionEffect>,
 ): Array<CompilerErrorDetailOptions> {
-  return functionEffects.map(eff => {
-    switch (eff.kind) {
-      case 'ReactMutation':
-      case 'GlobalMutation': {
-        return eff.error;
-      }
-      case 'ContextMutation': {
-        return {
-          severity: ErrorSeverity.Invariant,
-          reason: `Unexpected ContextMutation in top-level function effects`,
-          loc: eff.loc,
-        };
+  return functionEffects
+    .map(eff => {
+      switch (eff.kind) {
+        case 'ReactMutation':
+        case 'GlobalMutation': {
+          return eff.error;
+        }
+        case 'ContextMutation': {
+          return {
+            severity: ErrorSeverity.Invariant,
+            reason: `Unexpected ContextMutation in top-level function effects`,
+            loc: eff.loc,
+          };
+        }
+        case 'CaptureEffect': {
+          return null;
+        }
+        default:
+          assertExhaustive(
+            eff,
+            `Unexpected function effect kind \`${(eff as any).kind}\``,
+          );
       }
-      default:
-        assertExhaustive(
-          eff,
-          `Unexpected function effect kind \`${(eff as any).kind}\``,
-        );
-    }
-  });
+    })
+    .filter(eff => eff != null) as Array<CompilerErrorDetailOptions>;
 }
 
 function isEffectSafeOutsideRender(effect: FunctionEffect): boolean {