From 723b616c6778fa46fb817e2e4e000fdc7fd1b46d Mon Sep 17 00:00:00 2001 From: Joe Savona Date: Fri, 15 Dec 2023 13:47:22 -0800 Subject: [PATCH] enablePreserveMemo treats memo deps as frozen See discussion on #2448 for full context. In the new `@enablePreserveExistingMemoizationGuarantees` mode, the goal is to preserve the existing referential equality guarantees from the original code. #2448 lays the groundwork by explicitly marking the _output_ of each useMemo block as memoized, hinting to the compiler that the value cannot subsequently change. This ensures the mutable range doesn't extend _later_, possibly overlapping a hook call and causing memoization to gett pruned. This PR fixes the other direction. There are cases where free variables referenced in the useMemo block could have been inferred as mutated, which could then extend the _start_ of the range earlier past a hook: ```javascript const foo = createObject(); useBar(); const baz = useMemo(() => { const baz = createObject(); maybeMutate(foo, baz); return baz; }, [foo]); ``` Here the compiler would infer that both `baz` and `foo` are mutable at the `maybeMutate()` call, grouping them in the same scope. But that scope would span the `useBar()` call, and be pruned, meaning that `baz` went unmemoized. However, useMemo blocks shouldn't be mutating free variables. Only variables newly created within the useMemo block should be mutable. So this PR extends the feature to treat all free variables referenced in a useMemo block as frozen as of the block itself. --- .../src/HIR/PrintHIR.ts | 4 +- .../src/Inference/DropManualMemoization.ts | 46 +++++++-- .../src/Optimization/DeadCodeElimination.ts | 9 +- ...-preserve-memoization-guarantees.expect.md | 74 +++++++++++++++ ...le-dont-preserve-memoization-guarantees.js | 30 ++++++ ...-preserve-memoization-guarantees.expect.md | 95 +++++++++++++++++++ ...ariable-preserve-memoization-guarantees.js | 27 ++++++ .../packages/sprout/src/shared-runtime.ts | 4 + 8 files changed, 281 insertions(+), 8 deletions(-) create mode 100644 compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/useMemo-mabye-modified-free-variable-dont-preserve-memoization-guarantees.expect.md create mode 100644 compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/useMemo-mabye-modified-free-variable-dont-preserve-memoization-guarantees.js create mode 100644 compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/useMemo-mabye-modified-free-variable-preserve-memoization-guarantees.expect.md create mode 100644 compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/useMemo-mabye-modified-free-variable-preserve-memoization-guarantees.js diff --git a/compiler/packages/babel-plugin-react-forget/src/HIR/PrintHIR.ts b/compiler/packages/babel-plugin-react-forget/src/HIR/PrintHIR.ts index fbe57a451d..68f49329e8 100644 --- a/compiler/packages/babel-plugin-react-forget/src/HIR/PrintHIR.ts +++ b/compiler/packages/babel-plugin-react-forget/src/HIR/PrintHIR.ts @@ -494,6 +494,8 @@ export function printInstructionValue(instrValue: ReactiveValue): string { } case "ObjectMethod": case "FunctionExpression": { + const kind = + instrValue.kind === "FunctionExpression" ? "Function" : "ObjectMethod"; const name = getFunctionName(instrValue, ""); const fn = printFunction(instrValue.loweredFunc.func) .split("\n") @@ -505,7 +507,7 @@ export function printInstructionValue(instrValue: ReactiveValue): string { const context = instrValue.loweredFunc.func.context .map((dep) => printPlace(dep)) .join(","); - value = `Function ${name} @deps[${deps}] @context[${context}]:\n${fn}`; + value = `${kind} ${name} @deps[${deps}] @context[${context}]:\n${fn}`; break; } case "TaggedTemplateExpression": { diff --git a/compiler/packages/babel-plugin-react-forget/src/Inference/DropManualMemoization.ts b/compiler/packages/babel-plugin-react-forget/src/Inference/DropManualMemoization.ts index 39f8f1d6a7..8e3f29b46e 100644 --- a/compiler/packages/babel-plugin-react-forget/src/Inference/DropManualMemoization.ts +++ b/compiler/packages/babel-plugin-react-forget/src/Inference/DropManualMemoization.ts @@ -8,16 +8,17 @@ import { CompilerError } from ".."; import { Effect, + FunctionExpression, HIRFunction, IdentifierId, Instruction, Place, SpreadPattern, makeInstructionId, - markInstructionIds, } from "../HIR"; import { createTemporaryPlace } from "../HIR/HIRBuilder"; import { HookKind } from "../HIR/ObjectShape"; +import { eachInstructionValueOperand } from "../HIR/visitors"; /* * Removes manual memoization using the `useMemo` and `useCallback` APIs. This pass is designed @@ -28,6 +29,7 @@ import { HookKind } from "../HIR/ObjectShape"; * eg `React.useMemo()`. */ export function dropManualMemoization(func: HIRFunction): void { + const functions = new Map(); const hooks = new Map(); const react = new Set(); let hasChanges = false; @@ -35,10 +37,11 @@ export function dropManualMemoization(func: HIRFunction): void { let nextInstructions: Array | null = null; for (let i = 0; i < block.instructions.length; i++) { const instr = block.instructions[i]!; - if (nextInstructions !== null) { - nextInstructions.push(instr); - } switch (instr.value.kind) { + case "FunctionExpression": { + functions.set(instr.lvalue.identifier.id, instr.value); + break; + } case "LoadGlobal": { if ( instr.value.name === "useMemo" || @@ -109,6 +112,7 @@ export function dropManualMemoization(func: HIRFunction): void { args: [], loc: instr.value.loc, }; + if ( func.env.config.enablePreserveExistingMemoizationGuarantees ) { @@ -137,7 +141,29 @@ export function dropManualMemoization(func: HIRFunction): void { const temp = createTemporaryPlace(func.env); instr.lvalue = { ...temp }; nextInstructions = - nextInstructions ?? block.instructions.slice(0, i + 1); + nextInstructions ?? block.instructions.slice(0, i); + + const functionExpression = functions.get(fn.identifier.id); + if (functionExpression !== undefined) { + for (const operand of eachInstructionValueOperand( + functionExpression + )) { + const operandLValue = createTemporaryPlace(func.env); + nextInstructions.push({ + id: makeInstructionId(0), + lvalue: operandLValue, + value: { + kind: "Memoize", + value: { ...operand }, + loc: instr.loc, + }, + loc: instr.loc, + }); + } + } + + nextInstructions.push(instr); + nextInstructions.push({ id: makeInstructionId(0), lvalue, @@ -148,7 +174,12 @@ export function dropManualMemoization(func: HIRFunction): void { }, loc: instr.loc, }); + } else { + if (nextInstructions !== null) { + nextInstructions.push(instr); + } } + continue; } } else if (hookKind === "useCallback") { const [fn] = instr.value.args as Array< @@ -229,6 +260,9 @@ export function dropManualMemoization(func: HIRFunction): void { break; } } + if (nextInstructions !== null) { + nextInstructions.push(instr); + } } if (nextInstructions !== null) { block.instructions = nextInstructions; @@ -236,6 +270,6 @@ export function dropManualMemoization(func: HIRFunction): void { } } if (hasChanges) { - markInstructionIds(func.body); + // markInstructionIds(func.body); } } diff --git a/compiler/packages/babel-plugin-react-forget/src/Optimization/DeadCodeElimination.ts b/compiler/packages/babel-plugin-react-forget/src/Optimization/DeadCodeElimination.ts index a3bd094348..2a1c8d9b05 100644 --- a/compiler/packages/babel-plugin-react-forget/src/Optimization/DeadCodeElimination.ts +++ b/compiler/packages/babel-plugin-react-forget/src/Optimization/DeadCodeElimination.ts @@ -313,7 +313,14 @@ function pruneableValue(value: InstructionValue, state: State): boolean { case "StoreContext": { return false; } - case "Memoize": + case "Memoize": { + /** + * This instruction is used by the @enablePreserveExistingMemoizationGuarantees feature + * to preserve information about memoization semantics in the original code. We can't + * DCE without losing the memoization guarantees. + */ + return false; + } case "RegExpLiteral": case "LoadGlobal": case "ArrayExpression": diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/useMemo-mabye-modified-free-variable-dont-preserve-memoization-guarantees.expect.md b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/useMemo-mabye-modified-free-variable-dont-preserve-memoization-guarantees.expect.md new file mode 100644 index 0000000000..b3918c1910 --- /dev/null +++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/useMemo-mabye-modified-free-variable-dont-preserve-memoization-guarantees.expect.md @@ -0,0 +1,74 @@ + +## Input + +```javascript +// @enablePreserveExistingMemoizationGuarantees:false +import { useMemo } from "react"; +import { + identity, + makeObject_Primitives, + mutate, + useHook, +} from "shared-runtime"; + +function Component(props) { + // With the feature disabled these variables are inferred as being mutated inside the useMemo block + const free = makeObject_Primitives(); + const free2 = makeObject_Primitives(); + const part = free2.part; + + // This causes their range to extend to include this hook call, and in turn for the memoization to be pruned + useHook(); + const object = useMemo(() => { + const x = makeObject_Primitives(); + x.value = props.value; + mutate(x, free, part); + return x; + }, [props.value]); + return object; +} + +export const FIXTURE_ENTRYPOINT = { + fn: Component, + params: [{ value: 42 }], +}; + +``` + +## Code + +```javascript +// @enablePreserveExistingMemoizationGuarantees:false +import { useMemo } from "react"; +import { + identity, + makeObject_Primitives, + mutate, + useHook, +} from "shared-runtime"; + +function Component(props) { + const free = makeObject_Primitives(); + const free2 = makeObject_Primitives(); + const part = free2.part; + + useHook(); + let t39; + + const x = makeObject_Primitives(); + x.value = props.value; + mutate(x, free, part); + t39 = x; + const object = t39; + return object; +} + +export const FIXTURE_ENTRYPOINT = { + fn: Component, + params: [{ value: 42 }], +}; + +``` + +### Eval output +(kind: ok) {"a":0,"b":"value1","c":true,"value":42,"wat0":"joe"} \ No newline at end of file diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/useMemo-mabye-modified-free-variable-dont-preserve-memoization-guarantees.js b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/useMemo-mabye-modified-free-variable-dont-preserve-memoization-guarantees.js new file mode 100644 index 0000000000..27ef445b83 --- /dev/null +++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/useMemo-mabye-modified-free-variable-dont-preserve-memoization-guarantees.js @@ -0,0 +1,30 @@ +// @enablePreserveExistingMemoizationGuarantees:false +import { useMemo } from "react"; +import { + identity, + makeObject_Primitives, + mutate, + useHook, +} from "shared-runtime"; + +function Component(props) { + // With the feature disabled these variables are inferred as being mutated inside the useMemo block + const free = makeObject_Primitives(); + const free2 = makeObject_Primitives(); + const part = free2.part; + + // This causes their range to extend to include this hook call, and in turn for the memoization to be pruned + useHook(); + const object = useMemo(() => { + const x = makeObject_Primitives(); + x.value = props.value; + mutate(x, free, part); + return x; + }, [props.value]); + return object; +} + +export const FIXTURE_ENTRYPOINT = { + fn: Component, + params: [{ value: 42 }], +}; diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/useMemo-mabye-modified-free-variable-preserve-memoization-guarantees.expect.md b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/useMemo-mabye-modified-free-variable-preserve-memoization-guarantees.expect.md new file mode 100644 index 0000000000..74056b4b88 --- /dev/null +++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/useMemo-mabye-modified-free-variable-preserve-memoization-guarantees.expect.md @@ -0,0 +1,95 @@ + +## Input + +```javascript +// @enablePreserveExistingMemoizationGuarantees +import { useMemo } from "react"; +import { + identity, + makeObject_Primitives, + mutate, + useHook, +} from "shared-runtime"; + +function Component(props) { + const free = makeObject_Primitives(); + const free2 = makeObject_Primitives(); + const part = free2.part; + useHook(); + const object = useMemo(() => { + const x = makeObject_Primitives(); + x.value = props.value; + mutate(x, free, part); + return x; + }, [props.value]); + return object; +} + +export const FIXTURE_ENTRYPOINT = { + fn: Component, + params: [{ value: 42 }], +}; + +``` + +## Code + +```javascript +// @enablePreserveExistingMemoizationGuarantees +import { useMemo, unstable_useMemoCache as useMemoCache } from "react"; +import { + identity, + makeObject_Primitives, + mutate, + useHook, +} from "shared-runtime"; + +function Component(props) { + const $ = useMemoCache(4); + let t0; + if ($[0] === Symbol.for("react.memo_cache_sentinel")) { + t0 = makeObject_Primitives(); + $[0] = t0; + } else { + t0 = $[0]; + } + const free = t0; + let t1; + if ($[1] === Symbol.for("react.memo_cache_sentinel")) { + t1 = makeObject_Primitives(); + $[1] = t1; + } else { + t1 = $[1]; + } + const free2 = t1; + const part = free2.part; + useHook(); + + props.value; + free; + part; + let t44; + let x; + if ($[2] !== props.value) { + x = makeObject_Primitives(); + x.value = props.value; + mutate(x, free, part); + $[2] = props.value; + $[3] = x; + } else { + x = $[3]; + } + t44 = x; + const object = t44; + return object; +} + +export const FIXTURE_ENTRYPOINT = { + fn: Component, + params: [{ value: 42 }], +}; + +``` + +### Eval output +(kind: ok) {"a":0,"b":"value1","c":true,"value":42,"wat0":"joe"} \ No newline at end of file diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/useMemo-mabye-modified-free-variable-preserve-memoization-guarantees.js b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/useMemo-mabye-modified-free-variable-preserve-memoization-guarantees.js new file mode 100644 index 0000000000..d0431124b2 --- /dev/null +++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/useMemo-mabye-modified-free-variable-preserve-memoization-guarantees.js @@ -0,0 +1,27 @@ +// @enablePreserveExistingMemoizationGuarantees +import { useMemo } from "react"; +import { + identity, + makeObject_Primitives, + mutate, + useHook, +} from "shared-runtime"; + +function Component(props) { + const free = makeObject_Primitives(); + const free2 = makeObject_Primitives(); + const part = free2.part; + useHook(); + const object = useMemo(() => { + const x = makeObject_Primitives(); + x.value = props.value; + mutate(x, free, part); + return x; + }, [props.value]); + return object; +} + +export const FIXTURE_ENTRYPOINT = { + fn: Component, + params: [{ value: 42 }], +}; diff --git a/compiler/packages/sprout/src/shared-runtime.ts b/compiler/packages/sprout/src/shared-runtime.ts index 8c167e0d61..8f1133713f 100644 --- a/compiler/packages/sprout/src/shared-runtime.ts +++ b/compiler/packages/sprout/src/shared-runtime.ts @@ -153,6 +153,10 @@ export function throwInput(x: Object): never { throw x; } +export function useHook(): Object { + return makeObject_Primitives(); +} + const noAliasObject = Object.freeze({}); export function useNoAlias(...args: Array): object { return noAliasObject;