Update base for Update on "[compiler] Fewer assumptions about nonmutability when change detection enabled"

Summary: Change detection is desgined to determine whether rules of react have been violated, and to do so better we need to loosen Forgets assumptions about what kinds of values don't need to be memoized. For example, the compiler typically doesn't think of `o.x` as something that needs to be memoized, because it does not allocate. However, we want to compare `o.x` in the current render with `o.x` in a previous one, so we now insert a "memoization" (comparison, really) block around this value.

[ghstack-poisoned]

compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts

@@ -97,6 +97,7 @@ import {
   validateUseMemo,
 } from "../Validation";
 import { validateLocalsNotReassignedAfterRender } from "../Validation/ValidateLocalsNotReassignedAfterRender";
+import { memoizeExistingUseMemos } from "../ReactiveScopes/MemoizeExistingUseMemos";
 
 export type CompilerPipelineValue =
   | { kind: "ast"; name: string; value: CodegenFunction }
@@ -153,11 +154,7 @@ function* runWithEnvironment(
   validateContextVariableLValues(hir);
   validateUseMemo(hir);
 
-  if (
-    !env.config.enablePreserveExistingManualUseMemo &&
-    !env.config.disableMemoizationForDebugging &&
-    !env.config.enableChangeDetectionForDebugging
-  ) {
+  if (env.config.enablePreserveExistingManualUseMemo !== "hook") {
     dropManualMemoization(hir);
     yield log({ kind: "hir", name: "DropManualMemoization", value: hir });
   }
@@ -270,6 +267,18 @@ function* runWithEnvironment(
       value: hir,
     });
 
+    if (
+      env.config.enablePreserveExistingManualUseMemo === "scope" ||
+      env.config.disableMemoizationForDebugging
+    ) {
+      memoizeExistingUseMemos(hir);
+      yield log({
+        kind: "hir",
+        name: "MemoizeExistingUseMemos",
+        value: hir,
+      });
+    }
+
     alignReactiveScopesToBlockScopesHIR(hir);
     yield log({
       kind: "hir",

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

@@ -182,7 +182,9 @@ const EnvironmentConfigSchema = z.object({
    * that the memoized values remain memoized, the compiler will simply not prune existing calls to
    * useMemo/useCallback.
    */
-  enablePreserveExistingManualUseMemo: z.boolean().default(false),
+  enablePreserveExistingManualUseMemo: z
+    .nullable(z.enum(["hook", "scope"]))
+    .default(null),
 
   // 🌲
   enableForest: z.boolean().default(false),
@@ -456,6 +458,31 @@ export function parseConfigPragma(pragma: string): EnvironmentConfig {
       continue;
     }
 
+    if (
+      key === "enablePreserveExistingManualUseMemo" &&
+      (val === undefined || val === "true" || val === "scope")
+    ) {
+      maybeConfig[key] = "scope";
+      continue;
+    }
+
+    if (key === "enablePreserveExistingManualUseMemo" && val === "hook") {
+      maybeConfig[key] = "hook";
+      continue;
+    }
+
+    if (
+      key === "enablePreserveExistingManualUseMemo" &&
+      !(val === "false" || val === "off")
+    ) {
+      CompilerError.throwInvalidConfig({
+        reason: `Invalid setting '${val}' for 'enablePreserveExistingManualUseMemo'. Valid settings are 'hook', 'scope', or 'off'.`,
+        description: null,
+        loc: null,
+        suggestions: null,
+      });
+    }
+
     if (typeof defaultConfig[key as keyof EnvironmentConfig] !== "boolean") {
       // skip parsing non-boolean properties
       continue;

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

@@ -771,7 +771,12 @@ export type ManualMemoDependency = {
         kind: "NamedLocal";
         value: Place;
       }
-    | { kind: "Global"; identifierName: string };
+    | {
+        kind: "InlinedGlobal";
+        value: Place;
+        name: string;
+      }
+    | { kind: "Global"; binding: LoadGlobal };
   path: Array<string>;
 };
 
@@ -1429,6 +1434,8 @@ export type ReactiveScope = {
   merged: Set<ScopeId>;
 
   loc: SourceLocation;
+
+  source: boolean;
 };
 
 export type ReactiveScopeDependencies = Set<ReactiveScopeDependency>;

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

@@ -843,7 +843,14 @@ export function printManualMemoDependency(
 ): string {
   let rootStr;
   if (val.root.kind === "Global") {
-    rootStr = val.root.identifierName;
+    rootStr = val.root.binding.binding.name;
+  } else if (val.root.kind === "InlinedGlobal") {
+    const nameStr = nameOnly
+      ? val.root.value.identifier.name != null
+        ? printName(val.root.value.identifier.name)
+        : String(val.root.value.identifier.id)
+      : printIdentifier(val.root.value.identifier);
+    rootStr = `G(${val.root.name}=${nameStr})`;
   } else {
     CompilerError.invariant(val.root.value.identifier.name?.kind === "named", {
       reason: "DepsValidation: expected named local variable in depslist",

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

@@ -229,7 +229,10 @@ export function* eachInstructionValueOperand(
     case "StartMemoize": {
       if (instrValue.deps != null) {
         for (const dep of instrValue.deps) {
-          if (dep.root.kind === "NamedLocal") {
+          if (
+            dep.root.kind === "NamedLocal" ||
+            dep.root.kind === "InlinedGlobal"
+          ) {
             yield dep.root.value;
           }
         }
@@ -554,7 +557,10 @@ export function mapInstructionValueOperands(
     case "StartMemoize": {
       if (instrValue.deps != null) {
         for (const dep of instrValue.deps) {
-          if (dep.root.kind === "NamedLocal") {
+          if (
+            dep.root.kind === "NamedLocal" ||
+            dep.root.kind === "InlinedGlobal"
+          ) {
             dep.root.value = fn(dep.root.value);
           }
         }

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

@@ -58,7 +58,7 @@ export function collectMaybeMemoDependencies(
       return {
         root: {
           kind: "Global",
-          identifierName: value.binding.name,
+          binding: value,
         },
         path: [],
       };
@@ -173,24 +173,53 @@ function makeManualMemoizationMarkers(
   env: Environment,
   depsList: Array<ManualMemoDependency> | null,
   memoDecl: Place,
-  manualMemoId: number
-): [TInstruction<StartMemoize>, TInstruction<FinishMemoize>] {
+  manualMemoId: number,
+  isManualUseMemoEnabled: boolean
+): [
+  Array<TInstruction<StartMemoize> | TInstruction<LoadGlobal>>,
+  TInstruction<FinishMemoize>,
+] {
+  let globals: Array<TInstruction<StartMemoize> | TInstruction<LoadGlobal>> =
+    [];
+  for (const dep of depsList ?? []) {
+    if (dep.root.kind === "Global" && isManualUseMemoEnabled) {
+      const place = createTemporaryPlace(env, dep.root.binding.loc);
+      globals.push({
+        id: makeInstructionId(0),
+        lvalue: place,
+        value: {
+          kind: "LoadGlobal",
+          binding: dep.root.binding.binding,
+          loc: dep.root.binding.loc,
+        },
+        loc: dep.root.binding.loc,
+      });
+      dep.root = {
+        kind: "InlinedGlobal",
+        value: place,
+        name: dep.root.binding.binding.name,
+      };
+    }
+  }
   return [
-    {
-      id: makeInstructionId(0),
-      lvalue: createTemporaryPlace(env, fnExpr.loc),
-      value: {
-        kind: "StartMemoize",
-        manualMemoId,
-        /*
-         * Use deps list from source instead of inferred deps
-         * as dependencies
-         */
-        deps: depsList,
+    [
+      ...globals,
+      {
+        id: makeInstructionId(0),
+        lvalue: createTemporaryPlace(env, fnExpr.loc),
+        value: {
+          kind: "StartMemoize",
+          manualMemoId,
+          /*
+           * Use deps list from source instead of inferred deps
+           * as dependencies
+           */
+          deps: depsList,
+          loc: fnExpr.loc,
+        },
         loc: fnExpr.loc,
       },
-      loc: fnExpr.loc,
-    },
+    ],
     {
       id: makeInstructionId(0),
       lvalue: createTemporaryPlace(env, fnExpr.loc),
@@ -333,9 +362,14 @@ function extractManualMemoizationArgs(
  * eg `React.useMemo()`.
  */
 export function dropManualMemoization(func: HIRFunction): void {
+  const isManualUseMemoEnabled =
+    func.env.config.enablePreserveExistingManualUseMemo === "scope" ||
+    func.env.config.enableChangeDetectionForDebugging != null ||
+    func.env.config.disableMemoizationForDebugging;
   const isValidationEnabled =
     func.env.config.validatePreserveExistingMemoizationGuarantees ||
-    func.env.config.enablePreserveExistingMemoizationGuarantees;
+    func.env.config.enablePreserveExistingMemoizationGuarantees ||
+    isManualUseMemoEnabled;
   const sidemap: IdentifierSidemap = {
     functions: new Map(),
     manualMemos: new Map(),
@@ -356,7 +390,11 @@ export function dropManualMemoization(func: HIRFunction): void {
    */
   const queuedInserts: Map<
     InstructionId,
-    TInstruction<StartMemoize> | TInstruction<FinishMemoize>
+    Array<
+      | TInstruction<StartMemoize>
+      | TInstruction<FinishMemoize>
+      | TInstruction<LoadGlobal>
+    >
   > = new Map();
   for (const [_, block] of func.body.blocks) {
     for (let i = 0; i < block.instructions.length; i++) {
@@ -419,7 +457,8 @@ export function dropManualMemoization(func: HIRFunction): void {
               func.env,
               depsList,
               memoDecl,
-              nextManualMemoId++
+              nextManualMemoId++,
+              isManualUseMemoEnabled
             );
 
             /**
@@ -436,7 +475,7 @@ export function dropManualMemoization(func: HIRFunction): void {
              * ```
              */
             queuedInserts.set(manualMemo.loadInstr.id, startMarker);
-            queuedInserts.set(instr.id, finishMarker);
+            queuedInserts.set(instr.id, [finishMarker]);
           }
         }
       } else {
@@ -457,8 +496,16 @@ export function dropManualMemoization(func: HIRFunction): void {
         const insertInstr = queuedInserts.get(instr.id);
         if (insertInstr != null) {
           nextInstructions = nextInstructions ?? block.instructions.slice(0, i);
+          const postInstructions: Array<Instruction> = [];
+          insertInstr.forEach((instr) => {
+            if (instr.value.kind === "LoadGlobal") {
+              nextInstructions?.push(instr);
+            } else {
+              postInstructions.push(instr);
+            }
+          });
           nextInstructions.push(instr);
-          nextInstructions.push(insertInstr);
+          postInstructions.forEach((instr) => nextInstructions?.push(instr));
         } else if (nextInstructions != null) {
           nextInstructions.push(instr);
         }

compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/CodegenReactiveFunction.ts

@@ -622,7 +622,7 @@ function codegenReactiveScope(
     );
   }
 
-  if (cx.env.config.disableMemoizationForDebugging) {
+  if (cx.env.config.disableMemoizationForDebugging && !scope.source) {
     CompilerError.invariant(
       cx.env.config.enableChangeDetectionForDebugging == null,
       {

compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/InferReactiveScopeVariables.ts

@@ -111,6 +111,7 @@ export function inferReactiveScopeVariables(fn: HIRFunction): void {
         earlyReturnValue: null,
         merged: new Set(),
         loc: identifier.loc,
+        source: false,
       };
       scopes.set(groupIdentifier, scope);
     } else {
@@ -161,7 +162,10 @@ export function inferReactiveScopeVariables(fn: HIRFunction): void {
   }
 }
 
-function mergeLocation(l: SourceLocation, r: SourceLocation): SourceLocation {
+export function mergeLocation(
+  l: SourceLocation,
+  r: SourceLocation
+): SourceLocation {
   if (l === GeneratedSource) {
     return r;
   } else if (r === GeneratedSource) {
@@ -186,15 +190,20 @@ export function isMutable({ id }: Instruction, place: Place): boolean {
   return id >= range.start && id < range.end;
 }
 
-function mayAllocate(env: Environment, instruction: Instruction): boolean {
+export function mayAllocate(
+  env: Environment,
+  instruction: Instruction,
+  conservative: boolean
+): boolean {
   const { value } = instruction;
   switch (value.kind) {
     case "Destructure": {
-      return doesPatternContainSpreadElement(value.lvalue.pattern);
+      return (
+        doesPatternContainSpreadElement(value.lvalue.pattern) || conservative
+      );
     }
     case "PostfixUpdate":
     case "PrefixUpdate":
-    case "Await":
     case "DeclareLocal":
     case "DeclareContext":
     case "StoreLocal":
@@ -205,26 +214,31 @@ function mayAllocate(env: Environment, instruction: Instruction): boolean {
     case "LoadContext":
     case "StoreContext":
     case "PropertyDelete":
-    case "ComputedLoad":
     case "ComputedDelete":
     case "JSXText":
     case "TemplateLiteral":
     case "Primitive":
     case "GetIterator":
     case "IteratorNext":
-    case "NextPropertyOf":
     case "Debugger":
     case "StartMemoize":
     case "FinishMemoize":
     case "UnaryExpression":
     case "BinaryExpression":
-    case "PropertyLoad":
     case "StoreGlobal": {
       return false;
     }
+    case "PropertyLoad":
+    case "NextPropertyOf":
+    case "ComputedLoad":
+    case "Await": {
+      return conservative;
+    }
     case "CallExpression":
     case "MethodCall": {
-      return instruction.lvalue.identifier.type.kind !== "Primitive";
+      return (
+        conservative || instruction.lvalue.identifier.type.kind !== "Primitive"
+      );
     }
     case "RegExpLiteral":
     case "PropertyStore":
@@ -249,6 +263,82 @@ function mayAllocate(env: Environment, instruction: Instruction): boolean {
   }
 }
 
+export function collectMutableOperands(
+  fn: HIRFunction,
+  instr: Instruction,
+  conservative: boolean
+): Array<Identifier> {
+  const operands: Array<Identifier> = [];
+  const range = instr.lvalue.identifier.mutableRange;
+  if (range.end > range.start + 1 || mayAllocate(fn.env, instr, conservative)) {
+    operands.push(instr.lvalue!.identifier);
+  }
+  if (
+    instr.value.kind === "StoreLocal" ||
+    instr.value.kind === "StoreContext"
+  ) {
+    if (
+      instr.value.lvalue.place.identifier.mutableRange.end >
+      instr.value.lvalue.place.identifier.mutableRange.start + 1
+    ) {
+      operands.push(instr.value.lvalue.place.identifier);
+    }
+    if (
+      isMutable(instr, instr.value.value) &&
+      instr.value.value.identifier.mutableRange.start > 0
+    ) {
+      operands.push(instr.value.value.identifier);
+    }
+  } else if (instr.value.kind === "Destructure") {
+    for (const place of eachPatternOperand(instr.value.lvalue.pattern)) {
+      if (
+        place.identifier.mutableRange.end >
+        place.identifier.mutableRange.start + 1
+      ) {
+        operands.push(place.identifier);
+      }
+    }
+    if (
+      isMutable(instr, instr.value.value) &&
+      instr.value.value.identifier.mutableRange.start > 0
+    ) {
+      operands.push(instr.value.value.identifier);
+    }
+  } else if (instr.value.kind === "MethodCall") {
+    for (const operand of eachInstructionOperand(instr)) {
+      if (
+        isMutable(instr, operand) &&
+        /*
+         * exclude global variables from being added to scopes, we can't recreate them!
+         * TODO: improve handling of module-scoped variables and globals
+         */
+        operand.identifier.mutableRange.start > 0
+      ) {
+        operands.push(operand.identifier);
+      }
+    }
+    /*
+     * Ensure that the ComputedLoad to resolve the method is in the same scope as the
+     * call itself
+     */
+    operands.push(instr.value.property.identifier);
+  } else {
+    for (const operand of eachInstructionOperand(instr)) {
+      if (
+        isMutable(instr, operand) &&
+        /*
+         * exclude global variables from being added to scopes, we can't recreate them!
+         * TODO: improve handling of module-scoped variables and globals
+         */
+        operand.identifier.mutableRange.start > 0
+      ) {
+        operands.push(operand.identifier);
+      }
+    }
+  }
+  return operands;
+}
+
 export function findDisjointMutableValues(
   fn: HIRFunction
 ): DisjointSet<Identifier> {
@@ -276,74 +366,7 @@ export function findDisjointMutableValues(
     }
 
     for (const instr of block.instructions) {
-      const operands: Array<Identifier> = [];
-      const range = instr.lvalue.identifier.mutableRange;
-      if (range.end > range.start + 1 || mayAllocate(fn.env, instr)) {
-        operands.push(instr.lvalue!.identifier);
-      }
-      if (
-        instr.value.kind === "StoreLocal" ||
-        instr.value.kind === "StoreContext"
-      ) {
-        if (
-          instr.value.lvalue.place.identifier.mutableRange.end >
-          instr.value.lvalue.place.identifier.mutableRange.start + 1
-        ) {
-          operands.push(instr.value.lvalue.place.identifier);
-        }
-        if (
-          isMutable(instr, instr.value.value) &&
-          instr.value.value.identifier.mutableRange.start > 0
-        ) {
-          operands.push(instr.value.value.identifier);
-        }
-      } else if (instr.value.kind === "Destructure") {
-        for (const place of eachPatternOperand(instr.value.lvalue.pattern)) {
-          if (
-            place.identifier.mutableRange.end >
-            place.identifier.mutableRange.start + 1
-          ) {
-            operands.push(place.identifier);
-          }
-        }
-        if (
-          isMutable(instr, instr.value.value) &&
-          instr.value.value.identifier.mutableRange.start > 0
-        ) {
-          operands.push(instr.value.value.identifier);
-        }
-      } else if (instr.value.kind === "MethodCall") {
-        for (const operand of eachInstructionOperand(instr)) {
-          if (
-            isMutable(instr, operand) &&
-            /*
-             * exclude global variables from being added to scopes, we can't recreate them!
-             * TODO: improve handling of module-scoped variables and globals
-             */
-            operand.identifier.mutableRange.start > 0
-          ) {
-            operands.push(operand.identifier);
-          }
-        }
-        /*
-         * Ensure that the ComputedLoad to resolve the method is in the same scope as the
-         * call itself
-         */
-        operands.push(instr.value.property.identifier);
-      } else {
-        for (const operand of eachInstructionOperand(instr)) {
-          if (
-            isMutable(instr, operand) &&
-            /*
-             * exclude global variables from being added to scopes, we can't recreate them!
-             * TODO: improve handling of module-scoped variables and globals
-             */
-            operand.identifier.mutableRange.start > 0
-          ) {
-            operands.push(operand.identifier);
-          }
-        }
-      }
+      const operands = collectMutableOperands(fn, instr, false);
       if (operands.length !== 0) {
         scopeIdentifiers.union(operands);
       }

compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/MemoizeExistingUseMemos.ts

@@ -0,0 +1,121 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+import { CompilerError } from "../CompilerError";
+import {
+  BasicBlock,
+  HIRFunction,
+  Identifier,
+  makeInstructionId,
+  ReactiveScope,
+  ReactiveScopeDependencies,
+} from "../HIR";
+import { eachTerminalSuccessor } from "../HIR/visitors";
+import {
+  collectMutableOperands,
+  mergeLocation,
+} from "./InferReactiveScopeVariables";
+
+export function memoizeExistingUseMemos(fn: HIRFunction): void {
+  visitBlock(fn, fn.body.blocks.get(fn.body.entry)!, null, new Map());
+}
+
+let ctr = 0;
+function nextId(): number {
+  return ctr++;
+}
+
+type CurrentScope =
+  | null
+  | { kind: "pending"; deps: ReactiveScopeDependencies; id: number }
+  | { kind: "available"; scope: ReactiveScope; id: number };
+
+function visitBlock(
+  fn: HIRFunction,
+  block: BasicBlock,
+  scope: CurrentScope,
+  seen: Map<number, CurrentScope>
+): void {
+  const visited = seen.get(block.id);
+  if (visited === undefined) {
+    seen.set(block.id, scope);
+  } else {
+    CompilerError.invariant(
+      visited === null ? scope === null : visited.id === scope?.id,
+      {
+        reason:
+          "MemoizeExistingUseMemos: visiting the same block with different scopes",
+        loc: null,
+        suggestions: null,
+      }
+    );
+    return;
+  }
+
+  function extend(
+    currentScope: ReactiveScope,
+    operands: Iterable<Identifier>
+  ): void {
+    for (const operand of operands) {
+      currentScope.range.start = makeInstructionId(
+        Math.min(currentScope.range.start, operand.mutableRange.start)
+      );
+      currentScope.range.end = makeInstructionId(
+        Math.max(currentScope.range.end, operand.mutableRange.end)
+      );
+      currentScope.loc = mergeLocation(currentScope.loc, operand.loc);
+      operand.scope = currentScope;
+      operand.mutableRange = currentScope.range;
+    }
+  }
+
+  let currentScope = scope;
+  for (const instruction of block.instructions) {
+    if (instruction.value.kind === "StartMemoize") {
+      const deps: ReactiveScopeDependencies = new Set();
+      for (const dep of instruction.value.deps ?? []) {
+        CompilerError.invariant(dep.root.kind !== "Global", {
+          reason:
+            "MemoizeExistingUseMemos: Globals should have been replaced with InlineGlobals",
+          loc: instruction.loc,
+          suggestions: null,
+          description: null,
+        });
+        deps.add({ identifier: dep.root.value.identifier, path: dep.path });
+      }
+      currentScope = { kind: "pending", id: nextId(), deps };
+    } else if (instruction.value.kind === "FinishMemoize") {
+      currentScope = null;
+    } else if (currentScope != null) {
+      const operands = collectMutableOperands(fn, instruction, true);
+      if (operands.length > 0) {
+        if (currentScope.kind === "pending") {
+          currentScope = {
+            kind: "available",
+            id: currentScope.id,
+            scope: {
+              id: fn.env.nextScopeId,
+              range: { start: instruction.id, end: instruction.id },
+              dependencies: currentScope.deps,
+              declarations: new Map(),
+              reassignments: new Set(),
+              earlyReturnValue: null,
+              merged: new Set(),
+              loc: instruction.loc,
+              source: true,
+            },
+          };
+        }
+        extend(currentScope.scope, operands);
+      }
+    }
+  }
+
+  for (const successor of eachTerminalSuccessor(block.terminal)) {
+    visitBlock(fn, fn.body.blocks.get(successor)!, currentScope, seen);
+  }
+}

compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/MergeReactiveScopesThatInvalidateTogether.ts

@@ -15,7 +15,6 @@ import {
   ReactiveFunction,
   ReactiveScope,
   ReactiveScopeBlock,
-  ReactiveScopeDependencies,
   ReactiveScopeDependency,
   ReactiveStatement,
   Type,
@@ -109,7 +108,7 @@ class FindLastUsageVisitor extends ReactiveFunctionVisitor<void> {
   }
 }
 
-class Transform extends ReactiveFunctionTransform<ReactiveScopeDependencies | null> {
+class Transform extends ReactiveFunctionTransform<ReactiveScope | null> {
   lastUsage: Map<IdentifierId, InstructionId>;
 
   constructor(lastUsage: Map<IdentifierId, InstructionId>) {
@@ -119,12 +118,13 @@ class Transform extends ReactiveFunctionTransform<ReactiveScopeDependencies | nu
 
   override transformScope(
     scopeBlock: ReactiveScopeBlock,
-    state: ReactiveScopeDependencies | null
+    state: ReactiveScope | null
   ): Transformed<ReactiveStatement> {
-    this.visitScope(scopeBlock, scopeBlock.scope.dependencies);
+    this.visitScope(scopeBlock, scopeBlock.scope);
     if (
       state !== null &&
-      areEqualDependencies(state, scopeBlock.scope.dependencies)
+      areEqualDependencies(state.dependencies, scopeBlock.scope.dependencies) &&
+      state.source === scopeBlock.scope.source
     ) {
       return { kind: "replace-many", value: scopeBlock.instructions };
     } else {
@@ -132,10 +132,7 @@ class Transform extends ReactiveFunctionTransform<ReactiveScopeDependencies | nu
     }
   }
 
-  override visitBlock(
-    block: ReactiveBlock,
-    state: ReactiveScopeDependencies | null
-  ): void {
+  override visitBlock(block: ReactiveBlock, state: ReactiveScope | null): void {
     // Pass 1: visit nested blocks to potentially merge their scopes
     this.traverseBlock(block, state);
 
@@ -417,6 +414,9 @@ function canMergeScopes(
   current: ReactiveScopeBlock,
   next: ReactiveScopeBlock
 ): boolean {
+  if (current.scope.source !== next.scope.source) {
+    return false;
+  }
   // Don't merge scopes with reassignments
   if (
     current.scope.reassignments.size !== 0 ||

compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PropagateScopeDependencies.ts

@@ -685,7 +685,9 @@ class PropagationVisitor extends ReactiveFunctionVisitor<Context> {
     const scopeDependencies = context.enter(scope.scope, () => {
       this.visitBlock(scope.instructions, context);
     });
-    scope.scope.dependencies = scopeDependencies;
+    if (!scope.scope.source) {
+      scope.scope.dependencies = scopeDependencies;
+    }
   }
 
   override visitPrunedScope(

compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PruneInitializationDependencies.ts

@@ -183,7 +183,7 @@ class Visitor extends ReactiveFunctionVisitor<CreateUpdate> {
       ident.path.forEach((key) => {
         target &&= this.paths.get(target)?.get(key);
       });
-      if (target && this.map.get(target) === "Create") {
+      if (target && this.map.get(target) === "Create" && !scope.scope.source) {
         scope.scope.dependencies.delete(ident);
       }
     });

compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PruneNonEscapingScopes.ts

@@ -947,7 +947,7 @@ class PruneScopesTransform extends ReactiveFunctionTransform<
       Array.from(scopeBlock.scope.reassignments).some((identifier) =>
         state.has(identifier.id)
       );
-    if (hasMemoizedOutput) {
+    if (hasMemoizedOutput || scopeBlock.scope.source) {
       return { kind: "keep" };
     } else {
       this.prunedScopes.add(scopeBlock.scope.id);

compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PruneNonReactiveDependencies.ts

@@ -97,7 +97,7 @@ class Visitor extends ReactiveFunctionVisitor<ReactiveIdentifiers> {
     this.traverseScope(scopeBlock, state);
     for (const dep of scopeBlock.scope.dependencies) {
       const isReactive = state.has(dep.identifier.id);
-      if (!isReactive) {
+      if (!isReactive && !scopeBlock.scope.source) {
         scopeBlock.scope.dependencies.delete(dep);
       }
     }

compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PruneUnusedScopes.ts

@@ -43,6 +43,7 @@ class Transform extends ReactiveFunctionTransform<State> {
     this.visitScope(scopeBlock, scopeState);
     if (
       !scopeState.hasReturnStatement &&
+      !scopeBlock.scope.source &&
       scopeBlock.scope.reassignments.size === 0 &&
       (scopeBlock.scope.declarations.size === 0 ||
         /*

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

@@ -146,9 +146,13 @@ function compareDeps(
   const rootsEqual =
     (inferred.root.kind === "Global" &&
       source.root.kind === "Global" &&
-      inferred.root.identifierName === source.root.identifierName) ||
-    (inferred.root.kind === "NamedLocal" &&
-      source.root.kind === "NamedLocal" &&
+      inferred.root.binding.binding.name ===
+        source.root.binding.binding.name) ||
+    ((inferred.root.kind === "NamedLocal" ||
+      inferred.root.kind === "InlinedGlobal") &&
+      (source.root.kind === "NamedLocal" ||
+        source.root.kind === "InlinedGlobal") &&
+      source.root.kind === inferred.root.kind &&
       inferred.root.value.identifier.id === source.root.value.identifier.id);
   if (!rootsEqual) {
     return CompareDependencyResult.RootDifference;
@@ -378,7 +382,8 @@ class Visitor extends ReactiveFunctionVisitor<VisitorState> {
 
     if (
       state.manualMemoState != null &&
-      state.manualMemoState.depsFromSource != null
+      state.manualMemoState.depsFromSource != null &&
+      !scopeBlock.scope.source
     ) {
       for (const dep of scopeBlock.scope.dependencies) {
         validateInferredDep(

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/nomemo-path.expect.md

@@ -0,0 +1,66 @@
+
+## Input
+
+```javascript
+// @disableMemoizationForDebugging
+import { useMemo } from "react";
+
+const w = 42;
+
+function Component(a) {
+  let x = useMemo(() => a.x, [a, w]);
+  return <div>{x}</div>;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ x: 42 }],
+  isComponent: true,
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime"; // @disableMemoizationForDebugging
+import { useMemo } from "react";
+
+const w = 42;
+
+function Component(a) {
+  const $ = _c(5);
+  const t0 = w;
+  let t1;
+  let t2;
+  if ($[0] !== a || $[1] !== t0) {
+    t2 = a.x;
+    $[0] = a;
+    $[1] = t0;
+    $[2] = t2;
+  } else {
+    t2 = $[2];
+  }
+  t1 = t2;
+  const x = t1;
+  let t3;
+  if ($[3] !== x || true) {
+    t3 = <div>{x}</div>;
+    $[3] = x;
+    $[4] = t3;
+  } else {
+    t3 = $[4];
+  }
+  return t3;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ x: 42 }],
+  isComponent: true,
+};
+
+```
+      
+### Eval output
+(kind: ok) <div>42</div>

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/nomemo-path.js

@@ -0,0 +1,15 @@
+// @disableMemoizationForDebugging
+import { useMemo } from "react";
+
+const w = 42;
+
+function Component(a) {
+  let x = useMemo(() => a.x, [a, w]);
+  return <div>{x}</div>;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ x: 42 }],
+  isComponent: true,
+};

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useMemo-preserve-non-idempotent.expect.md

@@ -0,0 +1,77 @@
+
+## Input
+
+```javascript
+// @enablePreserveExistingManualUseMemo
+import { useMemo } from "react";
+let cur = 99;
+function random(id) {
+  "use no forget";
+  cur = cur + 1;
+  return cur;
+}
+
+export default function C(id) {
+  const r = useMemo(() => random(id.id), [id.id]);
+  const a = r + 1;
+  return <>{a}</>;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: C,
+  params: [{ id: 1 }],
+  sequentialRenders: [{ id: 1 }, { id: 1 }, { id: 1 }, { id: 1 }],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime"; // @enablePreserveExistingManualUseMemo
+import { useMemo } from "react";
+let cur = 99;
+function random(id) {
+  "use no forget";
+  cur = cur + 1;
+  return cur;
+}
+
+export default function C(id) {
+  const $ = _c(4);
+  let t0;
+  let t1;
+  if ($[0] !== id.id) {
+    t1 = random(id.id);
+    $[0] = id.id;
+    $[1] = t1;
+  } else {
+    t1 = $[1];
+  }
+  t0 = t1;
+  const r = t0;
+  const a = r + 1;
+  let t2;
+  if ($[2] !== a) {
+    t2 = <>{a}</>;
+    $[2] = a;
+    $[3] = t2;
+  } else {
+    t2 = $[3];
+  }
+  return t2;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: C,
+  params: [{ id: 1 }],
+  sequentialRenders: [{ id: 1 }, { id: 1 }, { id: 1 }, { id: 1 }],
+};
+
+```
+      
+### Eval output
+(kind: ok) 101
+101
+101
+101

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useMemo-preserve-non-idempotent.js

@@ -0,0 +1,20 @@
+// @enablePreserveExistingManualUseMemo
+import { useMemo } from "react";
+let cur = 99;
+function random(id) {
+  "use no forget";
+  cur = cur + 1;
+  return cur;
+}
+
+export default function C(id) {
+  const r = useMemo(() => random(id.id), [id.id]);
+  const a = r + 1;
+  return <>{a}</>;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: C,
+  params: [{ id: 1 }],
+  sequentialRenders: [{ id: 1 }, { id: 1 }, { id: 1 }, { id: 1 }],
+};

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useMemo-simple-preserved-nomemo.expect.md

@@ -25,31 +25,25 @@ import { c as _c } from "react/compiler-runtime"; // @disableMemoizationForDebug
 import { useMemo } from "react";
 
 function Component(t0) {
-  const $ = _c(5);
+  const $ = _c(3);
   const { a } = t0;
   let t1;
-  if ($[0] !== a || true) {
-    t1 = () => [a];
-    $[0] = a;
-    $[1] = t1;
-  } else {
-    t1 = $[1];
-  }
   let t2;
-  if ($[2] === Symbol.for("react.memo_cache_sentinel") || true) {
-    t2 = [];
-    $[2] = t2;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    t2 = [a];
+    $[0] = t2;
   } else {
-    t2 = $[2];
+    t2 = $[0];
   }
-  const x = useMemo(t1, t2);
+  t1 = t2;
+  const x = t1;
   let t3;
-  if ($[3] !== x || true) {
+  if ($[1] !== x || true) {
     t3 = <div>{x}</div>;
-    $[3] = x;
-    $[4] = t3;
+    $[1] = x;
+    $[2] = t3;
   } else {
-    t3 = $[4];
+    t3 = $[2];
   }
   return t3;
 }

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useMemo-simple-preserved.expect.md

@@ -25,31 +25,25 @@ import { c as _c } from "react/compiler-runtime"; // @enablePreserveExistingManu
 import { useMemo } from "react";
 
 function Component(t0) {
-  const $ = _c(5);
+  const $ = _c(3);
   const { a } = t0;
   let t1;
-  if ($[0] !== a) {
-    t1 = () => [a];
-    $[0] = a;
-    $[1] = t1;
-  } else {
-    t1 = $[1];
-  }
   let t2;
-  if ($[2] === Symbol.for("react.memo_cache_sentinel")) {
-    t2 = [];
-    $[2] = t2;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    t2 = [a];
+    $[0] = t2;
   } else {
-    t2 = $[2];
+    t2 = $[0];
   }
-  const x = useMemo(t1, t2);
+  t1 = t2;
+  const x = t1;
   let t3;
-  if ($[3] !== x) {
+  if ($[1] !== x) {
     t3 = <div>{x}</div>;
-    $[3] = x;
-    $[4] = t3;
+    $[1] = x;
+    $[2] = t3;
   } else {
-    t3 = $[4];
+    t3 = $[2];
   }
   return t3;
 }

compiler/packages/snap/src/compiler.ts

@@ -46,6 +46,7 @@ function makePluginOptions(
   // TODO(@mofeiZ) rewrite snap fixtures to @validatePreserveExistingMemo:false
   let validatePreserveExistingMemoizationGuarantees = false;
   let enableChangeDetectionForDebugging = null;
+  let enablePreserveExistingManualUseMemo: "hook" | "scope" | null = null;
   let customMacros = null;
 
   if (firstLine.indexOf("@compilationMode(annotation)") !== -1) {
@@ -130,6 +131,28 @@ function makePluginOptions(
       importSpecifierName: "$structuralCheck",
     };
   }
+
+  const useMemoMatch = /@enablePreserveExistingManualUseMemo:"([^"]+)"/.exec(
+    firstLine
+  );
+  if (
+    useMemoMatch &&
+    (useMemoMatch[1] === "hook" || useMemoMatch[1] === "scope")
+  ) {
+    enablePreserveExistingManualUseMemo = useMemoMatch[1];
+  } else if (
+    useMemoMatch &&
+    (useMemoMatch[1] === "false" || useMemoMatch[1] === "off")
+  ) {
+    enablePreserveExistingManualUseMemo = null;
+  } else if (useMemoMatch) {
+    throw new Error(
+      `Invalid setting '${useMemoMatch[1]}' for 'enablePreserveExistingManualUseMemo'. Valid settings are 'hook', 'scope', or 'off'.`
+    );
+  } else if (firstLine.includes("@enablePreserveExistingManualUseMemo")) {
+    enablePreserveExistingManualUseMemo = "scope";
+  }
+
   const hookPatternMatch = /@hookPattern:"([^"]+)"/.exec(firstLine);
   if (
     hookPatternMatch &&
@@ -207,6 +230,7 @@ function makePluginOptions(
       hookPattern,
       validatePreserveExistingMemoizationGuarantees,
       enableChangeDetectionForDebugging,
+      enablePreserveExistingManualUseMemo,
     },
     compilationMode,
     logger,