Update base for Update on "[compiler] Drop useMemos in memoization disabled mode, don't bail from source-level memo blocks"

Summary: When testing Forget with the always-bail-out mode, we now preserve useMemos as reactive scope rather than hook calls, and we don't short-circuit those scopes.

[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 }
@@ -154,7 +155,7 @@ function* runWithEnvironment(
   validateUseMemo(hir);
 
   if (
-    !env.config.enablePreserveExistingManualUseMemo &&
+    env.config.enablePreserveExistingManualUseMemo !== "hook" &&
     !env.config.disableMemoizationForDebugging &&
     !env.config.enableChangeDetectionForDebugging
   ) {
@@ -270,6 +271,15 @@ function* runWithEnvironment(
       value: hir,
     });
 
+    if (env.config.enablePreserveExistingManualUseMemo === "scope") {
+      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

@@ -1429,6 +1429,8 @@ export type ReactiveScope = {
   merged: Set<ScopeId>;
 
   loc: SourceLocation;
+
+  source: boolean;
 };
 
 export type ReactiveScopeDependencies = Set<ReactiveScopeDependency>;

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

@@ -334,6 +334,7 @@ function extractManualMemoizationArgs(
  */
 export function dropManualMemoization(func: HIRFunction): void {
   const isValidationEnabled =
+    func.env.config.enablePreserveExistingManualUseMemo === "scope" ||
     func.env.config.validatePreserveExistingMemoizationGuarantees ||
     func.env.config.enablePreserveExistingMemoizationGuarantees;
   const sidemap: IdentifierSidemap = {

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,108 @@
+/**
+ * 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,
+} 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"; 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") {
+      currentScope = { kind: "pending", id: nextId() };
+    } 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: new Set(),
+              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/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/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/__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.expect.md

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