[compiler] Stop relying on identifier mutable ranges after constructing scopes

Addresses discussion at https://github.com/facebook/react/pull/30399#discussion_r1684693021. Once we've constructed scopes it's invalid to use identifier mutable ranges. The only places we can do this which i can find are ValidateMemoizedEffectDeps (which is already flawed and disabled by default) and ValidatePreservedManualMemoization. I added a todo to the former, and fixed up the latter.

The idea of the fix is that for StartMemo dependencies, if they needed to be memoized (identifier.scope != null) then that scope should exist and should have already completed. If they didn't need a scope or can't have one created (eg their range spans a hook), then their scope would be pruned. So if the scope is set, not pruned, and not completed, then it's an error.

For declarations (FinishMemo) the existing logic applies unchanged.

ghstack-source-id: af5bfd88553de3e30621695f9d139c4dc5efb997
Pull Request resolved: https://github.com/facebook/react/pull/30428

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

@@ -9,7 +9,6 @@ import {
   GotoVariant,
   HIRFunction,
   InstructionId,
-  makeInstructionId,
   ReactiveScope,
   ReactiveScopeTerminal,
   ScopeId,
@@ -19,7 +18,6 @@ import {
   markPredecessors,
   reversePostorderBlocks,
 } from './HIRBuilder';
-import {eachInstructionLValue} from './visitors';
 
 /**
  * This pass assumes that all program blocks are properly nested with respect to fallthroughs
@@ -179,24 +177,6 @@ export function buildReactiveScopeTerminalsHIR(fn: HIRFunction): void {
    * Fix scope and identifier ranges to account for renumbered instructions
    */
   for (const [, block] of fn.body.blocks) {
-    for (const instruction of block.instructions) {
-      for (const lvalue of eachInstructionLValue(instruction)) {
-        /*
-         * Any lvalues whose mutable range was a single instruction must have
-         * started at the current instruction, so update the range to match
-         * the instruction's new id
-         */
-        if (
-          lvalue.identifier.mutableRange.end ===
-          lvalue.identifier.mutableRange.start + 1
-        ) {
-          lvalue.identifier.mutableRange.start = instruction.id;
-          lvalue.identifier.mutableRange.end = makeInstructionId(
-            instruction.id + 1,
-          );
-        }
-      }
-    }
     const terminal = block.terminal;
     if (terminal.kind === 'scope' || terminal.kind === 'pruned-scope') {
       /*

compiler/packages/babel-plugin-react-compiler/src/Validation/ValidateMemoizedEffectDependencies.ts

@@ -99,6 +99,10 @@ class Visitor extends ReactiveFunctionVisitor<CompilerError> {
       const deps = instruction.value.args[1]!;
       if (
         deps.kind === 'Identifier' &&
+        /*
+         * TODO: isMutable is not safe to call here as it relies on identifier mutableRange which is no longer valid at this point
+         * in the pipeline
+         */
         (isMutable(instruction as Instruction, deps) ||
           isUnmemoized(deps.identifier, this.scopes))
       ) {

compiler/packages/babel-plugin-react-compiler/src/Validation/ValidatePreservedManualMemoization.ts

@@ -10,10 +10,10 @@ import {
   GeneratedSource,
   Identifier,
   IdentifierId,
-  Instruction,
   InstructionValue,
   ManualMemoDependency,
   Place,
+  PrunedReactiveScopeBlock,
   ReactiveFunction,
   ReactiveInstruction,
   ReactiveScopeBlock,
@@ -25,7 +25,6 @@ import {
 import {printManualMemoDependency} from '../HIR/PrintHIR';
 import {eachInstructionValueOperand} from '../HIR/visitors';
 import {collectMaybeMemoDependencies} from '../Inference/DropManualMemoization';
-import {isMutable} from '../ReactiveScopes/InferReactiveScopeVariables';
 import {
   ReactiveFunctionVisitor,
   visitReactiveFunction,
@@ -277,6 +276,7 @@ function validateInferredDep(
 
 class Visitor extends ReactiveFunctionVisitor<VisitorState> {
   scopes: Set<ScopeId> = new Set();
+  prunedScopes: Set<ScopeId> = new Set();
   scopeMapping = new Map();
   temporaries: Map<IdentifierId, ManualMemoDependency> = new Map();
 
@@ -414,6 +414,14 @@ class Visitor extends ReactiveFunctionVisitor<VisitorState> {
     }
   }
 
+  override visitPrunedScope(
+    scopeBlock: PrunedReactiveScopeBlock,
+    state: VisitorState,
+  ): void {
+    this.traversePrunedScope(scopeBlock, state);
+    this.prunedScopes.add(scopeBlock.scope.id);
+  }
+
   override visitInstruction(
     instruction: ReactiveInstruction,
     state: VisitorState,
@@ -464,7 +472,10 @@ class Visitor extends ReactiveFunctionVisitor<VisitorState> {
         instruction.value as InstructionValue,
       )) {
         if (
-          isMutable(instruction as Instruction, value) ||
+          (isDep &&
+            value.identifier.scope != null &&
+            !this.scopes.has(value.identifier.scope.id) &&
+            !this.prunedScopes.has(value.identifier.scope.id)) ||
           (isDecl && isUnmemoized(value.identifier, this.scopes))
         ) {
           state.errors.push({

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-memo-validation/error.false-positive-useMemo-with-refs.flow.expect.md

@@ -1,38 +0,0 @@
-
-## Input
-
-```javascript
-// @flow @validatePreserveExistingMemoizationGuarantees
-import { identity } from "shared-runtime";
-
-component Component(
-  disableLocalRef,
-  ref,
-) {
-  const localRef = useFooRef();
-  const mergedRef = useMemo(() => {
-    return disableLocalRef ? ref : identity(ref, localRef);
-  }, [disableLocalRef, ref, localRef]);
-  return <div ref={mergedRef} />;
-}
-
-```
-
-
-## Error
-
-```
-   7 | ) {
-   8 |   const localRef = useFooRef();
->  9 |   const mergedRef = useMemo(() => {
-     |                             ^^^^^^^
-> 10 |     return disableLocalRef ? ref : identity(ref, localRef);
-     | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-> 11 |   }, [disableLocalRef, ref, localRef]);
-     | ^^^^ CannotPreserveMemoization: React Compiler has skipped optimizing this component because the existing manual memoization could not be preserved. This value may be mutated later, which could cause the value to change unexpectedly (9:11)
-  12 |   return <div ref={mergedRef} />;
-  13 | }
-  14 |
-```
-          
-      

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-memo-validation/useMemo-with-refs.flow.expect.md

@@ -0,0 +1,57 @@
+
+## Input
+
+```javascript
+// @flow @validatePreserveExistingMemoizationGuarantees
+import { identity } from "shared-runtime";
+
+component Component(
+  disableLocalRef,
+  ref,
+) {
+  const localRef = useFooRef();
+  const mergedRef = useMemo(() => {
+    return disableLocalRef ? ref : identity(ref, localRef);
+  }, [disableLocalRef, ref, localRef]);
+  return <div ref={mergedRef} />;
+}
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime";
+import { identity } from "shared-runtime";
+
+const Component = React.forwardRef(Component_withRef);
+function Component_withRef(t0, ref) {
+  const $ = _c(6);
+  const { disableLocalRef } = t0;
+  const localRef = useFooRef();
+  let t1;
+  let t2;
+  if ($[0] !== disableLocalRef || $[1] !== ref || $[2] !== localRef) {
+    t2 = disableLocalRef ? ref : identity(ref, localRef);
+    $[0] = disableLocalRef;
+    $[1] = ref;
+    $[2] = localRef;
+    $[3] = t2;
+  } else {
+    t2 = $[3];
+  }
+  t1 = t2;
+  const mergedRef = t1;
+  let t3;
+  if ($[4] !== mergedRef) {
+    t3 = <div ref={mergedRef} />;
+    $[4] = mergedRef;
+    $[5] = t3;
+  } else {
+    t3 = $[5];
+  }
+  return t3;
+}
+
+```
+      

compiler/packages/snap/src/SproutTodoFilter.ts

@@ -398,6 +398,7 @@ const skipFilter = new Set([
   'deeply-nested-function-expressions-with-params',
   'readonly-object-method-calls',
   'readonly-object-method-calls-mutable-lambda',
+  'preserve-memo-validation/useMemo-with-refs.flow',
 
   // TODO: we probably want to always skip these
   'rules-of-hooks/rules-of-hooks-0592bd574811',