Option to preserve existing memoization guarantees

Adds an option to preserve existing memoization guarantees for values produced with useMemo and useCallback. We still discard the calls to these hooks, but we preserve the information that the value is frozen at that point in the program. Because these values are produced solely within the useMemo/useCallback callback, their mutation cannot have any interspersed hook calls. This means that the values mutable range will never span a hook and end at the point of the useMemo, ensuring that they are memoized at the same point. The main things that can change (relative to the orignal code) are: * Forget will infer a precise set of dependencies, ignoring the user-provided values. In practice this should only occur if the original code had a lint violation, which Forget would bail out on. So in practice this shouldn't happen unless the code doesn't use the React linter. * Forget may start the memoization block earlier than the developer did if other values are mutated along with the value being produced. This can cause memoization to fail, but only in situations where it would have failed previously: ```javascript const a = []; useFoo(); const b = useMemo(() => { const c = a; c.push(1); return c; }, [a]); ``` In this example (sans Forget) the useMemo will invalidate on every render because `a` will always be a new array and its listed as a dependency of the useMemo. Forget would correctly determine that the memoization would have to work as follows: ```javascript let c; if (...) { const a = [] useFoo(); // OOPS we made a hook call conditional const t0 = a; t0.push(1); c = t0; ... } else { c = $[...] } ``` Because this is invalid, Forget would (later in the pipeline) strip out this memoization block and (as with the original) leave `c` un-memoized. In this same example, removing the hook would cause Forget to be able to memoize a value that wasn't memoized before: ```javascript const a = []; const b = useMemo(() => { const c = a; c.push(1); return c; }, [a]); ``` This invalidates every render without Forget, but would memoize correctly with Forget (it would expand the memoization block to include the declaration of `a`).
2025-11-01 09:12:30 +00:00 · 2023-12-15 13:47:20 -08:00
parent fc36043019
commit 2abd439b43
16 changed files with 337 additions and 19 deletions
@@ -102,6 +102,26 @@ export type Hook = z.infer<typeof HookSchema>;
 const EnvironmentConfigSchema = z.object({
  customHooks: z.map(z.string(), HookSchema).optional().default(new Map()),

+  /**
+   * Enable using information from existing useMemo/useCallback to understand when a value is done
+   * being mutated. With this mode enabled, Forget will still discard the actual useMemo/useCallback
+   * calls and may memoize slightly differently. However, it will assume that the values produced
+   * are not subsequently modified, guaranteeing that the value will be memoized.
+   *
+   * By preserving guarantees about when values are memoized, this option preserves any existing
+   * behavior that depends on referential equality in the original program. Notably, this preserves
+   * existing effect behavior (how often effects fire) for effects that rely on referential equality.
+   *
+   * When disabled, Forget will not only prune useMemo and useCallback calls but also completely ignore
+   * them, not using any information from them to guide compilation. Therefore, disabling this flag
+   * will produce output that mimics the result from removing all memoization.
+   *
+   * Our recommendation is to first try running your application with this flag enabled, then attempt
+   * to disable this flag and see what changes or breaks. This will mostly likely be effects that
+   * depend on referential equality, which can be refactored (TODO guide for this).
+   */
+  enablePreserveExistingMemoizationGuarantees: z.boolean().default(false),
+
  // 🌲
  enableForest: z.boolean().default(false),
  // <🌲>
@@ -845,6 +845,12 @@ export type InstructionValue =
    }
  // `debugger` statement
  | { kind: "Debugger"; loc: SourceLocation }
+  /*
+   * Represents semantic information from useMemo/useCallback that the developer
+   * has indicated a particular value should be memoized. This value is ignored
+   * unless the TODO flag is enabled.
+   */
+  | { kind: "Memoize"; value: Place; loc: SourceLocation }
  /*
   * Catch-all for statements such as type imports, nested class declarations, etc
   * which are not directly represented, but included for completeness and to allow
@@ -15,11 +15,14 @@ import {
  BasicBlock,
  BlockId,
  BlockKind,
+  Effect,
+  GeneratedSource,
  GotoVariant,
  HIR,
  Identifier,
  IdentifierId,
  Instruction,
+  Place,
  Terminal,
  makeBlockId,
  makeInstructionId,
@@ -856,3 +859,19 @@ export function removeUnnecessaryTryCatch(fn: HIR): void {
    }
  }
 }
+
+export function createTemporaryPlace(env: Environment): Place {
+  return {
+    kind: "Identifier",
+    identifier: {
+      id: env.nextIdentifierId,
+      mutableRange: { start: makeInstructionId(0), end: makeInstructionId(0) },
+      name: null,
+      scope: null,
+      type: makeType(),
+    },
+    reactive: false,
+    effect: Effect.Unknown,
+    loc: GeneratedSource,
+  };
+}
@@ -594,6 +594,10 @@ export function printInstructionValue(instrValue: ReactiveValue): string {
      } ${printPlace(instrValue.value)}`;
      break;
    }
+    case "Memoize": {
+      value = `Memoize ${printPlace(instrValue.value)}`;
+      break;
+    }
    default: {
      assertExhaustive(
        instrValue,
@@ -209,6 +209,10 @@ export function* eachInstructionValueOperand(
      yield instrValue.value;
      break;
    }
+    case "Memoize": {
+      yield instrValue.value;
+      break;
+    }
    case "Debugger":
    case "RegExpLiteral":
    case "LoadGlobal":
@@ -510,6 +514,10 @@ export function mapInstructionValueOperands(
      instrValue.value = fn(instrValue.value);
      break;
    }
+    case "Memoize": {
+      instrValue.value = fn(instrValue.value);
+      break;
+    }
    case "Debugger":
    case "RegExpLiteral":
    case "LoadGlobal":
@@ -10,9 +10,13 @@ import {
  Effect,
  HIRFunction,
  IdentifierId,
+  Instruction,
  Place,
  SpreadPattern,
+  makeInstructionId,
+  markInstructionIds,
 } from "../HIR";
+import { createTemporaryPlace } from "../HIR/HIRBuilder";
 import { HookKind } from "../HIR/ObjectShape";

 /*
@@ -26,8 +30,14 @@ import { HookKind } from "../HIR/ObjectShape";
 export function dropManualMemoization(func: HIRFunction): void {
  const hooks = new Map<IdentifierId, HookKind>();
  const react = new Set<IdentifierId>();
+  let hasChanges = false;
  for (const [_, block] of func.body.blocks) {
-    for (const instr of block.instructions) {
+    let nextInstructions: Array<Instruction> | null = null;
+    for (let i = 0; i < block.instructions.length; i++) {
+      const instr = block.instructions[i]!;
+      if (nextInstructions !== null) {
+        nextInstructions.push(instr);
+      }
      switch (instr.value.kind) {
        case "LoadGlobal": {
          if (
@@ -71,16 +81,22 @@ export function dropManualMemoization(func: HIRFunction): void {
                });
              }
              /*
-               * TODO(gsn): Consider inlining the function passed to useMemo,
-               * rather than just calling it directly.
-               *
               * Replace the hook callee with the fn arg.
               *
               * before:
-               *   foo = Call useMemo$2($9, $10)
+               *   $1 = LoadGlobal useMemo       // load the useMemo global
+               *   $2 = FunctionExpression ...   // memo function
+               *   $3 = ArrayExpression [ ... ]  // deps array
+               *   $4 = Call $1 ($2, $3 )        // invoke useMemo w fn and deps
               *
               * after:
-               *   foo = Call $9()
+               *   $1 = LoadGlobal useMemo       // load the useMemo global (dead code)
+               *   $2 = FunctionExpression ...   // memo function
+               *   $3 = ArrayExpression [ ... ]  // deps array (dead code)
+               *   $4 = Call $2 ()               // invoke the memo function itself
+               *
+               * Note that a later pass (InlineImmediatelyInvokedFunctionExpressions) will
+               * inline the useMemo callback along with any other immediately invoked IIFEs.
               */
              if (fn.kind === "Identifier") {
                instr.value = {
@@ -93,6 +109,46 @@ export function dropManualMemoization(func: HIRFunction): void {
                  args: [],
                  loc: instr.value.loc,
                };
+                if (
+                  func.env.config.enablePreserveExistingMemoizationGuarantees
+                ) {
+                  /**
+                   * When this flag is enabled we also compile in a 'Memoize' instruction
+                   * to preserve the intended memoization boundary:
+                   *
+                   * Normal output:
+                   *   $1 = LoadGlobal useMemo       // load the useMemo global (dead code)
+                   *   $2 = FunctionExpression ...   // memo function
+                   *   $3 = ArrayExpression [ ... ]  // deps array (dead code)
+                   *   $4 = Call $2 ()               // invoke the memo function itself
+                   *
+                   * Output w flag enabled:
+                   *   $1 = LoadGlobal useMemo       // load the useMemo global (dead code)
+                   *   $2 = FunctionExpression ...   // memo function
+                   *   $3 = ArrayExpression [ ... ]  // deps array (dead code)
+                   *   $5 = Call $2 ()               // invoke the memo function itself
+                   *   $4 = Memoize $5               // preserve memo information
+                   *
+                   * Note that we synthesize a new temporary for the call ($5) and use
+                   * the original lvalue for the result of the Memoize instruction, so that
+                   * we don't have to rewrite subsequent instructions.
+                   */
+                  const lvalue = instr.lvalue;
+                  const temp = createTemporaryPlace(func.env);
+                  instr.lvalue = { ...temp };
+                  nextInstructions =
+                    nextInstructions ?? block.instructions.slice(0, i + 1);
+                  nextInstructions.push({
+                    id: makeInstructionId(0),
+                    lvalue,
+                    value: {
+                      kind: "Memoize",
+                      value: temp,
+                      loc: instr.loc,
+                    },
+                    loc: instr.loc,
+                  });
+                }
              }
            } else if (hookKind === "useCallback") {
              const [fn] = instr.value.args as Array<
@@ -110,23 +166,63 @@ export function dropManualMemoization(func: HIRFunction): void {
               * Instead of a Call, just alias the callback directly.
               *
               * before:
-               *   foo = Call useCallback$8($19)
+               *   $1 = LoadGlobal useCallback
+               *   $2 = FunctionExpression ...   // the callback being memoized
+               *   $3 = ArrayExpression ...      // deps array
+               *   $3 = Call $1 ( $2, $3 )       // invoke useCallback
               *
               * after:
-               *   foo = $19
+               *   $1 = LoadGlobal useCallback   // dead code
+               *   $2 = FunctionExpression ...   // the callback being memoized
+               *   $3 = ArrayExpression ...      // deps array (dead code)
+               *   $3 = LoadLocal $2             // reference the function
               */
              if (fn.kind === "Identifier") {
-                instr.value = {
-                  kind: "LoadLocal",
-                  place: {
-                    kind: "Identifier",
-                    identifier: fn.identifier,
-                    effect: Effect.Unknown,
-                    reactive: false,
+                if (
+                  func.env.config.enablePreserveExistingMemoizationGuarantees
+                ) {
+                  /**
+                   * With the flag enabled the output changes to use a Memoize instruction instead
+                   * a loadlocal to load the function expression into the original temporary:
+                   *
+                   * Normal output:
+                   *   $1 = LoadGlobal useCallback   // dead code
+                   *   $2 = FunctionExpression ...   // the callback being memoized
+                   *   $3 = ArrayExpression ...      // deps array (dead code)
+                   *   $3 = LoadLocal $2             // reference the function
+                   *
+                   * With flag enabled:
+                   *   $1 = LoadGlobal useCallback   // dead code
+                   *   $2 = FunctionExpression ...   // the callback being memoized
+                   *   $3 = ArrayExpression ...      // deps array (dead code)
+                   *   $3 = Memoize $2             // reference the function
+                   *
+                   * Note the s/LoadLocal/Memoize/
+                   */
+                  instr.value = {
+                    kind: "Memoize",
+                    value: {
+                      kind: "Identifier",
+                      identifier: fn.identifier,
+                      effect: Effect.Unknown,
+                      reactive: false,
+                      loc: instr.value.loc,
+                    },
                    loc: instr.value.loc,
-                  },
-                  loc: instr.value.loc,
-                };
+                  };
+                } else {
+                  instr.value = {
+                    kind: "LoadLocal",
+                    place: {
+                      kind: "Identifier",
+                      identifier: fn.identifier,
+                      effect: Effect.Unknown,
+                      reactive: false,
+                      loc: instr.value.loc,
+                    },
+                    loc: instr.value.loc,
+                  };
+                }
              }
            }
          }
@@ -134,5 +230,12 @@ export function dropManualMemoization(func: HIRFunction): void {
        }
      }
    }
+    if (nextInstructions !== null) {
+      block.instructions = nextInstructions;
+      hasChanges = true;
+    }
+  }
+  if (hasChanges) {
+    markInstructionIds(func.body);
  }
 }
@@ -1162,6 +1162,17 @@ function inferBlock(
        state.alias(lvalue, instrValue.value);
        continue;
      }
+      case "Memoize": {
+        state.initialize(instrValue, {
+          kind: ValueKind.Frozen,
+          reason: new Set([ValueReason.Other]),
+        });
+        state.reference(instrValue.value, Effect.Freeze, ValueReason.Other);
+        const lvalue = instr.lvalue;
+        lvalue.effect = Effect.ConditionallyMutate;
+        state.alias(lvalue, instrValue.value);
+        continue;
+      }
      case "LoadLocal": {
        const lvalue = instr.lvalue;
        const effect =
@@ -313,6 +313,7 @@ function pruneableValue(value: InstructionValue, state: State): boolean {
    case "StoreContext": {
      return false;
    }
+    case "Memoize":
    case "RegExpLiteral":
    case "LoadGlobal":
    case "ArrayExpression":
@@ -1628,6 +1628,10 @@ function codegenInstructionValue(
      );
      break;
    }
+    case "Memoize": {
+      value = codegenPlaceToExpression(cx, instrValue.value);
+      break;
+    }
    case "Debugger":
    case "DeclareLocal":
    case "DeclareContext":
@@ -245,7 +245,8 @@ function mayAllocate(env: Environment, instruction: Instruction): boolean {
    case "Primitive":
    case "NextIterableOf":
    case "NextPropertyOf":
-    case "Debugger": {
+    case "Debugger":
+    case "Memoize": {
      return false;
    }
    case "UnaryExpression":
@@ -471,6 +471,7 @@ function computeMemoizationInputs(
        rvalues: [],
      };
    }
+    case "Memoize":
    case "Await":
    case "TypeCastExpression":
    case "NextIterableOf": {
@@ -280,6 +280,11 @@ function* generateInstructionTypes(
      break;
    }

+    case "Memoize": {
+      yield equation(left, value.value.identifier.type);
+      break;
+    }
+
    case "PropertyDelete":
    case "ComputedDelete": {
      yield equation(left, { kind: "Primitive" });
@@ -0,0 +1,54 @@
+
+## Input
+
+```javascript
+// @enablePreserveExistingMemoizationGuarantees:false
+import { useMemo } from "react";
+import { identity, makeObject_Primitives, mutate } from "shared-runtime";
+
+function Component(props) {
+  const object = useMemo(() => makeObject_Primitives(), []);
+  identity(object);
+  return object;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{}],
+};
+
+```
+
+## Code
+
+```javascript
+// @enablePreserveExistingMemoizationGuarantees:false
+import { useMemo, unstable_useMemoCache as useMemoCache } from "react";
+import { identity, makeObject_Primitives, mutate } from "shared-runtime";
+
+function Component(props) {
+  const $ = useMemoCache(2);
+  let t7;
+  let object;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    t7 = makeObject_Primitives();
+    object = t7;
+    identity(object);
+    $[0] = object;
+    $[1] = t7;
+  } else {
+    object = $[0];
+    t7 = $[1];
+  }
+  return object;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{}],
+};
+
+```
+      
+### Eval output
+(kind: ok) {"a":0,"b":"value1","c":true}
@@ -0,0 +1,14 @@
+// @enablePreserveExistingMemoizationGuarantees:false
+import { useMemo } from "react";
+import { identity, makeObject_Primitives, mutate } from "shared-runtime";
+
+function Component(props) {
+  const object = useMemo(() => makeObject_Primitives(), []);
+  identity(object);
+  return object;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{}],
+};
@@ -0,0 +1,53 @@
+
+## Input
+
+```javascript
+// @enablePreserveExistingMemoizationGuarantees
+import { useMemo } from "react";
+import { identity, makeObject_Primitives, mutate } from "shared-runtime";
+
+function Component(props) {
+  const object = useMemo(() => makeObject_Primitives(), []);
+  identity(object);
+  return object;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{}],
+};
+
+```
+
+## Code
+
+```javascript
+// @enablePreserveExistingMemoizationGuarantees
+import { useMemo, unstable_useMemoCache as useMemoCache } from "react";
+import { identity, makeObject_Primitives, mutate } from "shared-runtime";
+
+function Component(props) {
+  const $ = useMemoCache(1);
+  let t15;
+  let t0;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    t0 = makeObject_Primitives();
+    $[0] = t0;
+  } else {
+    t0 = $[0];
+  }
+  t15 = t0;
+  const object = t15;
+  identity(object);
+  return object;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{}],
+};
+
+```
+      
+### Eval output
+(kind: ok) {"a":0,"b":"value1","c":true}
@@ -0,0 +1,14 @@
+// @enablePreserveExistingMemoizationGuarantees
+import { useMemo } from "react";
+import { identity, makeObject_Primitives, mutate } from "shared-runtime";
+
+function Component(props) {
+  const object = useMemo(() => makeObject_Primitives(), []);
+  identity(object);
+  return object;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{}],
+};