[compiler] Fix mutable ranges for StoreContext

[ghstack-poisoned]

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

@@ -15,6 +15,7 @@ import {
   Hole,
   IdentifierId,
   Instruction,
+  InstructionKind,
   InstructionValue,
   isArrayType,
   isMapType,
@@ -768,7 +769,7 @@ function applyEffect(
   }
 }
 
-const DEBUG = true;
+const DEBUG = false;
 
 class InferenceState {
   env: Environment;
@@ -1452,7 +1453,21 @@ function computeSignatureForInstruction(
       }
       break;
     }
-    case 'DeclareContext':
+    case 'DeclareContext': {
+      // Context variables are conceptually like mutable boxes
+      effects.push({
+        kind: 'Create',
+        into: value.lvalue.place,
+        value: ValueKind.Mutable,
+      });
+      effects.push({
+        kind: 'Create',
+        into: lvalue,
+        // The result can't be referenced so this value doesn't matter
+        value: ValueKind.Primitive,
+      });
+      break;
+    }
     case 'DeclareLocal': {
       // TODO check this
       effects.push({
@@ -1481,13 +1496,25 @@ function computeSignatureForInstruction(
       break;
     }
     case 'LoadContext': {
+      // We load from the context
       effects.push({kind: 'Assign', from: value.place, into: lvalue});
       break;
     }
     case 'StoreContext': {
-      effects.push({kind: 'Mutate', value: value.lvalue.place});
+      // We're storing into the conceptual box
+      if (value.lvalue.kind === InstructionKind.Reassign) {
+        effects.push({kind: 'Mutate', value: value.lvalue.place});
+      } else {
+        // Context variables are conceptually like mutable boxes
+        effects.push({
+          kind: 'Create',
+          into: value.lvalue.place,
+          value: ValueKind.Mutable,
+        });
+      }
+      // Which aliases the value
       effects.push({
-        kind: 'Assign',
+        kind: 'Alias',
         from: value.value,
         into: value.lvalue.place,
       });

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

@@ -20,6 +20,7 @@ import {
 } from '../HIR/visitors';
 import DisjointSet from '../Utils/DisjointSet';
 import {assertExhaustive} from '../Utils/utils';
+import {debugAliases} from './InferMutableRanges';
 import {inferMutableRangesForAlias} from './InferMutableRangesForAlias';
 
 /**
@@ -44,7 +45,9 @@ export function inferMutationAliasingRanges(fn: HIRFunction): void {
 
     for (const instr of block.instructions) {
       for (const lvalue of eachInstructionLValue(instr)) {
-        lvalue.identifier.mutableRange.start = instr.id;
+        if (lvalue.identifier.mutableRange.start === 0) {
+          lvalue.identifier.mutableRange.start = instr.id;
+        }
         lvalue.identifier.mutableRange.end = makeInstructionId(
           Math.max(instr.id + 1, lvalue.identifier.mutableRange.end),
         );

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/error.invalid-useCallback-captures-reassigned-context.expect.md

@@ -0,0 +1,43 @@
+
+## Input
+
+```javascript
+// @validatePreserveExistingMemoizationGuarantees @enableNewMutationAliasingModel
+import {useCallback} from 'react';
+import {makeArray} from 'shared-runtime';
+
+// This case is already unsound in source, so we can safely bailout
+function Foo(props) {
+  let x = [];
+  x.push(props);
+
+  // makeArray() is captured, but depsList contains [props]
+  const cb = useCallback(() => [x], [x]);
+
+  x = makeArray();
+
+  return cb;
+}
+export const FIXTURE_ENTRYPOINT = {
+  fn: Foo,
+  params: [{}],
+};
+
+```
+
+
+## Error
+
+```
+   9 |
+  10 |   // makeArray() is captured, but depsList contains [props]
+> 11 |   const cb = useCallback(() => [x], [x]);
+     |                                      ^ CannotPreserveMemoization: React Compiler has skipped optimizing this component because the existing manual memoization could not be preserved. This dependency may be mutated later, which could cause the value to change unexpectedly (11:11)
+
+CannotPreserveMemoization: React Compiler has skipped optimizing this component because the existing manual memoization could not be preserved. This value was memoized in source but not in compilation output. (11:11)
+  12 |
+  13 |   x = makeArray();
+  14 |
+```
+          
+      

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/error.invalid-useCallback-captures-reassigned-context.js

@@ -0,0 +1,20 @@
+// @validatePreserveExistingMemoizationGuarantees @enableNewMutationAliasingModel
+import {useCallback} from 'react';
+import {makeArray} from 'shared-runtime';
+
+// This case is already unsound in source, so we can safely bailout
+function Foo(props) {
+  let x = [];
+  x.push(props);
+
+  // makeArray() is captured, but depsList contains [props]
+  const cb = useCallback(() => [x], [x]);
+
+  x = makeArray();
+
+  return cb;
+}
+export const FIXTURE_ENTRYPOINT = {
+  fn: Foo,
+  params: [{}],
+};

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/reactive-ref.expect.md

@@ -22,22 +22,28 @@ function ReactiveRefInEffect(props) {
 ```javascript
 import { c as _c } from "react/compiler-runtime"; // @enableNewMutationAliasingModel
 function ReactiveRefInEffect(props) {
-  const $ = _c(2);
+  const $ = _c(4);
   const ref1 = useRef("initial value");
   const ref2 = useRef("initial value");
   let ref;
-  if (props.foo) {
-    ref = ref1;
+  if ($[0] !== props.foo) {
+    if (props.foo) {
+      ref = ref1;
+    } else {
+      ref = ref2;
+    }
+    $[0] = props.foo;
+    $[1] = ref;
   } else {
-    ref = ref2;
+    ref = $[1];
   }
   let t0;
-  if ($[0] !== ref) {
+  if ($[2] !== ref) {
     t0 = () => print(ref);
-    $[0] = ref;
-    $[1] = t0;
+    $[2] = ref;
+    $[3] = t0;
   } else {
-    t0 = $[1];
+    t0 = $[3];
   }
   useEffect(t0);
 }