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 762f1a4112..f5aa5fa643 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 277921d5e2..5b7e53a237 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,7 @@ 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 +456,29 @@ export function parseConfigPragma(pragma: string): EnvironmentConfig { continue; } + if ( + key === "enableChangeDetectionForDebugging" && + (val === undefined || val === "true") + ) { + maybeConfig[key] = { + source: "react-compiler-runtime", + importSpecifierName: "$structuralCheck", + }; + continue; + } + + if (key === "enablePreserveExistingManualUseMemoAsHook" && + (val === undefined || val === "true")) { + maybeConfig[key] = "hook" + continue; + } + + if (key === "enablePreserveExistingManualUseMemoAsScope" && + (val === undefined || val === "true")) { + maybeConfig[key] = "scope" + continue; + } + 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 b79414de03..d4703e55cf 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 932cb4cc80..615244b872 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 2c9e67646b..d0c6d41e92 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,7 @@ 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,7 +187,7 @@ 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": { @@ -224,7 +225,7 @@ function mayAllocate(env: Environment, instruction: Instruction): boolean { } 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 +250,78 @@ 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 +349,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 0000000000..f13ffa8017 --- /dev/null +++ b/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/MemoizeExistingUseMemos.ts @@ -0,0 +1,88 @@ +/** + * 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() { + 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) { + 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) { + 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 e26fca1353..f570e8c14c 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/MergeReactiveScopesThatInvalidateTogether.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/MergeReactiveScopesThatInvalidateTogether.ts @@ -109,7 +109,7 @@ class FindLastUsageVisitor extends ReactiveFunctionVisitor { } } -class Transform extends ReactiveFunctionTransform { +class Transform extends ReactiveFunctionTransform { lastUsage: Map; constructor(lastUsage: Map) { @@ -119,12 +119,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 { @@ -134,7 +135,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 11cc928874..5e4ba867e4 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 0000000000..60cf09a0d6 --- /dev/null +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useMemo-preserve-non-idempotent.expect.md @@ -0,0 +1,82 @@ + +## Input + +```javascript +// @enablePreserveExistingManualUseMemoAsScope +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"; // @enablePreserveExistingManualUseMemoAsScope +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 0000000000..5941bd98a7 --- /dev/null +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useMemo-preserve-non-idempotent.js @@ -0,0 +1,25 @@ +// @enablePreserveExistingManualUseMemoAsScope +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 6c813c27a6..e7e88d2598 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 @@ -2,7 +2,7 @@ ## Input ```javascript -// @enablePreserveExistingManualUseMemo +// @enablePreserveExistingManualUseMemoAsScope import { useMemo } from "react"; function Component({ a }) { @@ -21,35 +21,30 @@ export const FIXTURE_ENTRYPOINT = { ## Code ```javascript -import { c as _c } from "react/compiler-runtime"; // @enablePreserveExistingManualUseMemo +import { c as _c } from "react/compiler-runtime"; // @enablePreserveExistingManualUseMemoAsScope import { useMemo } from "react"; function Component(t0) { - const $ = _c(5); + const $ = _c(4); const { a } = t0; let t1; - if ($[0] !== a) { - t1 = () => [a]; - $[0] = a; - $[1] = t1; - } else { - t1 = $[1]; - } let t2; - if ($[2] === Symbol.for("react.memo_cache_sentinel")) { - t2 = []; - $[2] = t2; + if ($[0] !== a) { + t2 = [a]; + $[0] = a; + $[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/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useMemo-simple-preserved.js b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useMemo-simple-preserved.js index a5731f2f09..52a88490f2 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useMemo-simple-preserved.js +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useMemo-simple-preserved.js @@ -1,4 +1,4 @@ -// @enablePreserveExistingManualUseMemo +// @enablePreserveExistingManualUseMemoAsScope import { useMemo } from "react"; function Component({ a }) { diff --git a/compiler/packages/snap/src/compiler.ts b/compiler/packages/snap/src/compiler.ts index 0848d51aec..534e239195 100644 --- a/compiler/packages/snap/src/compiler.ts +++ b/compiler/packages/snap/src/compiler.ts @@ -45,7 +45,8 @@ function makePluginOptions( let hookPattern: string | null = null; // TODO(@mofeiZ) rewrite snap fixtures to @validatePreserveExistingMemo:false let validatePreserveExistingMemoizationGuarantees = false; - let enableChangeDetectionForDebugging = null; + let enableChangeDetectionForDebugging= null; + let enablePreserveExistingManualUseMemo: "hook" | "scope" | null = null; let customMacros = null; if (firstLine.indexOf("@compilationMode(annotation)") !== -1) { @@ -130,6 +131,15 @@ function makePluginOptions( importSpecifierName: "$structuralCheck", }; } + if (firstLine.includes("@enablePreserveExistingManualUseMemoAsHook")) { + enablePreserveExistingManualUseMemo = "hook" + } else if (firstLine.includes("@enablePreserveExistingManualUseMemoAsScope")) { + enablePreserveExistingManualUseMemo = "scope" + } else if (firstLine.includes("@enablePreserveExistingManualUseMemo")) { + throw new Error( + 'Use either @enablePreserveExistingManualUseMemoAsScope or @enablePreserveExistingManualUseMemoAsHook' + ); + } const hookPatternMatch = /@hookPattern:"([^"]+)"/.exec(firstLine); if ( hookPatternMatch && @@ -207,6 +217,7 @@ function makePluginOptions( hookPattern, validatePreserveExistingMemoizationGuarantees, enableChangeDetectionForDebugging, + enablePreserveExistingManualUseMemo, }, compilationMode, logger,