From e791552bd51e2de25ceb15ea5f2f3f6659cdc169 Mon Sep 17 00:00:00 2001 From: Joe Savona Date: Thu, 16 Oct 2025 10:42:53 -0700 Subject: [PATCH] [compiler][wip] Fix for ValidatePreserveMemo edge case w multiple return in useMemo MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- .../src/HIR/HIRBuilder.ts | 2 +- .../src/Inference/DropManualMemoization.ts | 81 ++++++++------- .../PruneAlwaysInvalidatingScopes.ts | 9 ++ ...xpected-consistent-destructuring.expect.md | 2 +- ...ropped-infer-always-invalidating.expect.md | 49 ---------- ...eMemo-dropped-infer-always-invalidating.ts | 21 ---- ...ed-if-it-was-always-invalidating.expect.md | 56 +++++++++++ ...preserved-if-it-was-always-invalidating.ts | 21 ++++ ...ed-useMemo-with-multiple-returns.expect.md | 98 +++++++++++++++++++ ...e-unnamed-useMemo-with-multiple-returns.js | 28 ++++++ 10 files changed, 260 insertions(+), 107 deletions(-) delete mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-memo-validation/error.false-positive-useMemo-dropped-infer-always-invalidating.expect.md delete mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-memo-validation/error.false-positive-useMemo-dropped-infer-always-invalidating.ts create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-memo-validation/repro-useMemo-counts-as-preserved-if-it-was-always-invalidating.expect.md create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-memo-validation/repro-useMemo-counts-as-preserved-if-it-was-always-invalidating.ts create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro-inline-unnamed-useMemo-with-multiple-returns.expect.md create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro-inline-unnamed-useMemo-with-multiple-returns.js diff --git a/compiler/packages/babel-plugin-react-compiler/src/HIR/HIRBuilder.ts b/compiler/packages/babel-plugin-react-compiler/src/HIR/HIRBuilder.ts index d3ecb2abdc..1a2a56a7f7 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/HIR/HIRBuilder.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/HIR/HIRBuilder.ts @@ -988,7 +988,7 @@ export function createTemporaryPlace( identifier: makeTemporaryIdentifier(env.nextIdentifierId, loc), reactive: false, effect: Effect.Unknown, - loc: GeneratedSource, + loc, }; } diff --git a/compiler/packages/babel-plugin-react-compiler/src/Inference/DropManualMemoization.ts b/compiler/packages/babel-plugin-react-compiler/src/Inference/DropManualMemoization.ts index d31cd1446b..496959a256 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/Inference/DropManualMemoization.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/Inference/DropManualMemoization.ts @@ -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 | null, memoDecl: Place, manualMemoId: number, -): [TInstruction, TInstruction] { +): [Array, Array] { + 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 | TInstruction - > = new Map(); + const queuedInserts: Map> = 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 | 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); } diff --git a/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PruneAlwaysInvalidatingScopes.ts b/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PruneAlwaysInvalidatingScopes.ts index b470271cf0..66b2bdcb56 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PruneAlwaysInvalidatingScopes.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PruneAlwaysInvalidatingScopes.ts @@ -77,6 +77,15 @@ class Transform extends ReactiveFunctionTransform { } break; } + case 'FinishMemoize': { + if ( + !withinScope && + this.alwaysInvalidatingValues.has(value.decl.identifier) + ) { + value.pruned = true; + } + break; + } } return {kind: 'keep'}; } diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.bug-invariant-expected-consistent-destructuring.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.bug-invariant-expected-consistent-destructuring.expect.md index a30ccffcd7..bd787b2f0c 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.bug-invariant-expected-consistent-destructuring.expect.md +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.bug-invariant-expected-consistent-destructuring.expect.md @@ -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 | diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-memo-validation/error.false-positive-useMemo-dropped-infer-always-invalidating.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-memo-validation/error.false-positive-useMemo-dropped-infer-always-invalidating.expect.md deleted file mode 100644 index aead04aa95..0000000000 --- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-memo-validation/error.false-positive-useMemo-dropped-infer-always-invalidating.expect.md +++ /dev/null @@ -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 = { -``` - - \ No newline at end of file diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-memo-validation/error.false-positive-useMemo-dropped-infer-always-invalidating.ts b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-memo-validation/error.false-positive-useMemo-dropped-infer-always-invalidating.ts deleted file mode 100644 index 3265182c51..0000000000 --- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-memo-validation/error.false-positive-useMemo-dropped-infer-always-invalidating.ts +++ /dev/null @@ -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: [{}], -}; diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-memo-validation/repro-useMemo-counts-as-preserved-if-it-was-always-invalidating.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-memo-validation/repro-useMemo-counts-as-preserved-if-it-was-always-invalidating.expect.md new file mode 100644 index 0000000000..83c93faf89 --- /dev/null +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-memo-validation/repro-useMemo-counts-as-preserved-if-it-was-always-invalidating.expect.md @@ -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) [[{}]] \ No newline at end of file diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-memo-validation/repro-useMemo-counts-as-preserved-if-it-was-always-invalidating.ts b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-memo-validation/repro-useMemo-counts-as-preserved-if-it-was-always-invalidating.ts new file mode 100644 index 0000000000..28689b141f --- /dev/null +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-memo-validation/repro-useMemo-counts-as-preserved-if-it-was-always-invalidating.ts @@ -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: [{}], +}; diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro-inline-unnamed-useMemo-with-multiple-returns.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro-inline-unnamed-useMemo-with-multiple-returns.expect.md new file mode 100644 index 0000000000..bdec6f48d4 --- /dev/null +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro-inline-unnamed-useMemo-with-multiple-returns.expect.md @@ -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
{campaign.name}
; +} + +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 =
{campaign.name}
; + $[0] = campaign.name; + $[1] = t0; + } else { + t0 = $[1]; + } + return t0; +} + +export const FIXTURE_ENTRYPOINT = { + fn: Component, + params: [{}], +}; + +``` + +### Eval output +(kind: ok)
campaign
\ No newline at end of file diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro-inline-unnamed-useMemo-with-multiple-returns.js b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro-inline-unnamed-useMemo-with-multiple-returns.js new file mode 100644 index 0000000000..2c939ccce2 --- /dev/null +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro-inline-unnamed-useMemo-with-multiple-returns.js @@ -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
{campaign.name}
; +} + +export const FIXTURE_ENTRYPOINT = { + fn: Component, + params: [{}], +};