diff --git a/compiler/packages/babel-plugin-react-forget/src/HIR/PruneUnusedLabelsHIR.ts b/compiler/packages/babel-plugin-react-forget/src/HIR/PruneUnusedLabelsHIR.ts index 95392149d6..10714d5d87 100644 --- a/compiler/packages/babel-plugin-react-forget/src/HIR/PruneUnusedLabelsHIR.ts +++ b/compiler/packages/babel-plugin-react-forget/src/HIR/PruneUnusedLabelsHIR.ts @@ -7,7 +7,7 @@ export function pruneUnusedLabelsHIR(fn: HIRFunction): void { next: BlockId; fallthrough: BlockId; }> = []; - + const rewrites: Map = new Map(); for (const [blockId, block] of fn.body.blocks) { const terminal = block.terminal; if (terminal.kind === "label") { @@ -32,10 +32,11 @@ export function pruneUnusedLabelsHIR(fn: HIRFunction): void { } for (const { - label: labelId, + label: originalLabelId, next: nextId, fallthrough: fallthroughId, } of merged) { + const labelId = rewrites.get(originalLabelId) ?? originalLabelId; const label = fn.body.blocks.get(labelId)!; const next = fn.body.blocks.get(nextId)!; const fallthrough = fn.body.blocks.get(fallthroughId)!; @@ -52,7 +53,7 @@ export function pruneUnusedLabelsHIR(fn: HIRFunction): void { CompilerError.invariant( next.preds.size === 1 && fallthrough.preds.size === 1 && - next.preds.has(labelId) && + next.preds.has(originalLabelId) && fallthrough.preds.has(nextId), { reason: "Unexpected block predecessors when merging label blocks", @@ -64,5 +65,6 @@ export function pruneUnusedLabelsHIR(fn: HIRFunction): void { label.terminal = fallthrough.terminal; fn.body.blocks.delete(nextId); fn.body.blocks.delete(fallthroughId); + rewrites.set(fallthroughId, labelId); } } diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/consecutive-use-memo.expect.md b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/consecutive-use-memo.expect.md new file mode 100644 index 0000000000..7ed11d002e --- /dev/null +++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/consecutive-use-memo.expect.md @@ -0,0 +1,72 @@ + +## Input + +```javascript +import { useMemo } from "react"; +import { identity } from "shared-runtime"; + +function useHook({ a, b }) { + const valA = useMemo(() => identity({ a }), [a]); + const valB = useMemo(() => identity([b]), [b]); + return [valA, valB]; +} + +export const FIXTURE_ENTRYPOINT = { + fn: useHook, + params: [{ a: 2, b: 3 }], +}; + +``` + +## Code + +```javascript +import { useMemo, unstable_useMemoCache as useMemoCache } from "react"; +import { identity } from "shared-runtime"; + +function useHook(t0) { + const $ = useMemoCache(7); + const { a, b } = t0; + let t1; + let t2; + if ($[0] !== a) { + t2 = identity({ a }); + $[0] = a; + $[1] = t2; + } else { + t2 = $[1]; + } + t1 = t2; + const valA = t1; + let t3; + let t4; + if ($[2] !== b) { + t4 = identity([b]); + $[2] = b; + $[3] = t4; + } else { + t4 = $[3]; + } + t3 = t4; + const valB = t3; + let t5; + if ($[4] !== valA || $[5] !== valB) { + t5 = [valA, valB]; + $[4] = valA; + $[5] = valB; + $[6] = t5; + } else { + t5 = $[6]; + } + return t5; +} + +export const FIXTURE_ENTRYPOINT = { + fn: useHook, + params: [{ a: 2, b: 3 }], +}; + +``` + +### Eval output +(kind: ok) [{"a":2},[3]] \ No newline at end of file diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/consecutive-use-memo.ts b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/consecutive-use-memo.ts new file mode 100644 index 0000000000..ce2c8037e4 --- /dev/null +++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/consecutive-use-memo.ts @@ -0,0 +1,13 @@ +import { useMemo } from "react"; +import { identity } from "shared-runtime"; + +function useHook({ a, b }) { + const valA = useMemo(() => identity({ a }), [a]); + const valB = useMemo(() => identity([b]), [b]); + return [valA, valB]; +} + +export const FIXTURE_ENTRYPOINT = { + fn: useHook, + params: [{ a: 2, b: 3 }], +};