From 329809de818366cfccf0a6a69e00191b2639efeb Mon Sep 17 00:00:00 2001 From: Joe Savona Date: Wed, 17 May 2023 15:55:51 -0700 Subject: [PATCH] Dramatically simplify InlineUseMemo I realized a wayyyy simpler approach to inlining a lambda: wrap it in a labeled block. The transformation is roughly as follows: ```javascript // Before const x = useMemo(() => { if (a) { return b; } return c; }, [a, b, c]); return x; // After let x; label: { if (a) { x = b; break label; } x = c; break label; } return x; ``` The key to making this work is fixing up some edge cases in labeled blocks, hence the previous PRs. --- .../forget/src/Inference/InlineUseMemo.ts | 141 +----------------- .../useMemo-if-else-multiple-return.expect.md | 24 +-- .../useMemo-inlining-block-return.expect.md | 24 +-- .../useMemo-multiple-if-else.expect.md | 39 ++--- .../useMemo-switch-no-fallthrough.expect.md | 16 +- 5 files changed, 65 insertions(+), 179 deletions(-) diff --git a/compiler/forget/src/Inference/InlineUseMemo.ts b/compiler/forget/src/Inference/InlineUseMemo.ts index 9d346a4ba2..5a55ac9f70 100644 --- a/compiler/forget/src/Inference/InlineUseMemo.ts +++ b/compiler/forget/src/Inference/InlineUseMemo.ts @@ -13,20 +13,19 @@ import { Environment, FunctionExpression, GeneratedSource, - GotoTerminal, GotoVariant, - HIR, HIRFunction, Identifier, IdentifierId, InstructionKind, + LabelTerminal, Place, makeInstructionId, makeType, reversePostorderBlocks, } from "../HIR"; import { markInstructionIds, markPredecessors } from "../HIR/HIRBuilder"; -import { assertExhaustive, retainWhere } from "../Utils/utils"; +import { retainWhere } from "../Utils/utils"; /** * Rewrites `useMemo()` calls, rewriting so that the lambda body becomes part of the @@ -133,70 +132,18 @@ export function inlineUseMemo(fn: HIRFunction): void { // the useMemo block.instructions.length = ii; - // The block leading up to the useMemo needs to jump to the entry block of - // the useMemo control flow graph. These will be merged into a single block - // via MergeConsectuveBlocks - const newTerminal: GotoTerminal = { + // To account for complex control flow within the lambda, we treat the lambda + // as if it were a single labeled statement, and replace all returns with gotos + // to the label fallthrough. + const newTerminal: LabelTerminal = { block: body.loweredFunc.body.entry, id: makeInstructionId(0), - kind: "goto", - variant: GotoVariant.Break, + kind: "label", + fallthrough: continuationBlockId, loc: block.terminal.loc, }; block.terminal = newTerminal; - // If the final terminal type has a fallthrough, update it to point to the - // continuation block - const terminalBlock = getTerminalBlock( - body.loweredFunc.body, - body.loweredFunc.body.entry - ); - switch (terminalBlock.terminal.kind) { - case "if": - case "switch": - case "label": { - // These terminals can all appear as the final top-level terminal - // *and* have fallthroughs. If they are final, their fallthrough - // must be updated to point to the continuation block to main - // proper CFG structure (a block that succeeds all branches of a conditional - // must be marked as that conditional's fallthrough) - terminalBlock.terminal.fallthrough = continuationBlockId; - break; - } - case "return": - case "throw": { - // These can appear as the final top-level terminal - break; - } - // These all have non-nullable fallthroughs: there is always some code in the - // CFG that succeeds them which we should find instead - case "optional": - case "ternary": - case "logical": - case "while": - case "for": - case "for-of": - case "do-while": - // These are invalid terminals for a top-level block - case "branch": - case "goto": - case "unsupported": { - CompilerError.invariant( - `Unexpected final top-level terminal`, - terminalBlock.terminal.loc, - `Found ${terminalBlock.terminal.kind}, expected one of if, switch, label, return, or throw` - ); - } - default: { - assertExhaustive( - terminalBlock.terminal, - `Unexpected terminal kind '${ - (terminalBlock.terminal as any).kind - }'` - ); - } - } - // We store the result in the useMemo temporary const result = instr.lvalue; @@ -241,78 +188,6 @@ export function inlineUseMemo(fn: HIRFunction): void { } } -// Finds the final top-level terminal node for a CFG, by following any -// fallthrough nodes. -function getTerminalBlock(cfg: HIR, start: BlockId): BasicBlock { - let current = cfg.blocks.get(start)!; - while (true) { - const { terminal } = current; - switch (terminal.kind) { - case "if": { - if ( - terminal.fallthrough !== null && - terminal.fallthrough === terminal.alternate - ) { - // Here we don't know if the fallthrough and alternate are the same because there was - // no alternate or because both the alternate exists and the fallthrough is just unreachable - // So we check if the fallthrough returns/throws (the if is the final top-level terminal) - // or whether execution actually may continue. - const fallthrough = getTerminalBlock(cfg, terminal.fallthrough); - if ( - fallthrough.terminal.kind === "return" || - fallthrough.terminal.kind === "throw" - ) { - return current; - } else { - current = fallthrough; - continue; - } - } else { - return current; - } - } - case "switch": - case "label": { - if (terminal.fallthrough !== null) { - current = cfg.blocks.get(terminal.fallthrough)!; - continue; - } else { - return current; - } - } - case "optional": - case "ternary": - case "logical": - case "while": - case "for": - case "for-of": - case "do-while": { - current = cfg.blocks.get(terminal.fallthrough)!; - continue; - } - case "return": - case "throw": { - return current; - } - case "unsupported": - case "branch": - case "goto": { - CompilerError.invariant( - `Unexpected block terminal`, - terminal.loc, - `Top-level blocks may not end in a ${terminal.kind} terminal` - ); - } - default: { - assertExhaustive( - terminal, - `Unexpected terminal kind '${(terminal as any).kind}'` - ); - } - } - } -} - /** * Rewrites the block so that all `return` terminals are replaced: * * Add a StoreLocal = diff --git a/compiler/forget/src/__tests__/fixtures/compiler/useMemo-if-else-multiple-return.expect.md b/compiler/forget/src/__tests__/fixtures/compiler/useMemo-if-else-multiple-return.expect.md index ff2405d0f9..47d8a8f0bc 100644 --- a/compiler/forget/src/__tests__/fixtures/compiler/useMemo-if-else-multiple-return.expect.md +++ b/compiler/forget/src/__tests__/fixtures/compiler/useMemo-if-else-multiple-return.expect.md @@ -21,18 +21,20 @@ import { unstable_useMemoCache as useMemoCache } from "react"; function Component(props) { const $ = useMemoCache(4); let t21 = undefined; - if (props.cond) { - const c_0 = $[0] !== props.a; - let t0; - if (c_0) { - t0 = makeObject(props.a); - $[0] = props.a; - $[1] = t0; - } else { - t0 = $[1]; + bb7: { + if (props.cond) { + const c_0 = $[0] !== props.a; + let t0; + if (c_0) { + t0 = makeObject(props.a); + $[0] = props.a; + $[1] = t0; + } else { + t0 = $[1]; + } + t21 = t0; + break bb7; } - t21 = t0; - } else { const c_2 = $[2] !== props.b; let t1; if (c_2) { diff --git a/compiler/forget/src/__tests__/fixtures/compiler/useMemo-inlining-block-return.expect.md b/compiler/forget/src/__tests__/fixtures/compiler/useMemo-inlining-block-return.expect.md index 06b74a672e..e14c4179a0 100644 --- a/compiler/forget/src/__tests__/fixtures/compiler/useMemo-inlining-block-return.expect.md +++ b/compiler/forget/src/__tests__/fixtures/compiler/useMemo-inlining-block-return.expect.md @@ -20,18 +20,20 @@ import { unstable_useMemoCache as useMemoCache } from "react"; function component(a, b) { const $ = useMemoCache(2); let t14 = undefined; - if (a) { - const c_0 = $[0] !== b; - let t0; - if (c_0) { - t0 = { b }; - $[0] = b; - $[1] = t0; - } else { - t0 = $[1]; + bb6: { + if (a) { + const c_0 = $[0] !== b; + let t0; + if (c_0) { + t0 = { b }; + $[0] = b; + $[1] = t0; + } else { + t0 = $[1]; + } + t14 = t0; + break bb6; } - t14 = t0; - } else { t14 = undefined; } const x = t14; diff --git a/compiler/forget/src/__tests__/fixtures/compiler/useMemo-multiple-if-else.expect.md b/compiler/forget/src/__tests__/fixtures/compiler/useMemo-multiple-if-else.expect.md index 7c8254cb6a..c99ce367cb 100644 --- a/compiler/forget/src/__tests__/fixtures/compiler/useMemo-multiple-if-else.expect.md +++ b/compiler/forget/src/__tests__/fixtures/compiler/useMemo-multiple-if-else.expect.md @@ -24,25 +24,30 @@ function Component(props) { ```javascript import { unstable_useMemoCache as useMemoCache } from "react"; function Component(props) { - const $ = useMemoCache(2); - const c_0 = $[0] !== props; - let t32; - if (c_0) { - const y = []; - if (props.cond) { - y.push(props.a); - } - t32 = undefined; - if (props.cond2) { - t32 = y; - } else { + const $ = useMemoCache(3); + let t32 = undefined; + bb9: { + const c_0 = $[0] !== props; + let y; + if (c_0) { + y = []; + if (props.cond) { + y.push(props.a); + } + if (props.cond2) { + t32 = y; + break bb9; + } + y.push(props.b); - t32 = y; + $[0] = props; + $[1] = y; + $[2] = t32; + } else { + y = $[1]; + t32 = $[2]; } - $[0] = props; - $[1] = t32; - } else { - t32 = $[1]; + t32 = y; } const x = t32; return x; diff --git a/compiler/forget/src/__tests__/fixtures/compiler/useMemo-switch-no-fallthrough.expect.md b/compiler/forget/src/__tests__/fixtures/compiler/useMemo-switch-no-fallthrough.expect.md index 25f584d49e..e6bbc878ef 100644 --- a/compiler/forget/src/__tests__/fixtures/compiler/useMemo-switch-no-fallthrough.expect.md +++ b/compiler/forget/src/__tests__/fixtures/compiler/useMemo-switch-no-fallthrough.expect.md @@ -23,13 +23,15 @@ function Component(props) { ```javascript function Component(props) { let t18 = undefined; - bb8: switch (props.key) { - case "key": { - t18 = props.value; - break bb8; - } - default: { - t18 = props.defaultValue; + bb8: { + switch (props.key) { + case "key": { + t18 = props.value; + break bb8; + } + default: { + t18 = props.defaultValue; + } } } const x = t18;