diff --git a/compiler/packages/babel-plugin-react-forget/src/HIR/Environment.ts b/compiler/packages/babel-plugin-react-forget/src/HIR/Environment.ts index af3d326124..fc9ac6d7b1 100644 --- a/compiler/packages/babel-plugin-react-forget/src/HIR/Environment.ts +++ b/compiler/packages/babel-plugin-react-forget/src/HIR/Environment.ts @@ -102,6 +102,26 @@ export type Hook = z.infer; const EnvironmentConfigSchema = z.object({ customHooks: z.map(z.string(), HookSchema).optional().default(new Map()), + /** + * Enable using information from existing useMemo/useCallback to understand when a value is done + * being mutated. With this mode enabled, Forget will still discard the actual useMemo/useCallback + * calls and may memoize slightly differently. However, it will assume that the values produced + * are not subsequently modified, guaranteeing that the value will be memoized. + * + * By preserving guarantees about when values are memoized, this option preserves any existing + * behavior that depends on referential equality in the original program. Notably, this preserves + * existing effect behavior (how often effects fire) for effects that rely on referential equality. + * + * When disabled, Forget will not only prune useMemo and useCallback calls but also completely ignore + * them, not using any information from them to guide compilation. Therefore, disabling this flag + * will produce output that mimics the result from removing all memoization. + * + * Our recommendation is to first try running your application with this flag enabled, then attempt + * to disable this flag and see what changes or breaks. This will mostly likely be effects that + * depend on referential equality, which can be refactored (TODO guide for this). + */ + enablePreserveExistingMemoizationGuarantees: z.boolean().default(false), + // 🌲 enableForest: z.boolean().default(false), // <🌲> diff --git a/compiler/packages/babel-plugin-react-forget/src/HIR/HIR.ts b/compiler/packages/babel-plugin-react-forget/src/HIR/HIR.ts index c4e69ebaea..3f4ee9857b 100644 --- a/compiler/packages/babel-plugin-react-forget/src/HIR/HIR.ts +++ b/compiler/packages/babel-plugin-react-forget/src/HIR/HIR.ts @@ -845,6 +845,12 @@ export type InstructionValue = } // `debugger` statement | { kind: "Debugger"; loc: SourceLocation } + /* + * Represents semantic information from useMemo/useCallback that the developer + * has indicated a particular value should be memoized. This value is ignored + * unless the TODO flag is enabled. + */ + | { kind: "Memoize"; value: Place; loc: SourceLocation } /* * Catch-all for statements such as type imports, nested class declarations, etc * which are not directly represented, but included for completeness and to allow diff --git a/compiler/packages/babel-plugin-react-forget/src/HIR/HIRBuilder.ts b/compiler/packages/babel-plugin-react-forget/src/HIR/HIRBuilder.ts index e79891ff98..390bf7989f 100644 --- a/compiler/packages/babel-plugin-react-forget/src/HIR/HIRBuilder.ts +++ b/compiler/packages/babel-plugin-react-forget/src/HIR/HIRBuilder.ts @@ -15,11 +15,14 @@ import { BasicBlock, BlockId, BlockKind, + Effect, + GeneratedSource, GotoVariant, HIR, Identifier, IdentifierId, Instruction, + Place, Terminal, makeBlockId, makeInstructionId, @@ -856,3 +859,19 @@ export function removeUnnecessaryTryCatch(fn: HIR): void { } } } + +export function createTemporaryPlace(env: Environment): Place { + return { + kind: "Identifier", + identifier: { + id: env.nextIdentifierId, + mutableRange: { start: makeInstructionId(0), end: makeInstructionId(0) }, + name: null, + scope: null, + type: makeType(), + }, + reactive: false, + effect: Effect.Unknown, + loc: GeneratedSource, + }; +} 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 20c4ccfd93..fbe57a451d 100644 --- a/compiler/packages/babel-plugin-react-forget/src/HIR/PrintHIR.ts +++ b/compiler/packages/babel-plugin-react-forget/src/HIR/PrintHIR.ts @@ -594,6 +594,10 @@ export function printInstructionValue(instrValue: ReactiveValue): string { } ${printPlace(instrValue.value)}`; break; } + case "Memoize": { + value = `Memoize ${printPlace(instrValue.value)}`; + break; + } default: { assertExhaustive( instrValue, diff --git a/compiler/packages/babel-plugin-react-forget/src/HIR/visitors.ts b/compiler/packages/babel-plugin-react-forget/src/HIR/visitors.ts index bf76b62b64..54007e5e09 100644 --- a/compiler/packages/babel-plugin-react-forget/src/HIR/visitors.ts +++ b/compiler/packages/babel-plugin-react-forget/src/HIR/visitors.ts @@ -209,6 +209,10 @@ export function* eachInstructionValueOperand( yield instrValue.value; break; } + case "Memoize": { + yield instrValue.value; + break; + } case "Debugger": case "RegExpLiteral": case "LoadGlobal": @@ -510,6 +514,10 @@ export function mapInstructionValueOperands( instrValue.value = fn(instrValue.value); break; } + case "Memoize": { + instrValue.value = fn(instrValue.value); + break; + } case "Debugger": case "RegExpLiteral": case "LoadGlobal": 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 03be0f92a9..39f8f1d6a7 100644 --- a/compiler/packages/babel-plugin-react-forget/src/Inference/DropManualMemoization.ts +++ b/compiler/packages/babel-plugin-react-forget/src/Inference/DropManualMemoization.ts @@ -10,9 +10,13 @@ import { Effect, HIRFunction, IdentifierId, + Instruction, Place, SpreadPattern, + makeInstructionId, + markInstructionIds, } from "../HIR"; +import { createTemporaryPlace } from "../HIR/HIRBuilder"; import { HookKind } from "../HIR/ObjectShape"; /* @@ -26,8 +30,14 @@ import { HookKind } from "../HIR/ObjectShape"; export function dropManualMemoization(func: HIRFunction): void { const hooks = new Map(); const react = new Set(); + let hasChanges = false; for (const [_, block] of func.body.blocks) { - for (const instr of block.instructions) { + 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 "LoadGlobal": { if ( @@ -71,16 +81,22 @@ export function dropManualMemoization(func: HIRFunction): void { }); } /* - * TODO(gsn): Consider inlining the function passed to useMemo, - * rather than just calling it directly. - * * Replace the hook callee with the fn arg. * * before: - * foo = Call useMemo$2($9, $10) + * $1 = LoadGlobal useMemo // load the useMemo global + * $2 = FunctionExpression ... // memo function + * $3 = ArrayExpression [ ... ] // deps array + * $4 = Call $1 ($2, $3 ) // invoke useMemo w fn and deps * * after: - * foo = Call $9() + * $1 = LoadGlobal useMemo // load the useMemo global (dead code) + * $2 = FunctionExpression ... // memo function + * $3 = ArrayExpression [ ... ] // deps array (dead code) + * $4 = Call $2 () // invoke the memo function itself + * + * Note that a later pass (InlineImmediatelyInvokedFunctionExpressions) will + * inline the useMemo callback along with any other immediately invoked IIFEs. */ if (fn.kind === "Identifier") { instr.value = { @@ -93,6 +109,46 @@ export function dropManualMemoization(func: HIRFunction): void { args: [], loc: instr.value.loc, }; + if ( + func.env.config.enablePreserveExistingMemoizationGuarantees + ) { + /** + * When this flag is enabled we also compile in a 'Memoize' instruction + * to preserve the intended memoization boundary: + * + * Normal output: + * $1 = LoadGlobal useMemo // load the useMemo global (dead code) + * $2 = FunctionExpression ... // memo function + * $3 = ArrayExpression [ ... ] // deps array (dead code) + * $4 = Call $2 () // invoke the memo function itself + * + * Output w flag enabled: + * $1 = LoadGlobal useMemo // load the useMemo global (dead code) + * $2 = FunctionExpression ... // memo function + * $3 = ArrayExpression [ ... ] // deps array (dead code) + * $5 = Call $2 () // invoke the memo function itself + * $4 = Memoize $5 // preserve memo information + * + * Note that we synthesize a new temporary for the call ($5) and use + * the original lvalue for the result of the Memoize instruction, so that + * we don't have to rewrite subsequent instructions. + */ + const lvalue = instr.lvalue; + const temp = createTemporaryPlace(func.env); + instr.lvalue = { ...temp }; + nextInstructions = + nextInstructions ?? block.instructions.slice(0, i + 1); + nextInstructions.push({ + id: makeInstructionId(0), + lvalue, + value: { + kind: "Memoize", + value: temp, + loc: instr.loc, + }, + loc: instr.loc, + }); + } } } else if (hookKind === "useCallback") { const [fn] = instr.value.args as Array< @@ -110,23 +166,63 @@ export function dropManualMemoization(func: HIRFunction): void { * Instead of a Call, just alias the callback directly. * * before: - * foo = Call useCallback$8($19) + * $1 = LoadGlobal useCallback + * $2 = FunctionExpression ... // the callback being memoized + * $3 = ArrayExpression ... // deps array + * $3 = Call $1 ( $2, $3 ) // invoke useCallback * * after: - * foo = $19 + * $1 = LoadGlobal useCallback // dead code + * $2 = FunctionExpression ... // the callback being memoized + * $3 = ArrayExpression ... // deps array (dead code) + * $3 = LoadLocal $2 // reference the function */ if (fn.kind === "Identifier") { - instr.value = { - kind: "LoadLocal", - place: { - kind: "Identifier", - identifier: fn.identifier, - effect: Effect.Unknown, - reactive: false, + if ( + func.env.config.enablePreserveExistingMemoizationGuarantees + ) { + /** + * With the flag enabled the output changes to use a Memoize instruction instead + * a loadlocal to load the function expression into the original temporary: + * + * Normal output: + * $1 = LoadGlobal useCallback // dead code + * $2 = FunctionExpression ... // the callback being memoized + * $3 = ArrayExpression ... // deps array (dead code) + * $3 = LoadLocal $2 // reference the function + * + * With flag enabled: + * $1 = LoadGlobal useCallback // dead code + * $2 = FunctionExpression ... // the callback being memoized + * $3 = ArrayExpression ... // deps array (dead code) + * $3 = Memoize $2 // reference the function + * + * Note the s/LoadLocal/Memoize/ + */ + instr.value = { + kind: "Memoize", + value: { + kind: "Identifier", + identifier: fn.identifier, + effect: Effect.Unknown, + reactive: false, + loc: instr.value.loc, + }, loc: instr.value.loc, - }, - loc: instr.value.loc, - }; + }; + } else { + instr.value = { + kind: "LoadLocal", + place: { + kind: "Identifier", + identifier: fn.identifier, + effect: Effect.Unknown, + reactive: false, + loc: instr.value.loc, + }, + loc: instr.value.loc, + }; + } } } } @@ -134,5 +230,12 @@ export function dropManualMemoization(func: HIRFunction): void { } } } + if (nextInstructions !== null) { + block.instructions = nextInstructions; + hasChanges = true; + } + } + if (hasChanges) { + markInstructionIds(func.body); } } diff --git a/compiler/packages/babel-plugin-react-forget/src/Inference/InferReferenceEffects.ts b/compiler/packages/babel-plugin-react-forget/src/Inference/InferReferenceEffects.ts index e0111eb16a..a6afdf6d51 100644 --- a/compiler/packages/babel-plugin-react-forget/src/Inference/InferReferenceEffects.ts +++ b/compiler/packages/babel-plugin-react-forget/src/Inference/InferReferenceEffects.ts @@ -1162,6 +1162,17 @@ function inferBlock( state.alias(lvalue, instrValue.value); continue; } + case "Memoize": { + state.initialize(instrValue, { + kind: ValueKind.Frozen, + reason: new Set([ValueReason.Other]), + }); + state.reference(instrValue.value, Effect.Freeze, ValueReason.Other); + const lvalue = instr.lvalue; + lvalue.effect = Effect.ConditionallyMutate; + state.alias(lvalue, instrValue.value); + continue; + } case "LoadLocal": { const lvalue = instr.lvalue; const effect = 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 6ad2e086a8..a3bd094348 100644 --- a/compiler/packages/babel-plugin-react-forget/src/Optimization/DeadCodeElimination.ts +++ b/compiler/packages/babel-plugin-react-forget/src/Optimization/DeadCodeElimination.ts @@ -313,6 +313,7 @@ function pruneableValue(value: InstructionValue, state: State): boolean { case "StoreContext": { return false; } + case "Memoize": case "RegExpLiteral": case "LoadGlobal": case "ArrayExpression": diff --git a/compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/CodegenReactiveFunction.ts b/compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/CodegenReactiveFunction.ts index 701554d469..15380ed6bc 100644 --- a/compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/CodegenReactiveFunction.ts +++ b/compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/CodegenReactiveFunction.ts @@ -1628,6 +1628,10 @@ function codegenInstructionValue( ); break; } + case "Memoize": { + value = codegenPlaceToExpression(cx, instrValue.value); + break; + } case "Debugger": case "DeclareLocal": case "DeclareContext": diff --git a/compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/InferReactiveScopeVariables.ts b/compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/InferReactiveScopeVariables.ts index e6692f089d..0e38cbf251 100644 --- a/compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/InferReactiveScopeVariables.ts +++ b/compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/InferReactiveScopeVariables.ts @@ -245,7 +245,8 @@ function mayAllocate(env: Environment, instruction: Instruction): boolean { case "Primitive": case "NextIterableOf": case "NextPropertyOf": - case "Debugger": { + case "Debugger": + case "Memoize": { return false; } case "UnaryExpression": diff --git a/compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/PruneNonEscapingScopes.ts b/compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/PruneNonEscapingScopes.ts index 9f1af74010..e97ae5e9e0 100644 --- a/compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/PruneNonEscapingScopes.ts +++ b/compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/PruneNonEscapingScopes.ts @@ -471,6 +471,7 @@ function computeMemoizationInputs( rvalues: [], }; } + case "Memoize": case "Await": case "TypeCastExpression": case "NextIterableOf": { diff --git a/compiler/packages/babel-plugin-react-forget/src/TypeInference/InferTypes.ts b/compiler/packages/babel-plugin-react-forget/src/TypeInference/InferTypes.ts index 9b677e4e95..c39abb1cde 100644 --- a/compiler/packages/babel-plugin-react-forget/src/TypeInference/InferTypes.ts +++ b/compiler/packages/babel-plugin-react-forget/src/TypeInference/InferTypes.ts @@ -280,6 +280,11 @@ function* generateInstructionTypes( break; } + case "Memoize": { + yield equation(left, value.value.identifier.type); + break; + } + case "PropertyDelete": case "ComputedDelete": { yield equation(left, { kind: "Primitive" }); diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/useMemo-maybe-modified-later-dont-preserve-memoization-guarantees.expect.md b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/useMemo-maybe-modified-later-dont-preserve-memoization-guarantees.expect.md new file mode 100644 index 0000000000..7da8ef440c --- /dev/null +++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/useMemo-maybe-modified-later-dont-preserve-memoization-guarantees.expect.md @@ -0,0 +1,54 @@ + +## Input + +```javascript +// @enablePreserveExistingMemoizationGuarantees:false +import { useMemo } from "react"; +import { identity, makeObject_Primitives, mutate } from "shared-runtime"; + +function Component(props) { + const object = useMemo(() => makeObject_Primitives(), []); + identity(object); + return object; +} + +export const FIXTURE_ENTRYPOINT = { + fn: Component, + params: [{}], +}; + +``` + +## Code + +```javascript +// @enablePreserveExistingMemoizationGuarantees:false +import { useMemo, unstable_useMemoCache as useMemoCache } from "react"; +import { identity, makeObject_Primitives, mutate } from "shared-runtime"; + +function Component(props) { + const $ = useMemoCache(2); + let t7; + let object; + if ($[0] === Symbol.for("react.memo_cache_sentinel")) { + t7 = makeObject_Primitives(); + object = t7; + identity(object); + $[0] = object; + $[1] = t7; + } else { + object = $[0]; + t7 = $[1]; + } + return object; +} + +export const FIXTURE_ENTRYPOINT = { + fn: Component, + params: [{}], +}; + +``` + +### Eval output +(kind: ok) {"a":0,"b":"value1","c":true} \ No newline at end of file diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/useMemo-maybe-modified-later-dont-preserve-memoization-guarantees.js b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/useMemo-maybe-modified-later-dont-preserve-memoization-guarantees.js new file mode 100644 index 0000000000..44c08441e3 --- /dev/null +++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/useMemo-maybe-modified-later-dont-preserve-memoization-guarantees.js @@ -0,0 +1,14 @@ +// @enablePreserveExistingMemoizationGuarantees:false +import { useMemo } from "react"; +import { identity, makeObject_Primitives, mutate } from "shared-runtime"; + +function Component(props) { + const object = useMemo(() => makeObject_Primitives(), []); + identity(object); + return object; +} + +export const FIXTURE_ENTRYPOINT = { + fn: Component, + params: [{}], +}; diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/useMemo-maybe-modified-later-preserve-memoization-guarantees.expect.md b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/useMemo-maybe-modified-later-preserve-memoization-guarantees.expect.md new file mode 100644 index 0000000000..d652d52c94 --- /dev/null +++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/useMemo-maybe-modified-later-preserve-memoization-guarantees.expect.md @@ -0,0 +1,53 @@ + +## Input + +```javascript +// @enablePreserveExistingMemoizationGuarantees +import { useMemo } from "react"; +import { identity, makeObject_Primitives, mutate } from "shared-runtime"; + +function Component(props) { + const object = useMemo(() => makeObject_Primitives(), []); + identity(object); + return object; +} + +export const FIXTURE_ENTRYPOINT = { + fn: Component, + params: [{}], +}; + +``` + +## Code + +```javascript +// @enablePreserveExistingMemoizationGuarantees +import { useMemo, unstable_useMemoCache as useMemoCache } from "react"; +import { identity, makeObject_Primitives, mutate } from "shared-runtime"; + +function Component(props) { + const $ = useMemoCache(1); + let t15; + let t0; + if ($[0] === Symbol.for("react.memo_cache_sentinel")) { + t0 = makeObject_Primitives(); + $[0] = t0; + } else { + t0 = $[0]; + } + t15 = t0; + const object = t15; + identity(object); + return object; +} + +export const FIXTURE_ENTRYPOINT = { + fn: Component, + params: [{}], +}; + +``` + +### Eval output +(kind: ok) {"a":0,"b":"value1","c":true} \ No newline at end of file diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/useMemo-maybe-modified-later-preserve-memoization-guarantees.js b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/useMemo-maybe-modified-later-preserve-memoization-guarantees.js new file mode 100644 index 0000000000..961fc974a2 --- /dev/null +++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/useMemo-maybe-modified-later-preserve-memoization-guarantees.js @@ -0,0 +1,14 @@ +// @enablePreserveExistingMemoizationGuarantees +import { useMemo } from "react"; +import { identity, makeObject_Primitives, mutate } from "shared-runtime"; + +function Component(props) { + const object = useMemo(() => makeObject_Primitives(), []); + identity(object); + return object; +} + +export const FIXTURE_ENTRYPOINT = { + fn: Component, + params: [{}], +};