diff --git a/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts b/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts index 762f1a4112a..f5aa5fa643e 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts +++ b/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", diff --git a/compiler/packages/babel-plugin-react-compiler/src/HIR/Environment.ts b/compiler/packages/babel-plugin-react-compiler/src/HIR/Environment.ts index 277921d5e25..63a3f9502a2 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/HIR/Environment.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/HIR/Environment.ts @@ -182,7 +182,9 @@ 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 +458,31 @@ export function parseConfigPragma(pragma: string): EnvironmentConfig { continue; } + if ( + key === "enablePreserveExistingManualUseMemo" && + (val === undefined || val === "true" || val === "scope") + ) { + maybeConfig[key] = "scope"; + continue; + } + + if (key === "enablePreserveExistingManualUseMemo" && val === "hook") { + maybeConfig[key] = "hook"; + continue; + } + + if ( + key === "enablePreserveExistingManualUseMemo" && + !(val === "false" || val === "off") + ) { + CompilerError.throwInvalidConfig({ + reason: `Invalid setting '${val}' for 'enablePreserveExistingManualUseMemo'. Valid settings are 'hook', 'scope', or 'off'.`, + description: null, + loc: null, + suggestions: null, + }); + } + if (typeof defaultConfig[key as keyof EnvironmentConfig] !== "boolean") { // skip parsing non-boolean properties continue; diff --git a/compiler/packages/babel-plugin-react-compiler/src/HIR/HIR.ts b/compiler/packages/babel-plugin-react-compiler/src/HIR/HIR.ts index b79414de031..d4703e55cfa 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/HIR/HIR.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/HIR/HIR.ts @@ -1429,6 +1429,8 @@ export type ReactiveScope = { merged: Set; loc: SourceLocation; + + source: boolean; }; export type ReactiveScopeDependencies = Set; 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 932cb4cc806..615244b8728 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/Inference/DropManualMemoization.ts +++ b/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 = { diff --git a/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/InferReactiveScopeVariables.ts b/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/InferReactiveScopeVariables.ts index 2c9e67646b1..2d1bdfbdb70 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/InferReactiveScopeVariables.ts +++ b/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,10 @@ 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,15 +190,20 @@ 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": { - return doesPatternContainSpreadElement(value.lvalue.pattern); + return ( + doesPatternContainSpreadElement(value.lvalue.pattern) || conservative + ); } case "PostfixUpdate": case "PrefixUpdate": - case "Await": case "DeclareLocal": case "DeclareContext": case "StoreLocal": @@ -205,26 +214,31 @@ function mayAllocate(env: Environment, instruction: Instruction): boolean { case "LoadContext": case "StoreContext": case "PropertyDelete": - case "ComputedLoad": case "ComputedDelete": case "JSXText": case "TemplateLiteral": case "Primitive": case "GetIterator": case "IteratorNext": - case "NextPropertyOf": case "Debugger": case "StartMemoize": case "FinishMemoize": case "UnaryExpression": case "BinaryExpression": - case "PropertyLoad": case "StoreGlobal": { return false; } + case "PropertyLoad": + case "NextPropertyOf": + case "ComputedLoad": + case "Await": { + return conservative; + } 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 +263,82 @@ function mayAllocate(env: Environment, instruction: Instruction): boolean { } } +export function collectMutableOperands( + fn: HIRFunction, + instr: Instruction, + conservative: boolean +): Array { + const operands: Array = []; + 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 { @@ -276,74 +366,7 @@ export function findDisjointMutableValues( } for (const instr of block.instructions) { - const operands: Array = []; - 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); } diff --git a/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/MemoizeExistingUseMemos.ts b/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/MemoizeExistingUseMemos.ts new file mode 100644 index 00000000000..a1c4b77f9af --- /dev/null +++ b/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/MemoizeExistingUseMemos.ts @@ -0,0 +1,108 @@ +/** + * 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(): number { + 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 +): void { + 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 + ): void { + 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); + } +} diff --git a/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/MergeReactiveScopesThatInvalidateTogether.ts b/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/MergeReactiveScopesThatInvalidateTogether.ts index e26fca13538..61378689aee 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/MergeReactiveScopesThatInvalidateTogether.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/MergeReactiveScopesThatInvalidateTogether.ts @@ -15,7 +15,6 @@ import { ReactiveFunction, ReactiveScope, ReactiveScopeBlock, - ReactiveScopeDependencies, ReactiveScopeDependency, ReactiveStatement, Type, @@ -109,7 +108,7 @@ class FindLastUsageVisitor extends ReactiveFunctionVisitor { } } -class Transform extends ReactiveFunctionTransform { +class Transform extends ReactiveFunctionTransform { lastUsage: Map; constructor(lastUsage: Map) { @@ -119,12 +118,13 @@ class Transform extends ReactiveFunctionTransform { - 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 { @@ -132,10 +132,7 @@ class Transform extends ReactiveFunctionTransform state.has(identifier.id) ); - if (hasMemoizedOutput) { + if (hasMemoizedOutput || scopeBlock.scope.source) { return { kind: "keep" }; } else { this.prunedScopes.add(scopeBlock.scope.id); diff --git a/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PruneUnusedScopes.ts b/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PruneUnusedScopes.ts index 11cc928874e..5e4ba867e43 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PruneUnusedScopes.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PruneUnusedScopes.ts @@ -43,6 +43,7 @@ class Transform extends ReactiveFunctionTransform { this.visitScope(scopeBlock, scopeState); if ( !scopeState.hasReturnStatement && + !scopeBlock.scope.source && scopeBlock.scope.reassignments.size === 0 && (scopeBlock.scope.declarations.size === 0 || /* diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useMemo-preserve-non-idempotent.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useMemo-preserve-non-idempotent.expect.md new file mode 100644 index 00000000000..deb36c578ba --- /dev/null +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useMemo-preserve-non-idempotent.expect.md @@ -0,0 +1,77 @@ + +## Input + +```javascript +// @enablePreserveExistingManualUseMemo +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"; // @enablePreserveExistingManualUseMemo +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 \ No newline at end of file diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useMemo-preserve-non-idempotent.js b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useMemo-preserve-non-idempotent.js new file mode 100644 index 00000000000..63a5be8eee1 --- /dev/null +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useMemo-preserve-non-idempotent.js @@ -0,0 +1,20 @@ +// @enablePreserveExistingManualUseMemo +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 }], +}; diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useMemo-simple-preserved.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useMemo-simple-preserved.expect.md index 6c813c27a68..0fb5bd59a47 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useMemo-simple-preserved.expect.md +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useMemo-simple-preserved.expect.md @@ -25,31 +25,26 @@ import { c as _c } from "react/compiler-runtime"; // @enablePreserveExistingManu import { useMemo } from "react"; function Component(t0) { - const $ = _c(5); + const $ = _c(4); const { a } = t0; let t1; + let t2; if ($[0] !== a) { - t1 = () => [a]; + t2 = [a]; $[0] = a; - $[1] = t1; - } else { - t1 = $[1]; - } - let t2; - if ($[2] === Symbol.for("react.memo_cache_sentinel")) { - t2 = []; - $[2] = t2; + $[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 =
{x}
; - $[3] = x; - $[4] = t3; + $[2] = x; + $[3] = t3; } else { - t3 = $[4]; + t3 = $[3]; } return t3; } diff --git a/compiler/packages/snap/src/compiler.ts b/compiler/packages/snap/src/compiler.ts index 0848d51aec6..8678d7c5786 100644 --- a/compiler/packages/snap/src/compiler.ts +++ b/compiler/packages/snap/src/compiler.ts @@ -46,6 +46,7 @@ function makePluginOptions( // TODO(@mofeiZ) rewrite snap fixtures to @validatePreserveExistingMemo:false let validatePreserveExistingMemoizationGuarantees = false; let enableChangeDetectionForDebugging = null; + let enablePreserveExistingManualUseMemo: "hook" | "scope" | null = null; let customMacros = null; if (firstLine.indexOf("@compilationMode(annotation)") !== -1) { @@ -130,6 +131,28 @@ function makePluginOptions( importSpecifierName: "$structuralCheck", }; } + + const useMemoMatch = /@enablePreserveExistingManualUseMemo:"([^"]+)"/.exec( + firstLine + ); + if ( + useMemoMatch && + (useMemoMatch[1] === "hook" || useMemoMatch[1] === "scope") + ) { + enablePreserveExistingManualUseMemo = useMemoMatch[1]; + } else if ( + useMemoMatch && + (useMemoMatch[1] === "false" || useMemoMatch[1] === "off") + ) { + enablePreserveExistingManualUseMemo = null; + } else if (useMemoMatch) { + throw new Error( + `Invalid setting '${useMemoMatch[1]}' for 'enablePreserveExistingManualUseMemo'. Valid settings are 'hook', 'scope', or 'off'.` + ); + } else if (firstLine.includes("@enablePreserveExistingManualUseMemo")) { + enablePreserveExistingManualUseMemo = "scope"; + } + const hookPatternMatch = /@hookPattern:"([^"]+)"/.exec(firstLine); if ( hookPatternMatch && @@ -207,6 +230,7 @@ function makePluginOptions( hookPattern, validatePreserveExistingMemoizationGuarantees, enableChangeDetectionForDebugging, + enablePreserveExistingManualUseMemo, }, compilationMode, logger,