[compiler] useMemo calls directly induce memoization blocks

[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,7 @@ 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 +456,29 @@ export function parseConfigPragma(pragma: string): EnvironmentConfig {
       continue;
     }
 
+    if (
+      key === "enableChangeDetectionForDebugging" &&
+      (val === undefined || val === "true")
+    ) {
+      maybeConfig[key] = {
+        source: "react-compiler-runtime",
+        importSpecifierName: "$structuralCheck",
+      };
+      continue;
+    }
+
+    if (key === "enablePreserveExistingManualUseMemoAsHook" &&
+       (val === undefined || val === "true")) {
+      maybeConfig[key] = "hook"
+      continue;
+    }
+
+    if (key === "enablePreserveExistingManualUseMemoAsScope" &&
+       (val === undefined || val === "true")) {
+      maybeConfig[key] = "scope"
+      continue;
+    }
+
     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,7 @@ 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,7 +187,7 @@ 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": {
@@ -224,7 +225,7 @@ function mayAllocate(env: Environment, instruction: Instruction): boolean {
     }
     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 +250,78 @@ 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 +349,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,88 @@
+/**
+ * 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() {
+  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>) {
+  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>) {
+    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

@@ -109,7 +109,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 +119,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 {
@@ -134,7 +135,7 @@ class Transform extends ReactiveFunctionTransform<ReactiveScopeDependencies | nu
 
   override visitBlock(
     block: ReactiveBlock,
-    state: ReactiveScopeDependencies | null
+    state: ReactiveScope | null
   ): void {
     // Pass 1: visit nested blocks to potentially merge their scopes
     this.traverseBlock(block, state);
@@ -417,6 +418,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,82 @@
+
+## Input
+
+```javascript
+// @enablePreserveExistingManualUseMemoAsScope
+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"; // @enablePreserveExistingManualUseMemoAsScope
+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,25 @@
+// @enablePreserveExistingManualUseMemoAsScope
+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

@@ -2,7 +2,7 @@
 ## Input
 
 ```javascript
-// @enablePreserveExistingManualUseMemo
+// @enablePreserveExistingManualUseMemoAsScope
 import { useMemo } from "react";
 
 function Component({ a }) {
@@ -21,35 +21,30 @@ export const FIXTURE_ENTRYPOINT = {
 ## Code
 
 ```javascript
-import { c as _c } from "react/compiler-runtime"; // @enablePreserveExistingManualUseMemo
+import { c as _c } from "react/compiler-runtime"; // @enablePreserveExistingManualUseMemoAsScope
 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/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useMemo-simple-preserved.js

@@ -1,4 +1,4 @@
-// @enablePreserveExistingManualUseMemo
+// @enablePreserveExistingManualUseMemoAsScope
 import { useMemo } from "react";
 
 function Component({ a }) {

compiler/packages/snap/src/compiler.ts

@@ -45,7 +45,8 @@ function makePluginOptions(
   let hookPattern: string | null = null;
   // TODO(@mofeiZ) rewrite snap fixtures to @validatePreserveExistingMemo:false
   let validatePreserveExistingMemoizationGuarantees = false;
-  let enableChangeDetectionForDebugging = null;
+  let enableChangeDetectionForDebugging= null;
+  let enablePreserveExistingManualUseMemo: "hook" | "scope" | null  = null;
   let customMacros = null;
 
   if (firstLine.indexOf("@compilationMode(annotation)") !== -1) {
@@ -130,6 +131,15 @@ function makePluginOptions(
       importSpecifierName: "$structuralCheck",
     };
   }
+  if (firstLine.includes("@enablePreserveExistingManualUseMemoAsHook")) {
+    enablePreserveExistingManualUseMemo = "hook"
+  } else if (firstLine.includes("@enablePreserveExistingManualUseMemoAsScope")) {
+    enablePreserveExistingManualUseMemo = "scope"
+  } else if (firstLine.includes("@enablePreserveExistingManualUseMemo")) {
+    throw new Error(
+      'Use either @enablePreserveExistingManualUseMemoAsScope or @enablePreserveExistingManualUseMemoAsHook'
+    );
+  }
   const hookPatternMatch = /@hookPattern:"([^"]+)"/.exec(firstLine);
   if (
     hookPatternMatch &&
@@ -207,6 +217,7 @@ function makePluginOptions(
       hookPattern,
       validatePreserveExistingMemoizationGuarantees,
       enableChangeDetectionForDebugging,
+      enablePreserveExistingManualUseMemo,
     },
     compilationMode,
     logger,