[compiler] Stop relying on identifier mutable ranges after constructi…

…ng scopes

Addresses discussion at facebook#30399 (comment). 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: facebook#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/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',