[compiler] Delay mutation of function expr context variables until function is called

See comments in code. The idea is that rather than immediately processing function expression effects when declaring the function, we record Capture effects for context variables that may be captured/mutated in the function. Then, transitive mutations of the function value itself will extend the range of these values via the normal captured value comutation inference established earlier in the stack (if capture a -> b, then transitiveMutate(b) => mutate(a)). So capture contextVar -> function and transitiveMutate(function) => mutate(contextVar).

ghstack-source-id: 3558955
Pull Request resolved: #33369

Author: josephsavona

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

@@ -10,11 +10,11 @@ import {
   Effect,
   HIRFunction,
   Identifier,
+  IdentifierId,
   LoweredFunction,
   Place,
   isRefOrRefValue,
   makeInstructionId,
-  printFunction,
 } from '../HIR';
 import {deadCodeElimination} from '../Optimization';
 import {inferReactiveScopeVariables} from '../ReactiveScopes';
@@ -26,7 +26,7 @@ import {
   eachInstructionLValue,
   eachInstructionValueOperand,
 } from '../HIR/visitors';
-import {Iterable_some} from '../Utils/utils';
+import {assertExhaustive, Iterable_some} from '../Utils/utils';
 import {inferMutationAliasingEffects} from './InferMutationAliasingEffects';
 import {inferMutationAliasingFunctionEffects} from './InferMutationAliasingFunctionEffects';
 import {inferMutationAliasingRanges} from './InferMutationAliasingRanges';
@@ -73,6 +73,49 @@ function lowerWithMutationAliasing(fn: HIRFunction): void {
   });
   const effects = inferMutationAliasingFunctionEffects(fn);
   fn.aliasingEffects = effects;
+
+  const capturedOrMutated = new Set<IdentifierId>();
+  for (const effect of effects ?? []) {
+    switch (effect.kind) {
+      case 'Alias':
+      case 'Capture':
+      case 'CreateFrom': {
+        capturedOrMutated.add(effect.from.identifier.id);
+        break;
+      }
+      case 'Apply': {
+        capturedOrMutated.add(effect.function.place.identifier.id);
+        break;
+      }
+      case 'Mutate':
+      case 'MutateConditionally':
+      case 'MutateTransitive':
+      case 'MutateTransitiveConditionally': {
+        capturedOrMutated.add(effect.value.identifier.id);
+        break;
+      }
+      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;
+    } else {
+      operand.effect = Effect.Read;
+    }
+  }
 }
 
 function lower(func: HIRFunction): DisjointSet<Identifier> {

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

@@ -5,7 +5,7 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-import {CompilerError, ValueKind} from '..';
+import {CompilerError, Effect, ValueKind} from '..';
 import {
   BasicBlock,
   BlockId,
@@ -1038,8 +1038,43 @@ function computeSignatureForInstruction(
         into: lvalue,
         value: ValueKind.Mutable,
       });
-      if (value.loweredFunc.func.aliasingEffects != null) {
-        effects.push(...value.loweredFunc.func.aliasingEffects);
+      /**
+       * We've already analyzed the function expression in AnalyzeFunctions. There, we assign
+       * a Capture effect to any context variable that appears (locally) to be aliased and/or
+       * mutated. The precise effects are annotated on the function expression's aliasingEffects
+       * property, but we don't want to execute those effects yet. We can only use those when
+       * we know exactly how the function is invoked — via an Apply effect from a custom signature.
+       *
+       * But in the general case, functions can be passed around and possibly called in ways where
+       * we don't know how to interpret their precise effects. For example:
+       *
+       * ```
+       * const a = {};
+       * // We don't want to consider a as mutating here either, this just declares the function
+       * const f = () => { maybeMutate(a) };
+       * // We don't want to consider a as mutating here either, it can't possibly call f yet
+       * const x = [f];
+       * // Here we have to assume that f can be called (transitively), and have to consider a
+       * // as mutating
+       * callAllFunctionInArray(x);
+       * ```
+       *
+       * So for any context variables that were inferred as captured or mutated, we record a
+       * Capture effect. If the resulting function is transitively mutated, this will mean
+       * that those operands are also considered mutated. If the function is never called,
+       * they won't be!
+       *
+       * Note that if the type of the context variables are frozen, global, or primitive, the
+       * Capture will either get pruned or downgraded to an ImmutableCapture.
+       */
+      for (const operand of value.loweredFunc.func.context) {
+        if (operand.effect === Effect.Capture) {
+          effects.push({
+            kind: 'Capture',
+            from: operand,
+            into: lvalue,
+          });
+        }
       }
       break;
     }

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/basic-mutation-via-function-expression.expect.md

@@ -10,7 +10,8 @@ function Component({a, b}) {
     y.x = x;
     mutate(y);
   };
-  return <div onClick={f}>{x}</div>;
+  f();
+  return <div>{x}</div>;
 }
 
 ```
@@ -20,36 +21,26 @@ function Component({a, b}) {
 ```javascript
 import { c as _c } from "react/compiler-runtime"; // @enableNewMutationAliasingModel
 function Component(t0) {
-  const $ = _c(7);
+  const $ = _c(3);
   const { a, b } = t0;
   let t1;
-  let x;
   if ($[0] !== a || $[1] !== b) {
-    x = { a };
+    const x = { a };
     const y = [b];
-    t1 = () => {
+    const f = () => {
       y.x = x;
       mutate(y);
     };
+
+    f();
+    t1 = <div>{x}</div>;
     $[0] = a;
     $[1] = b;
     $[2] = t1;
-    $[3] = x;
   } else {
     t1 = $[2];
-    x = $[3];
-  }
-  const f = t1;
-  let t2;
-  if ($[4] !== f || $[5] !== x) {
-    t2 = <div onClick={f}>{x}</div>;
-    $[4] = f;
-    $[5] = x;
-    $[6] = t2;
-  } else {
-    t2 = $[6];
   }
-  return t2;
+  return t1;
 }
 
 ```

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/potential-mutation-in-function-expression.expect.md

@@ -20,37 +20,42 @@ function Component({a, b, c}) {
 ```javascript
 import { c as _c } from "react/compiler-runtime"; // @enableNewMutationAliasingModel
 function Component(t0) {
-  const $ = _c(8);
+  const $ = _c(9);
   const { a, b, c } = t0;
   let t1;
-  let x;
-  if ($[0] !== a || $[1] !== b || $[2] !== c) {
-    x = [a, b];
-    t1 = () => {
+  if ($[0] !== a || $[1] !== b) {
+    t1 = [a, b];
+    $[0] = a;
+    $[1] = b;
+    $[2] = t1;
+  } else {
+    t1 = $[2];
+  }
+  const x = t1;
+  let t2;
+  if ($[3] !== c || $[4] !== x) {
+    t2 = () => {
       maybeMutate(x);
 
       console.log(c);
     };
-    $[0] = a;
-    $[1] = b;
-    $[2] = c;
-    $[3] = t1;
+    $[3] = c;
     $[4] = x;
+    $[5] = t2;
   } else {
-    t1 = $[3];
-    x = $[4];
+    t2 = $[5];
   }
-  const f = t1;
-  let t2;
-  if ($[5] !== f || $[6] !== x) {
-    t2 = <Foo onClick={f} value={x} />;
-    $[5] = f;
-    $[6] = x;
-    $[7] = t2;
+  const f = t2;
+  let t3;
+  if ($[6] !== f || $[7] !== x) {
+    t3 = <Foo onClick={f} value={x} />;
+    $[6] = f;
+    $[7] = x;
+    $[8] = t3;
   } else {
-    t2 = $[7];
+    t3 = $[8];
   }
-  return t2;
+  return t3;
 }
 
 ```