diff --git a/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PruneNonEscapingScopes.ts b/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PruneNonEscapingScopes.ts index 2951d50485300..b2e91fa302728 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PruneNonEscapingScopes.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PruneNonEscapingScopes.ts @@ -9,6 +9,7 @@ import {CompilerError} from '../CompilerError'; import { DeclarationId, Environment, + Identifier, InstructionId, Pattern, Place, @@ -24,7 +25,7 @@ import { isMutableEffect, } from '../HIR'; import {getFunctionCallSignature} from '../Inference/InferReferenceEffects'; -import {assertExhaustive} from '../Utils/utils'; +import {assertExhaustive, getOrInsertDefault} from '../Utils/utils'; import {getPlaceScope} from './BuildReactiveBlocks'; import { ReactiveFunctionTransform, @@ -935,6 +936,11 @@ class PruneScopesTransform extends ReactiveFunctionTransform< Set > { prunedScopes: Set = new Set(); + /** + * Track reassignments so we can correctly set `pruned` flags for + * inlined useMemos. + */ + reassignments: Map> = new Map(); override transformScope( scopeBlock: ReactiveScopeBlock, @@ -977,24 +983,45 @@ class PruneScopesTransform extends ReactiveFunctionTransform< } } + /** + * If we pruned the scope for a non-escaping value, we know it doesn't + * need to be memoized. Remove associated `Memoize` instructions so that + * we don't report false positives on "missing" memoization of these values. + */ override transformInstruction( instruction: ReactiveInstruction, state: Set, ): Transformed { this.traverseInstruction(instruction, state); - /** - * If we pruned the scope for a non-escaping value, we know it doesn't - * need to be memoized. Remove associated `Memoize` instructions so that - * we don't report false positives on "missing" memoization of these values. - */ - if (instruction.value.kind === 'FinishMemoize') { - const identifier = instruction.value.decl.identifier; + const value = instruction.value; + if (value.kind === 'StoreLocal' && value.lvalue.kind === 'Reassign') { + const ids = getOrInsertDefault( + this.reassignments, + value.lvalue.place.identifier.declarationId, + new Set(), + ); + ids.add(value.value.identifier); + } else if (value.kind === 'FinishMemoize') { + let decls; + if (value.decl.identifier.scope == null) { + /** + * If the manual memo was a useMemo that got inlined, iterate through + * all reassignments to the iife temporary to ensure they're memoized. + */ + decls = this.reassignments.get(value.decl.identifier.declarationId) ?? [ + value.decl.identifier, + ]; + } else { + decls = [value.decl.identifier]; + } + if ( - identifier.scope !== null && - this.prunedScopes.has(identifier.scope.id) + [...decls].every( + decl => decl.scope == null || this.prunedScopes.has(decl.scope.id), + ) ) { - instruction.value.pruned = true; + value.pruned = true; } } diff --git a/compiler/packages/babel-plugin-react-compiler/src/Validation/ValidatePreservedManualMemoization.ts b/compiler/packages/babel-plugin-react-compiler/src/Validation/ValidatePreservedManualMemoization.ts index bbe0f17f59425..50d1c986cea5a 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/Validation/ValidatePreservedManualMemoization.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/Validation/ValidatePreservedManualMemoization.ts @@ -30,6 +30,7 @@ import { ReactiveFunctionVisitor, visitReactiveFunction, } from '../ReactiveScopes/visitors'; +import {getOrInsertDefault} from '../Utils/utils'; /** * Validates that all explicit manual memoization (useMemo/useCallback) was accurately @@ -52,6 +53,16 @@ export function validatePreservedManualMemoization(fn: ReactiveFunction): void { const DEBUG = false; type ManualMemoBlockState = { + /** + * Tracks reassigned temporaries. + * This is necessary because useMemo calls are usually inlined. + * Inlining produces a `let` declaration, followed by reassignments + * to the newly declared variable (one per return statement). + * Since InferReactiveScopes does not merge scopes across reassigned + * variables (except in the case of a mutate-after-phi), we need to + * track reassignments to validate we're retaining manual memo. + */ + reassignments: Map>; // The source of the original memoization, used when reporting errors loc: SourceLocation; @@ -425,6 +436,18 @@ class Visitor extends ReactiveFunctionVisitor { */ this.recordTemporaries(instruction, state); const value = instruction.value; + if ( + value.kind === 'StoreLocal' && + value.lvalue.kind === 'Reassign' && + state.manualMemoState != null + ) { + const ids = getOrInsertDefault( + state.manualMemoState.reassignments, + value.lvalue.place.identifier.declarationId, + new Set(), + ); + ids.add(value.value.identifier); + } if (value.kind === 'StartMemoize') { let depsFromSource: Array | null = null; if (value.deps != null) { @@ -442,6 +465,7 @@ class Visitor extends ReactiveFunctionVisitor { decls: new Set(), depsFromSource, manualMemoId: value.manualMemoId, + reassignments: new Map(), }; /** @@ -491,20 +515,34 @@ class Visitor extends ReactiveFunctionVisitor { suggestions: null, }, ); + const reassignments = state.manualMemoState.reassignments; state.manualMemoState = null; if (!value.pruned) { for (const {identifier, loc} of eachInstructionValueOperand( value as InstructionValue, )) { - if (isUnmemoized(identifier, this.scopes)) { - state.errors.push({ - reason: - '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.', - description: null, - severity: ErrorSeverity.CannotPreserveMemoization, - loc, - suggestions: null, - }); + let decls; + if (identifier.scope == null) { + /** + * If the manual memo was a useMemo that got inlined, iterate through + * all reassignments to the iife temporary to ensure they're memoized. + */ + decls = reassignments.get(identifier.declarationId) ?? [identifier]; + } else { + decls = [identifier]; + } + + for (const identifier of decls) { + if (isUnmemoized(identifier, this.scopes)) { + state.errors.push({ + reason: + '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.', + description: null, + severity: ErrorSeverity.CannotPreserveMemoization, + loc, + suggestions: null, + }); + } } } } diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/bug-maybe-mutable-ref-memo-not-preserved.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/bug-maybe-mutable-ref-memo-not-preserved.expect.md deleted file mode 100644 index ce19a1dc4eac0..0000000000000 --- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/bug-maybe-mutable-ref-memo-not-preserved.expect.md +++ /dev/null @@ -1,45 +0,0 @@ - -## Input - -```javascript -// @validatePreserveExistingMemoizationGuarantees:true - -import {useRef, useMemo} from 'react'; -import {makeArray} from 'shared-runtime'; - -function useFoo() { - const r = useRef(); - return useMemo(() => makeArray(r), []); -} - -export const FIXTURE_ENTRYPOINT = { - fn: useFoo, - params: [], -}; - -``` - -## Code - -```javascript -// @validatePreserveExistingMemoizationGuarantees:true - -import { useRef, useMemo } from "react"; -import { makeArray } from "shared-runtime"; - -function useFoo() { - const r = useRef(); - let t0; - t0 = makeArray(r); - return t0; -} - -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/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 new file mode 100644 index 0000000000000..855058d25b685 --- /dev/null +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-memo-validation/error.false-positive-useMemo-dropped-infer-always-invalidating.expect.md @@ -0,0 +1,42 @@ + +## 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 + +``` + 13 | x.push(props); + 14 | +> 15 | return useMemo(() => [x], [x]); + | ^^^^^^^^^^^^^^^^^^^^^^^ CannotPreserveMemoization: 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. (15:15) + 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/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 similarity index 81% rename from compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-memo-validation/useMemo-dropped-infer-always-invalidating.ts rename to compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-memo-validation/error.false-positive-useMemo-dropped-infer-always-invalidating.ts index bb6f655a2c36e..3265182c51dfd 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-memo-validation/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 @@ -5,8 +5,8 @@ import {useHook} from 'shared-runtime'; // useMemo values may not be memoized in Forget output if we // infer that their deps always invalidate. -// This is still correct as the useMemo in source was effectively -// a no-op already. +// This is technically a false positive as the useMemo in source +// was effectively a no-op function useFoo(props) { const x = []; useHook(); diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-memo-validation/error.maybe-mutable-ref-not-preserved.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-memo-validation/error.maybe-mutable-ref-not-preserved.expect.md new file mode 100644 index 0000000000000..61fe8b8deb7ee --- /dev/null +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-memo-validation/error.maybe-mutable-ref-not-preserved.expect.md @@ -0,0 +1,35 @@ + +## Input + +```javascript +// @validatePreserveExistingMemoizationGuarantees:true + +import {useRef, useMemo} from 'react'; +import {makeArray} from 'shared-runtime'; + +function useFoo() { + const r = useRef(); + return useMemo(() => makeArray(r), []); +} + +export const FIXTURE_ENTRYPOINT = { + fn: useFoo, + params: [], +}; + +``` + + +## Error + +``` + 6 | function useFoo() { + 7 | const r = useRef(); +> 8 | return useMemo(() => makeArray(r), []); + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ CannotPreserveMemoization: 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. (8:8) + 9 | } + 10 | + 11 | export const FIXTURE_ENTRYPOINT = { +``` + + \ No newline at end of file diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/bug-maybe-mutable-ref-memo-not-preserved.ts b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-memo-validation/error.maybe-mutable-ref-not-preserved.ts similarity index 100% rename from compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/bug-maybe-mutable-ref-memo-not-preserved.ts rename to compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-memo-validation/error.maybe-mutable-ref-not-preserved.ts diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-memo-validation/prune-nonescaping-useMemo-mult-returns-primitive.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-memo-validation/prune-nonescaping-useMemo-mult-returns-primitive.expect.md new file mode 100644 index 0000000000000..8850f869d07be --- /dev/null +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-memo-validation/prune-nonescaping-useMemo-mult-returns-primitive.expect.md @@ -0,0 +1,52 @@ + +## Input + +```javascript +// @validatePreserveExistingMemoizationGuarantees + +import {useMemo} from 'react'; +import {identity} from 'shared-runtime'; + +function useFoo(cond) { + useMemo(() => { + if (cond) { + return 2; + } else { + return identity(5); + } + }, [cond]); +} + +export const FIXTURE_ENTRYPOINT = { + fn: useFoo, + params: [true], +}; + +``` + +## Code + +```javascript +// @validatePreserveExistingMemoizationGuarantees + +import { useMemo } from "react"; +import { identity } from "shared-runtime"; + +function useFoo(cond) { + let t0; + if (cond) { + t0 = 2; + } else { + t0 = identity(5); + } +} + +export const FIXTURE_ENTRYPOINT = { + fn: useFoo, + params: [true], +}; + +``` + +### 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/prune-nonescaping-useMemo-mult-returns-primitive.ts b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-memo-validation/prune-nonescaping-useMemo-mult-returns-primitive.ts new file mode 100644 index 0000000000000..a79c84823ccc2 --- /dev/null +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-memo-validation/prune-nonescaping-useMemo-mult-returns-primitive.ts @@ -0,0 +1,19 @@ +// @validatePreserveExistingMemoizationGuarantees + +import {useMemo} from 'react'; +import {identity} from 'shared-runtime'; + +function useFoo(cond) { + useMemo(() => { + if (cond) { + return 2; + } else { + return identity(5); + } + }, [cond]); +} + +export const FIXTURE_ENTRYPOINT = { + fn: useFoo, + params: [true], +}; diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-memo-validation/prune-nonescaping-useMemo-mult-returns.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-memo-validation/prune-nonescaping-useMemo-mult-returns.expect.md new file mode 100644 index 0000000000000..af959b9865bcb --- /dev/null +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-memo-validation/prune-nonescaping-useMemo-mult-returns.expect.md @@ -0,0 +1,52 @@ + +## Input + +```javascript +// @validatePreserveExistingMemoizationGuarantees + +import {useMemo} from 'react'; +import {identity} from 'shared-runtime'; + +function useFoo(cond) { + useMemo(() => { + if (cond) { + return identity(10); + } else { + return identity(5); + } + }, [cond]); +} + +export const FIXTURE_ENTRYPOINT = { + fn: useFoo, + params: [true], +}; + +``` + +## Code + +```javascript +// @validatePreserveExistingMemoizationGuarantees + +import { useMemo } from "react"; +import { identity } from "shared-runtime"; + +function useFoo(cond) { + let t0; + if (cond) { + t0 = identity(10); + } else { + t0 = identity(5); + } +} + +export const FIXTURE_ENTRYPOINT = { + fn: useFoo, + params: [true], +}; + +``` + +### 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/prune-nonescaping-useMemo-mult-returns.ts b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-memo-validation/prune-nonescaping-useMemo-mult-returns.ts new file mode 100644 index 0000000000000..3772934106115 --- /dev/null +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-memo-validation/prune-nonescaping-useMemo-mult-returns.ts @@ -0,0 +1,19 @@ +// @validatePreserveExistingMemoizationGuarantees + +import {useMemo} from 'react'; +import {identity} from 'shared-runtime'; + +function useFoo(cond) { + useMemo(() => { + if (cond) { + return identity(10); + } else { + return identity(5); + } + }, [cond]); +} + +export const FIXTURE_ENTRYPOINT = { + fn: useFoo, + params: [true], +}; diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-memo-validation/prune-nonescaping-useMemo.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-memo-validation/prune-nonescaping-useMemo.expect.md new file mode 100644 index 0000000000000..27759f5db30fd --- /dev/null +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-memo-validation/prune-nonescaping-useMemo.expect.md @@ -0,0 +1,50 @@ + +## Input + +```javascript +// @validatePreserveExistingMemoizationGuarantees + +import {useMemo} from 'react'; +import {identity} from 'shared-runtime'; + +/** + * This is technically a false positive, although it makes sense + * to bailout as source code might be doing something sketchy. + */ +function useFoo(x) { + useMemo(() => identity(x), [x]); +} + +export const FIXTURE_ENTRYPOINT = { + fn: useFoo, + params: [2], +}; + +``` + +## Code + +```javascript +// @validatePreserveExistingMemoizationGuarantees + +import { useMemo } from "react"; +import { identity } from "shared-runtime"; + +/** + * This is technically a false positive, although it makes sense + * to bailout as source code might be doing something sketchy. + */ +function useFoo(x) { + let t0; + t0 = identity(x); +} + +export const FIXTURE_ENTRYPOINT = { + fn: useFoo, + params: [2], +}; + +``` + +### 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/prune-nonescaping-useMemo.ts b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-memo-validation/prune-nonescaping-useMemo.ts new file mode 100644 index 0000000000000..95059470add65 --- /dev/null +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-memo-validation/prune-nonescaping-useMemo.ts @@ -0,0 +1,17 @@ +// @validatePreserveExistingMemoizationGuarantees + +import {useMemo} from 'react'; +import {identity} from 'shared-runtime'; + +/** + * This is technically a false positive, although it makes sense + * to bailout as source code might be doing something sketchy. + */ +function useFoo(x) { + useMemo(() => identity(x), [x]); +} + +export const FIXTURE_ENTRYPOINT = { + fn: useFoo, + params: [2], +}; diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-memo-validation/useMemo-dropped-infer-always-invalidating.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-memo-validation/useMemo-dropped-infer-always-invalidating.expect.md deleted file mode 100644 index f4e77ed2a7784..0000000000000 --- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-memo-validation/useMemo-dropped-infer-always-invalidating.expect.md +++ /dev/null @@ -1,58 +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 still correct as the useMemo in source was effectively -// a no-op already. -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"; - -// useMemo values may not be memoized in Forget output if we -// infer that their deps always invalidate. -// This is still correct as the useMemo in source was effectively -// a no-op already. -function useFoo(props) { - const x = []; - useHook(); - x.push(props); - let t0; - t0 = [x]; - return t0; -} - -export const FIXTURE_ENTRYPOINT = { - fn: useFoo, - params: [{}], -}; - -``` - -### Eval output -(kind: ok) [[{}]] \ No newline at end of file