[compiler][wip] Infer alias effects for function expressions

This is a stab at addressing a pattern that mofeiz and I have both stumbled across. Today, FunctionExpression's context list describes values from the outer context that are accessed in the function, and with what effect they were accessed. This allows us to describe the fact that a value from the outer context is known to be mutated inside a function expression, or is known to be captured (aliased) into some other value in the function expression. However, the basic `Effect` kind is insufficient to describe the full semantics. Notably, it doesn't let us describe more complex aliasing relationships.

From an example mofeiz added:

```js
const x = {};
const y = {};
const f = () => {
  const a = [y];
  const b = x;
  // this sets y.x = x
  a[0].x = b;
}
f();
mutate(y.x);  // which means this mutates x!
```

Here, the Effect on the context operands are `[mutate y, read x]`. The `mutate y` is bc of the array push. But the `read x` is surprising — `x` is captured into `y`, but there is no subsequent mutation of y or x, so we consider this a read. But as the comments indicate, the final line mutates x! We need to reflect the fact that even though x isn't mutated inside the function, it is aliased into y, such that if y is subsequently mutated that this should count as a mutation of x too.

The idea of this PR is to extend the FunctionEffect type with a CaptureEffect variant which lists out the aliasing groups that occur inside the function expression. This allows us to bubble up the results of alias analysis from inside a function. The idea is to:

* Return the alias sets from InferMutableRanges
* Augment them with capturing of the form above, handling cases such as the `a[0].x = b`
* For each alias group, record a CaptureEffect for any group that contains 2+ context operands
* Extend the alias sets in the _outer_ function with the CaptureEffect sets from FunctionExpression/ObjectMethod instructions.

As part of this, I realized that our handling of PropertyStore's effect wasn't quite right. We used a store effect for the object, but only if it was a known object type — otherwise we recorded it as a mutation. But a PropertyStore really always is a store — it only mutates direct aliases of a value, not any interior objects that are captured. So I updated to always use store for known properties, and use mutate for computed properties. The latter is still also wrong, but i want to debug the change there separately.

ghstack-source-id: d8bf611
Pull Request resolved: #33151

Author: josephsavona

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 {