From b7325ff1f06619750dd0e78a12c1fa860d3b644f Mon Sep 17 00:00:00 2001 From: Mike Vitousek Date: Mon, 1 Jul 2024 18:02:19 -0700 Subject: [PATCH] [compiler] Use dependencies from source for useMemo scopes Summary: This modified the behavior of the compiler when preserving source-level useMemos as reactive scopes to use the depencies from the source as well. This accounts for the possibility that the useMemo in the source is more general than is required by local code, but is necessary due to rules of react violations. With this change, ideally the disableMemoization mode will behave exactly like the uncompiled code. [ghstack-poisoned] --- .../src/HIR/HIR.ts | 7 +- .../src/HIR/PrintHIR.ts | 4 +- .../src/HIR/visitors.ts | 10 ++- .../src/Inference/DropManualMemoization.ts | 90 ++++++++++++++----- .../ReactiveScopes/MemoizeExistingUseMemos.ts | 19 +++- .../PropagateScopeDependencies.ts | 4 +- .../PruneInitializationDependencies.ts | 2 +- .../PruneNonReactiveDependencies.ts | 2 +- .../ValidatePreservedManualMemoization.ts | 13 ++- .../fixtures/compiler/nomemo-path.expect.md | 66 ++++++++++++++ .../fixtures/compiler/nomemo-path.js | 15 ++++ .../useMemo-simple-preserved-nomemo.expect.md | 17 ++-- .../useMemo-simple-preserved.expect.md | 17 ++-- 13 files changed, 211 insertions(+), 55 deletions(-) create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/nomemo-path.expect.md create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/nomemo-path.js 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 d4703e55cfa..6cfa10113f1 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/HIR/HIR.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/HIR/HIR.ts @@ -771,7 +771,12 @@ export type ManualMemoDependency = { kind: "NamedLocal"; value: Place; } - | { kind: "Global"; identifierName: string }; + | { + kind: "InlinedGlobal"; + value: Place; + name: string; + } + | { kind: "Global"; binding: LoadGlobal }; path: Array; }; diff --git a/compiler/packages/babel-plugin-react-compiler/src/HIR/PrintHIR.ts b/compiler/packages/babel-plugin-react-compiler/src/HIR/PrintHIR.ts index df7d2698f4b..9e5b8c610d4 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/HIR/PrintHIR.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/HIR/PrintHIR.ts @@ -843,7 +843,9 @@ export function printManualMemoDependency( ): string { let rootStr; if (val.root.kind === "Global") { - rootStr = val.root.identifierName; + rootStr = val.root.binding.binding.name; + } else if (val.root.kind === "InlinedGlobal") { + rootStr = `G(${val.root.name})`; } else { CompilerError.invariant(val.root.value.identifier.name?.kind === "named", { reason: "DepsValidation: expected named local variable in depslist", diff --git a/compiler/packages/babel-plugin-react-compiler/src/HIR/visitors.ts b/compiler/packages/babel-plugin-react-compiler/src/HIR/visitors.ts index beda6e4a205..91f07184b04 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/HIR/visitors.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/HIR/visitors.ts @@ -229,7 +229,10 @@ export function* eachInstructionValueOperand( case "StartMemoize": { if (instrValue.deps != null) { for (const dep of instrValue.deps) { - if (dep.root.kind === "NamedLocal") { + if ( + dep.root.kind === "NamedLocal" || + dep.root.kind === "InlinedGlobal" + ) { yield dep.root.value; } } @@ -554,7 +557,10 @@ export function mapInstructionValueOperands( case "StartMemoize": { if (instrValue.deps != null) { for (const dep of instrValue.deps) { - if (dep.root.kind === "NamedLocal") { + if ( + dep.root.kind === "NamedLocal" || + dep.root.kind === "InlinedGlobal" + ) { dep.root.value = fn(dep.root.value); } } 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 41df9e963ff..0798dc717ff 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/Inference/DropManualMemoization.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/Inference/DropManualMemoization.ts @@ -58,7 +58,7 @@ export function collectMaybeMemoDependencies( return { root: { kind: "Global", - identifierName: value.binding.name, + binding: value, }, path: [], }; @@ -173,24 +173,53 @@ function makeManualMemoizationMarkers( env: Environment, depsList: Array | null, memoDecl: Place, - manualMemoId: number -): [TInstruction, TInstruction] { + manualMemoId: number, + isManualUseMemoEnabled: boolean +): [ + Array | TInstruction>, + TInstruction, +] { + let globals: Array | TInstruction> = + []; + for (const dep of depsList ?? []) { + if (dep.root.kind === "Global" && isManualUseMemoEnabled) { + const place = createTemporaryPlace(env, dep.root.binding.loc); + globals.push({ + id: makeInstructionId(0), + lvalue: place, + value: { + kind: "LoadGlobal", + binding: dep.root.binding.binding, + loc: dep.root.binding.loc, + }, + loc: dep.root.binding.loc, + }); + dep.root = { + kind: "InlinedGlobal", + value: place, + name: dep.root.binding.binding.name, + }; + } + } 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, + [ + ...globals, + { + 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, + }, loc: fnExpr.loc, }, - loc: fnExpr.loc, - }, + ], { id: makeInstructionId(0), lvalue: createTemporaryPlace(env, fnExpr.loc), @@ -333,12 +362,14 @@ function extractManualMemoizationArgs( * eg `React.useMemo()`. */ export function dropManualMemoization(func: HIRFunction): void { - const isValidationEnabled = + const isManualUseMemoEnabled = func.env.config.enablePreserveExistingManualUseMemo === "scope" || + func.env.config.enableChangeDetectionForDebugging != null || + func.env.config.disableMemoizationForDebugging; + const isValidationEnabled = func.env.config.validatePreserveExistingMemoizationGuarantees || func.env.config.enablePreserveExistingMemoizationGuarantees || - func.env.config.enableChangeDetectionForDebugging || - func.env.config.disableMemoizationForDebugging; + isManualUseMemoEnabled; const sidemap: IdentifierSidemap = { functions: new Map(), manualMemos: new Map(), @@ -359,7 +390,11 @@ export function dropManualMemoization(func: HIRFunction): void { */ const queuedInserts: Map< InstructionId, - TInstruction | TInstruction + Array< + | TInstruction + | TInstruction + | TInstruction + > > = new Map(); for (const [_, block] of func.body.blocks) { for (let i = 0; i < block.instructions.length; i++) { @@ -422,7 +457,8 @@ export function dropManualMemoization(func: HIRFunction): void { func.env, depsList, memoDecl, - nextManualMemoId++ + nextManualMemoId++, + isManualUseMemoEnabled ); /** @@ -439,7 +475,7 @@ export function dropManualMemoization(func: HIRFunction): void { * ``` */ queuedInserts.set(manualMemo.loadInstr.id, startMarker); - queuedInserts.set(instr.id, finishMarker); + queuedInserts.set(instr.id, [finishMarker]); } } } else { @@ -460,8 +496,16 @@ export function dropManualMemoization(func: HIRFunction): void { const insertInstr = queuedInserts.get(instr.id); if (insertInstr != null) { nextInstructions = nextInstructions ?? block.instructions.slice(0, i); + const postInstructions: Array = []; + insertInstr.forEach((instr) => { + if (instr.value.kind === "LoadGlobal") { + nextInstructions?.push(instr); + } else { + postInstructions.push(instr); + } + }); nextInstructions.push(instr); - nextInstructions.push(insertInstr); + postInstructions.forEach((instr) => nextInstructions?.push(instr)); } else if (nextInstructions != null) { nextInstructions.push(instr); } diff --git a/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/MemoizeExistingUseMemos.ts b/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/MemoizeExistingUseMemos.ts index a1c4b77f9af..8db9c29ff8b 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/MemoizeExistingUseMemos.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/MemoizeExistingUseMemos.ts @@ -12,6 +12,7 @@ import { Identifier, makeInstructionId, ReactiveScope, + ReactiveScopeDependencies, } from "../HIR"; import { eachTerminalSuccessor } from "../HIR/visitors"; import { @@ -30,7 +31,7 @@ function nextId(): number { type CurrentScope = | null - | { kind: "pending"; id: number } + | { kind: "pending"; deps: ReactiveScopeDependencies; id: number } | { kind: "available"; scope: ReactiveScope; id: number }; function visitBlock( @@ -75,7 +76,18 @@ function visitBlock( let currentScope = scope; for (const instruction of block.instructions) { if (instruction.value.kind === "StartMemoize") { - currentScope = { kind: "pending", id: nextId() }; + const deps: ReactiveScopeDependencies = new Set(); + for (const dep of instruction.value.deps ?? []) { + CompilerError.invariant(dep.root.kind !== "Global", { + reason: + "MemoizeExistingUseMemos: Globals should have been replaced with InlineGlobals", + loc: instruction.loc, + suggestions: null, + description: null, + }); + deps.add({ identifier: dep.root.value.identifier, path: dep.path }); + } + currentScope = { kind: "pending", id: nextId(), deps }; } else if (instruction.value.kind === "FinishMemoize") { currentScope = null; } else if (currentScope != null) { @@ -88,7 +100,7 @@ function visitBlock( scope: { id: fn.env.nextScopeId, range: { start: instruction.id, end: instruction.id }, - dependencies: new Set(), + dependencies: currentScope.deps, declarations: new Map(), reassignments: new Set(), earlyReturnValue: null, @@ -102,6 +114,7 @@ function visitBlock( } } } + 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/PropagateScopeDependencies.ts b/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PropagateScopeDependencies.ts index bc684ffe36f..b7cabe7f543 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PropagateScopeDependencies.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PropagateScopeDependencies.ts @@ -685,7 +685,9 @@ class PropagationVisitor extends ReactiveFunctionVisitor { const scopeDependencies = context.enter(scope.scope, () => { this.visitBlock(scope.instructions, context); }); - scope.scope.dependencies = scopeDependencies; + if (!scope.scope.source) { + scope.scope.dependencies = scopeDependencies; + } } override visitPrunedScope( diff --git a/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PruneInitializationDependencies.ts b/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PruneInitializationDependencies.ts index b9939addcf1..50ec6099adc 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PruneInitializationDependencies.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PruneInitializationDependencies.ts @@ -183,7 +183,7 @@ class Visitor extends ReactiveFunctionVisitor { ident.path.forEach((key) => { target &&= this.paths.get(target)?.get(key); }); - if (target && this.map.get(target) === "Create") { + if (target && this.map.get(target) === "Create" && !scope.scope.source) { scope.scope.dependencies.delete(ident); } }); diff --git a/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PruneNonReactiveDependencies.ts b/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PruneNonReactiveDependencies.ts index 2e1748fc6a4..b6cb3453a85 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PruneNonReactiveDependencies.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PruneNonReactiveDependencies.ts @@ -97,7 +97,7 @@ class Visitor extends ReactiveFunctionVisitor { this.traverseScope(scopeBlock, state); for (const dep of scopeBlock.scope.dependencies) { const isReactive = state.has(dep.identifier.id); - if (!isReactive) { + if (!isReactive && !scopeBlock.scope.source) { scopeBlock.scope.dependencies.delete(dep); } } 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 365e3aa88d4..7e4f21004fb 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/Validation/ValidatePreservedManualMemoization.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/Validation/ValidatePreservedManualMemoization.ts @@ -146,9 +146,13 @@ function compareDeps( const rootsEqual = (inferred.root.kind === "Global" && source.root.kind === "Global" && - inferred.root.identifierName === source.root.identifierName) || - (inferred.root.kind === "NamedLocal" && - source.root.kind === "NamedLocal" && + inferred.root.binding.binding.name === + source.root.binding.binding.name) || + ((inferred.root.kind === "NamedLocal" || + inferred.root.kind === "InlinedGlobal") && + (source.root.kind === "NamedLocal" || + source.root.kind === "InlinedGlobal") && + source.root.kind === inferred.root.kind && inferred.root.value.identifier.id === source.root.value.identifier.id); if (!rootsEqual) { return CompareDependencyResult.RootDifference; @@ -378,7 +382,8 @@ class Visitor extends ReactiveFunctionVisitor { if ( state.manualMemoState != null && - state.manualMemoState.depsFromSource != null + state.manualMemoState.depsFromSource != null && + !scopeBlock.scope.source ) { for (const dep of scopeBlock.scope.dependencies) { validateInferredDep( diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/nomemo-path.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/nomemo-path.expect.md new file mode 100644 index 00000000000..9dd98c8cba7 --- /dev/null +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/nomemo-path.expect.md @@ -0,0 +1,66 @@ + +## Input + +```javascript +// @disableMemoizationForDebugging +import { useMemo } from "react"; + +const w = 42; + +function Component(a) { + let x = useMemo(() => a.x, [a, w]); + return
{x}
; +} + +export const FIXTURE_ENTRYPOINT = { + fn: Component, + params: [{ x: 42 }], + isComponent: true, +}; + +``` + +## Code + +```javascript +import { c as _c } from "react/compiler-runtime"; // @disableMemoizationForDebugging +import { useMemo } from "react"; + +const w = 42; + +function Component(a) { + const $ = _c(5); + const t0 = w; + let t1; + let t2; + if ($[0] !== a || $[1] !== t0) { + t2 = a.x; + $[0] = a; + $[1] = t0; + $[2] = t2; + } else { + t2 = $[2]; + } + t1 = t2; + const x = t1; + let t3; + if ($[3] !== x || true) { + t3 =
{x}
; + $[3] = x; + $[4] = t3; + } else { + t3 = $[4]; + } + return t3; +} + +export const FIXTURE_ENTRYPOINT = { + fn: Component, + params: [{ x: 42 }], + isComponent: true, +}; + +``` + +### Eval output +(kind: ok)
42
\ No newline at end of file diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/nomemo-path.js b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/nomemo-path.js new file mode 100644 index 00000000000..34fcb759e1a --- /dev/null +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/nomemo-path.js @@ -0,0 +1,15 @@ +// @disableMemoizationForDebugging +import { useMemo } from "react"; + +const w = 42; + +function Component(a) { + let x = useMemo(() => a.x, [a, w]); + return
{x}
; +} + +export const FIXTURE_ENTRYPOINT = { + fn: Component, + params: [{ x: 42 }], + isComponent: true, +}; diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useMemo-simple-preserved-nomemo.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useMemo-simple-preserved-nomemo.expect.md index 307567a39cf..d1c7afdfd05 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useMemo-simple-preserved-nomemo.expect.md +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useMemo-simple-preserved-nomemo.expect.md @@ -25,26 +25,25 @@ import { c as _c } from "react/compiler-runtime"; // @disableMemoizationForDebug import { useMemo } from "react"; function Component(t0) { - const $ = _c(4); + const $ = _c(3); const { a } = t0; let t1; let t2; - if ($[0] !== a) { + if ($[0] === Symbol.for("react.memo_cache_sentinel")) { t2 = [a]; - $[0] = a; - $[1] = t2; + $[0] = t2; } else { - t2 = $[1]; + t2 = $[0]; } t1 = t2; const x = t1; let t3; - if ($[2] !== x || true) { + if ($[1] !== x || true) { t3 =
{x}
; - $[2] = x; - $[3] = t3; + $[1] = x; + $[2] = t3; } else { - t3 = $[3]; + t3 = $[2]; } return t3; } 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 e7e88d25986..0a78356a5e5 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,26 +25,25 @@ import { c as _c } from "react/compiler-runtime"; // @enablePreserveExistingManu import { useMemo } from "react"; function Component(t0) { - const $ = _c(4); + const $ = _c(3); const { a } = t0; let t1; let t2; - if ($[0] !== a) { + if ($[0] === Symbol.for("react.memo_cache_sentinel")) { t2 = [a]; - $[0] = a; - $[1] = t2; + $[0] = t2; } else { - t2 = $[1]; + t2 = $[0]; } t1 = t2; const x = t1; let t3; - if ($[2] !== x) { + if ($[1] !== x) { t3 =
{x}
; - $[2] = x; - $[3] = t3; + $[1] = x; + $[2] = t3; } else { - t3 = $[3]; + t3 = $[2]; } return t3; }