From 7fa4752e774a99937c34c076b0d38dfc8faf0565 Mon Sep 17 00:00:00 2001
From: Joe Savona <joesavona@fb.com>
Date: Tue, 23 Jul 2024 11:25:37 +0900
Subject: [PATCH] [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: 9ad6f5dc1773ab0ef2b4cfd01be106c6d929be51
Pull Request resolved: https://github.com/facebook/react/pull/30428
---
 .../src/HIR/BuildReactiveScopeTerminalsHIR.ts  | 18 ------------------
 .../ValidateMemoizedEffectDependencies.ts      |  2 ++
 .../ValidatePreservedManualMemoization.ts      | 15 ++++++++++++++-
 3 files changed, 16 insertions(+), 19 deletions(-)

diff --git a/compiler/packages/babel-plugin-react-compiler/src/HIR/BuildReactiveScopeTerminalsHIR.ts b/compiler/packages/babel-plugin-react-compiler/src/HIR/BuildReactiveScopeTerminalsHIR.ts
index 6819c2c6c5bfb..a761152853ff3 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/HIR/BuildReactiveScopeTerminalsHIR.ts
+++ b/compiler/packages/babel-plugin-react-compiler/src/HIR/BuildReactiveScopeTerminalsHIR.ts
@@ -179,24 +179,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') {
       /*
diff --git a/compiler/packages/babel-plugin-react-compiler/src/Validation/ValidateMemoizedEffectDependencies.ts b/compiler/packages/babel-plugin-react-compiler/src/Validation/ValidateMemoizedEffectDependencies.ts
index dcd12abf34205..259e03c2d3243 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/Validation/ValidateMemoizedEffectDependencies.ts
+++ b/compiler/packages/babel-plugin-react-compiler/src/Validation/ValidateMemoizedEffectDependencies.ts
@@ -99,6 +99,8 @@ 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))
       ) {
diff --git a/compiler/packages/babel-plugin-react-compiler/src/Validation/ValidatePreservedManualMemoization.ts b/compiler/packages/babel-plugin-react-compiler/src/Validation/ValidatePreservedManualMemoization.ts
index bbf271c50e5bb..4eb67230c41f3 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/Validation/ValidatePreservedManualMemoization.ts
+++ b/compiler/packages/babel-plugin-react-compiler/src/Validation/ValidatePreservedManualMemoization.ts
@@ -14,6 +14,7 @@ import {
   InstructionValue,
   ManualMemoDependency,
   Place,
+  PrunedReactiveScopeBlock,
   ReactiveFunction,
   ReactiveInstruction,
   ReactiveScopeBlock,
@@ -277,6 +278,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 +416,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 +474,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({