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

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;
 }
 
 ```