[compiler][wip] Fix for ValidatePreserveMemo edge case w multiple return in useMemo

Potential fix for #34750. With a useMemo with multiple returns we can sometimes end up thinking the FinishMemoize value wasn't memoized, even though it definitely is. But the same case works if you rewrite to have a let binding, assign to the let instead of returning, and then have a single return with the let — which is the exact shape we convert to when we inline the useMemo! The only difference is that when we inline, we end up with an extra LoadLocal. Adding an extra LoadLocal "fixes" the issue but this is definitely a hack. However, along the way this helped me find that we also have a bug in PruneAlwaysInvalidatingScopes, where we were forgetting to set FinishMemoize instructions to pruned if their declaration always invalidates. Putting up for feedback, not necessarily expecting to land as-is.
2025-11-01 09:12:30 +00:00 · 2025-10-16 10:42:53 -07:00
parent d79680bfa1
commit e791552bd5
10 changed files with 260 additions and 107 deletions
@@ -988,7 +988,7 @@ export function createTemporaryPlace(
    identifier: makeTemporaryIdentifier(env.nextIdentifierId, loc),
    reactive: false,
    effect: Effect.Unknown,
-    loc: GeneratedSource,
+    loc,
  };
 }

@@ -11,7 +11,6 @@ import {
  CallExpression,
  Effect,
  Environment,
-  FinishMemoize,
  FunctionExpression,
  HIRFunction,
  IdentifierId,
@@ -25,7 +24,6 @@ import {
  Place,
  PropertyLoad,
  SpreadPattern,
-  StartMemoize,
  TInstruction,
  getHookKindForType,
  makeInstructionId,
@@ -184,36 +182,52 @@ function makeManualMemoizationMarkers(
  depsList: Array<ManualMemoDependency> | null,
  memoDecl: Place,
  manualMemoId: number,
-): [TInstruction<StartMemoize>, TInstruction<FinishMemoize>] {
+): [Array<Instruction>, Array<Instruction>] {
+  const temp = createTemporaryPlace(env, memoDecl.loc);
  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,
+    [
+      {
+        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,
+        },
+        effects: null,
        loc: fnExpr.loc,
      },
-      effects: null,
-      loc: fnExpr.loc,
-    },
-    {
-      id: makeInstructionId(0),
-      lvalue: createTemporaryPlace(env, fnExpr.loc),
-      value: {
-        kind: 'FinishMemoize',
-        manualMemoId,
-        decl: {...memoDecl},
-        loc: fnExpr.loc,
+    ],
+    [
+      {
+        id: makeInstructionId(0),
+        lvalue: {...temp},
+        value: {
+          kind: 'LoadLocal',
+          place: {...memoDecl},
+          loc: memoDecl.loc,
+        },
+        effects: null,
+        loc: memoDecl.loc,
      },
-      effects: null,
-      loc: fnExpr.loc,
-    },
+      {
+        id: makeInstructionId(0),
+        lvalue: createTemporaryPlace(env, memoDecl.loc),
+        value: {
+          kind: 'FinishMemoize',
+          manualMemoId,
+          decl: {...temp},
+          loc: memoDecl.loc,
+        },
+        effects: null,
+        loc: memoDecl.loc,
+      },
+    ],
  ];
 }

@@ -409,10 +423,7 @@ export function dropManualMemoization(
   *   LoadLocal fnArg
   * - (if validation is enabled) collect manual memoization markers
   */
-  const queuedInserts: Map<
-    InstructionId,
-    TInstruction<StartMemoize> | TInstruction<FinishMemoize>
-  > = new Map();
+  const queuedInserts: Map<InstructionId, Array<Instruction>> = new Map();
  for (const [_, block] of func.body.blocks) {
    for (let i = 0; i < block.instructions.length; i++) {
      const instr = block.instructions[i]!;
@@ -557,11 +568,11 @@ export function dropManualMemoization(
      let nextInstructions: Array<Instruction> | null = null;
      for (let i = 0; i < block.instructions.length; i++) {
        const instr = block.instructions[i];
-        const insertInstr = queuedInserts.get(instr.id);
-        if (insertInstr != null) {
+        const insertInstructions = queuedInserts.get(instr.id);
+        if (insertInstructions != null) {
          nextInstructions = nextInstructions ?? block.instructions.slice(0, i);
          nextInstructions.push(instr);
-          nextInstructions.push(insertInstr);
+          nextInstructions.push(...insertInstructions);
        } else if (nextInstructions != null) {
          nextInstructions.push(instr);
        }
@@ -77,6 +77,15 @@ class Transform extends ReactiveFunctionTransform<boolean> {
        }
        break;
      }
+      case 'FinishMemoize': {
+        if (
+          !withinScope &&
+          this.alwaysInvalidatingValues.has(value.decl.identifier)
+        ) {
+          value.pruned = true;
+        }
+        break;
+      }
    }
    return {kind: 'keep'};
  }
@@ -29,7 +29,7 @@ Found 1 error:

 Invariant: Expected consistent kind for destructuring

-Other places were `Reassign` but 'mutate? #t8$46[7:9]{reactive}' is const.
+Other places were `Reassign` but 'mutate? #t8$47[7:9]{reactive}' is const.

 error.bug-invariant-expected-consistent-destructuring.ts:9:9
   7 |
@@ -1,49 +0,0 @@
-
-## Input
-
-```javascript
-// @validatePreserveExistingMemoizationGuarantees
-
-import {useMemo} from 'react';
-import {useHook} from 'shared-runtime';
-
-// useMemo values may not be memoized in Forget output if we
-// infer that their deps always invalidate.
-// This is technically a false positive as the useMemo in source
-// was effectively a no-op
-function useFoo(props) {
-  const x = [];
-  useHook();
-  x.push(props);
-
-  return useMemo(() => [x], [x]);
-}
-
-export const FIXTURE_ENTRYPOINT = {
-  fn: useFoo,
-  params: [{}],
-};
-
-```
-
-
-## Error
-
-```
-Found 1 error:
-
-Compilation Skipped: Existing memoization could not be preserved
-
-React Compiler has skipped optimizing this component because the existing manual memoization could not be preserved. This value was memoized in source but not in compilation output.
-
-error.false-positive-useMemo-dropped-infer-always-invalidating.ts:15:9
-  13 |   x.push(props);
-  14 |
-> 15 |   return useMemo(() => [x], [x]);
-     |          ^^^^^^^^^^^^^^^^^^^^^^^ Could not preserve existing memoization
-  16 | }
-  17 |
-  18 | export const FIXTURE_ENTRYPOINT = {
-```
-          
-      
@@ -1,21 +0,0 @@
-// @validatePreserveExistingMemoizationGuarantees
-
-import {useMemo} from 'react';
-import {useHook} from 'shared-runtime';
-
-// useMemo values may not be memoized in Forget output if we
-// infer that their deps always invalidate.
-// This is technically a false positive as the useMemo in source
-// was effectively a no-op
-function useFoo(props) {
-  const x = [];
-  useHook();
-  x.push(props);
-
-  return useMemo(() => [x], [x]);
-}
-
-export const FIXTURE_ENTRYPOINT = {
-  fn: useFoo,
-  params: [{}],
-};
@@ -0,0 +1,56 @@
+
+## Input
+
+```javascript
+// @validatePreserveExistingMemoizationGuarantees
+
+import {useMemo} from 'react';
+import {useHook} from 'shared-runtime';
+
+// If we can prove that a useMemo was ineffective because it would always invalidate,
+// then we shouldn't throw a "couldn't preserve existing memoization" error
+// TODO: consider reporting a separate error to the user for this case, if you're going
+// to memoize manually, then you probably want to know that it's a no-op
+function useFoo(props) {
+  const x = [];
+  useHook();
+  x.push(props);
+
+  return useMemo(() => [x], [x]);
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useFoo,
+  params: [{}],
+};
+
+```
+
+## Code
+
+```javascript
+// @validatePreserveExistingMemoizationGuarantees
+
+import { useMemo } from "react";
+import { useHook } from "shared-runtime";
+
+// If we can prove that a useMemo was ineffective because it would always invalidate,
+// then we shouldn't throw a "couldn't preserve existing memoization" error
+// TODO: consider reporting a separate error to the user for this case, if you're going
+// to memoize manually, then you probably want to know that it's a no-op
+function useFoo(props) {
+  const x = [];
+  useHook();
+  x.push(props);
+  return [x];
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useFoo,
+  params: [{}],
+};
+
+```
+      
+### Eval output
+(kind: ok) [[{}]]
@@ -0,0 +1,21 @@
+// @validatePreserveExistingMemoizationGuarantees
+
+import {useMemo} from 'react';
+import {useHook} from 'shared-runtime';
+
+// If we can prove that a useMemo was ineffective because it would always invalidate,
+// then we shouldn't throw a "couldn't preserve existing memoization" error
+// TODO: consider reporting a separate error to the user for this case, if you're going
+// to memoize manually, then you probably want to know that it's a no-op
+function useFoo(props) {
+  const x = [];
+  useHook();
+  x.push(props);
+
+  return useMemo(() => [x], [x]);
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useFoo,
+  params: [{}],
+};
@@ -0,0 +1,98 @@
+
+## Input
+
+```javascript
+import {useMemo} from 'react';
+import {identity, useIdentity} from 'shared-runtime';
+
+// Adapted from https://github.com/facebook/react/issues/34750
+function useLocalCampaignBySlug(slug: string) {
+  const campaigns = useIdentity({a: {slug: 'a', name: 'campaign'}});
+  // The useMemo result is never assigned to a local so we did not previously ensure
+  // that there was a variable declaration for it when promoting the result temporary
+  return useMemo(() => {
+    for (const id of Object.keys(campaigns)) {
+      const campaign = campaigns[id];
+      if (campaign.slug === slug) {
+        return identity(campaign);
+      }
+    }
+    return null;
+  }, [campaigns, slug]);
+}
+
+function Component() {
+  const campaign = useLocalCampaignBySlug('a');
+  return <div>{campaign.name}</div>;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{}],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime";
+import { useMemo } from "react";
+import { identity, useIdentity } from "shared-runtime";
+
+// Adapted from https://github.com/facebook/react/issues/34750
+function useLocalCampaignBySlug(slug) {
+  const $ = _c(4);
+  let t0;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    t0 = { a: { slug: "a", name: "campaign" } };
+    $[0] = t0;
+  } else {
+    t0 = $[0];
+  }
+  const campaigns = useIdentity(t0);
+  let t1;
+  if ($[1] !== campaigns || $[2] !== slug) {
+    bb0: {
+      for (const id of Object.keys(campaigns)) {
+        const campaign = campaigns[id];
+        if (campaign.slug === slug) {
+          t1 = identity(campaign);
+          break bb0;
+        }
+      }
+
+      t1 = null;
+    }
+    $[1] = campaigns;
+    $[2] = slug;
+    $[3] = t1;
+  } else {
+    t1 = $[3];
+  }
+  return t1;
+}
+
+function Component() {
+  const $ = _c(2);
+  const campaign = useLocalCampaignBySlug("a");
+  let t0;
+  if ($[0] !== campaign.name) {
+    t0 = <div>{campaign.name}</div>;
+    $[0] = campaign.name;
+    $[1] = t0;
+  } else {
+    t0 = $[1];
+  }
+  return t0;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{}],
+};
+
+```
+      
+### Eval output
+(kind: ok) <div>campaign</div>
@@ -0,0 +1,28 @@
+import {useMemo} from 'react';
+import {identity, useIdentity} from 'shared-runtime';
+
+// Adapted from https://github.com/facebook/react/issues/34750
+function useLocalCampaignBySlug(slug: string) {
+  const campaigns = useIdentity({a: {slug: 'a', name: 'campaign'}});
+  // The useMemo result is never assigned to a local so we did not previously ensure
+  // that there was a variable declaration for it when promoting the result temporary
+  return useMemo(() => {
+    for (const id of Object.keys(campaigns)) {
+      const campaign = campaigns[id];
+      if (campaign.slug === slug) {
+        return identity(campaign);
+      }
+    }
+    return null;
+  }, [campaigns, slug]);
+}
+
+function Component() {
+  const campaign = useLocalCampaignBySlug('a');
+  return <div>{campaign.name}</div>;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{}],
+};