diff --git a/compiler/forget/src/CompilerPipeline.ts b/compiler/forget/src/CompilerPipeline.ts index 4e0a8bb4f8..fba60ba925 100644 --- a/compiler/forget/src/CompilerPipeline.ts +++ b/compiler/forget/src/CompilerPipeline.ts @@ -21,6 +21,7 @@ import { dropMemoCalls, inferMutableRanges, inferReferenceEffects, + inlineUseMemo, } from "./Inference"; import { constantPropagation, deadCodeElimination } from "./Optimization"; import { @@ -60,6 +61,9 @@ export function* run( const hir = lower(func, env).unwrap(); yield log({ kind: "hir", name: "HIR", value: hir }); + inlineUseMemo(hir); + yield log({ kind: "hir", name: "RewriteUseMemo", value: hir }); + mergeConsecutiveBlocks(hir); yield log({ kind: "hir", name: "MergeConsecutiveBlocks", value: hir }); diff --git a/compiler/forget/src/HIR/MergeConsecutiveBlocks.ts b/compiler/forget/src/HIR/MergeConsecutiveBlocks.ts index 86b48a59be..dc1351eac7 100644 --- a/compiler/forget/src/HIR/MergeConsecutiveBlocks.ts +++ b/compiler/forget/src/HIR/MergeConsecutiveBlocks.ts @@ -13,7 +13,7 @@ import { HIRFunction, Instruction, } from "./HIR"; -import { removeUnreachableFallthroughs } from "./HIRBuilder"; +import { markPredecessors, removeUnreachableFallthroughs } from "./HIRBuilder"; /** * Merges sequences of blocks that will always execute consecutively — @@ -85,6 +85,7 @@ export function mergeConsecutiveBlocks(fn: HIRFunction): void { merged.merge(block.id, predecessorId); fn.body.blocks.delete(block.id); } + markPredecessors(fn.body); removeUnreachableFallthroughs(fn.body); } diff --git a/compiler/forget/src/Inference/InlineUseMemo.ts b/compiler/forget/src/Inference/InlineUseMemo.ts new file mode 100644 index 0000000000..28f45446a8 --- /dev/null +++ b/compiler/forget/src/Inference/InlineUseMemo.ts @@ -0,0 +1,343 @@ +/** + * Copyright (c) Facebook, Inc. and its affiliates. + * + * This source code is licensed under the MIT license found in the + * LICENSE file in the root directory of this source tree. + */ + +import { CompilerError } from "../CompilerError"; +import { + BasicBlock, + BlockId, + Effect, + Environment, + FunctionExpression, + GotoTerminal, + GotoVariant, + HIR, + HIRFunction, + IdentifierId, + InstructionKind, + makeInstructionId, + makeType, + Place, + reversePostorderBlocks, + shrink, +} from "../HIR"; +import { markInstructionIds, markPredecessors } from "../HIR/HIRBuilder"; +import { assertExhaustive, retainWhere } from "../Utils/utils"; + +/** + * Rewrites `useMemo()` calls, rewriting so that the lambda body becomes part of the + * outer block's instructions. + * + * Example: + * + * ```javascript + * // Before + * const x = useMemo(() => foo(y, z), [y, z]) + * + * // After + * const x = foo(y, z); + * ``` + * + * The main challenge is dealing with the possibility of complex control flow within + * the lambda body. The approach is roughly: + * - split the block with the useMemo call in two: + * - the first block is everything up to the memo call plus the lambda body + * - the second block is everything after the memo call + * - use the temporary from the useMemo call result value as the place to store + * the useMemo result + * - for every return terminal in the lambda body: + * - add a StoreLocal to the temporary, assigning the return value + * - replace the terminal w a goto to the second block + * + * NOTE: *this pass must be run prior to EnterSSA*. Prior to entering SSA form identifiers + * in the top-level function and any function expressions will have consistent + * correlation between `Identifier` instances and IdentifierIds. After entering SSA + * form we drop this correspondence. It's much easier to write this inlining pass + * without having to worry about SSA form. + */ +export function inlineUseMemo(fn: HIRFunction): void { + // Track all function expressions in case they appear as the argument to a useMemo + const functions = new Map(); + // Track all references to `useMemo` + const useMemoGlobals = new Set(); + // Identifiers (lvalues) for known useMemo functions, so that we can prune them + // at the end of the pass + const useMemoFunctions = new Set(); + + // Iterate the *existing* blocks from the outer component to find useMemo calls + // and inline them. During iteration we will modify `fn` (by inlining the CFG + // of useMemo callbacks) so we explicitly copy references to just the original + // function's blocks first. As blocks are split to make room for useMemo calls, + // the split portions of the blocks will be added to this queue. + const queue = Array.from(fn.body.blocks.values()); + queue: for (const block of queue) { + for (let ii = 0; ii < block.instructions.length; ii++) { + const instr = block.instructions[ii]!; + switch (instr.value.kind) { + case "LoadGlobal": { + if (instr.value.name === "useMemo") { + useMemoGlobals.add(instr.lvalue.identifier.id); + } + break; + } + case "FunctionExpression": { + functions.set(instr.lvalue.identifier.id, instr.value); + break; + } + case "CallExpression": { + if (useMemoGlobals.has(instr.value.callee.identifier.id)) { + const [lambda] = instr.value.args; + if (lambda.kind === "Spread") { + continue; + } + const body = functions.get(lambda.identifier.id); + if (body === undefined) { + CompilerError.invariant( + "Expected first argument to useMemo() to be a function expression", + fn.loc + ); + } + // We know this function is used for useMemo and can prune it later + useMemoFunctions.add(lambda.identifier.id); + + // Create a new block which will contain code following the useMemo call + const continuationBlockId = fn.env.nextBlockId; + const continuationBlock: BasicBlock = { + id: continuationBlockId, + instructions: block.instructions.slice(ii + 1), + kind: block.kind, + phis: new Set(), + preds: new Set(), + terminal: block.terminal, + }; + fn.body.blocks.set(continuationBlockId, continuationBlock); + + // Trim the original block to contain instructions up to (but not including) + // 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 = { + block: body.loweredFunc.body.entry, + id: makeInstructionId(0), + kind: "goto", + variant: GotoVariant.Break, + 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-call": + 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 + }'` + ); + } + } + + // Rewrite blocks from the lambda to replace any `return` with a + // store the useMemo temporary and `goto` the continuation block + for (const [id, block] of body.loweredFunc.body.blocks) { + block.preds.clear(); + rewriteBlock(fn.env, block, continuationBlockId, instr.lvalue); + fn.body.blocks.set(id, block); + } + + // Ensure we visit the continuation block, since there may have been + // sequential useMemos that need to be visited. + queue.push(continuationBlock); + continue queue; + } + } + } + } + } + + if (useMemoFunctions.size !== 0) { + // Remove instructions that define lambdas which we inlined + for (const [, block] of fn.body.blocks) { + retainWhere( + block.instructions, + (instr) => !useMemoFunctions.has(instr.lvalue.identifier.id) + ); + } + + // If terminals have changed then blocks may have become newly unreachable. + // Re-run minification of the graph (incl reordering instruction ids) + shrink(fn.body); + reversePostorderBlocks(fn.body); + markInstructionIds(fn.body); + markPredecessors(fn.body); + } +} + +// 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-call": + 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 = + * * Replace the terminal with a Goto to + */ +function rewriteBlock( + env: Environment, + block: BasicBlock, + returnTarget: BlockId, + returnValue: Place +): void { + const { terminal } = block; + if (terminal.kind !== "return") { + return; + } + if (terminal.value !== null) { + block.instructions.push({ + id: makeInstructionId(0), + loc: terminal.loc, + lvalue: { + effect: Effect.Unknown, + identifier: { + id: env.nextIdentifierId, + mutableRange: { + start: makeInstructionId(0), + end: makeInstructionId(0), + }, + name: null, + scope: null, + type: makeType(), + }, + kind: "Identifier", + loc: terminal.loc, + }, + value: { + kind: "StoreLocal", + lvalue: { kind: InstructionKind.Const, place: { ...returnValue } }, + value: terminal.value, + loc: terminal.loc, + }, + }); + } + block.terminal = { + kind: "goto", + block: returnTarget, + id: makeInstructionId(0), + variant: GotoVariant.Break, + loc: block.terminal.loc, + }; +} diff --git a/compiler/forget/src/Inference/index.ts b/compiler/forget/src/Inference/index.ts index 653e8e1f7e..8892a8e4e8 100644 --- a/compiler/forget/src/Inference/index.ts +++ b/compiler/forget/src/Inference/index.ts @@ -5,7 +5,8 @@ * LICENSE file in the root directory of this source tree. */ +export { default as analyseFunctions } from "./AnalyseFunctions"; export { default as dropMemoCalls } from "./DropMemoCalls"; export { inferMutableRanges } from "./InferMutableRanges"; -export { default as analyseFunctions } from "./AnalyseFunctions"; export { default as inferReferenceEffects } from "./InferReferenceEffects"; +export { inlineUseMemo } from "./InlineUseMemo"; 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 new file mode 100644 index 0000000000..82ece6e3f1 --- /dev/null +++ b/compiler/forget/src/__tests__/fixtures/compiler/useMemo-if-else-multiple-return.expect.md @@ -0,0 +1,56 @@ + +## Input + +```javascript +function Component(props) { + const x = useMemo(() => { + if (props.cond) { + return makeObject(props.a); + } + return makeObject(props.b); + }); + return x; +} + +``` + +## Code + +```javascript +function Component(props) { + const $ = React.unstable_useMemoCache(5); + 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]; + } + let t1; + if ($[2] === Symbol.for("react.memo_cache_sentinel")) { + t1 = t0; + $[2] = t1; + } else { + t1 = $[2]; + } + } else { + const c_3 = $[3] !== props.b; + let t2; + if (c_3) { + t2 = makeObject(props.b); + $[3] = props.b; + $[4] = t2; + } else { + t2 = $[4]; + } + t1 = t2; + } + const x = t1; + return x; +} + +``` + \ No newline at end of file diff --git a/compiler/forget/src/__tests__/fixtures/compiler/useMemo-if-else-multiple-return.js b/compiler/forget/src/__tests__/fixtures/compiler/useMemo-if-else-multiple-return.js new file mode 100644 index 0000000000..a39f3b2826 --- /dev/null +++ b/compiler/forget/src/__tests__/fixtures/compiler/useMemo-if-else-multiple-return.js @@ -0,0 +1,9 @@ +function Component(props) { + const x = useMemo(() => { + if (props.cond) { + return makeObject(props.a); + } + return makeObject(props.b); + }); + return x; +} diff --git a/compiler/forget/src/__tests__/fixtures/compiler/useMemo-independently-memoizeable.expect.md b/compiler/forget/src/__tests__/fixtures/compiler/useMemo-independently-memoizeable.expect.md index fb6918b129..df99148842 100644 --- a/compiler/forget/src/__tests__/fixtures/compiler/useMemo-independently-memoizeable.expect.md +++ b/compiler/forget/src/__tests__/fixtures/compiler/useMemo-independently-memoizeable.expect.md @@ -18,43 +18,52 @@ function Component(props) { ```javascript function Component(props) { - const $ = React.unstable_useMemoCache(7); + const $ = React.unstable_useMemoCache(10); const c_0 = $[0] !== props.a; let t0; if (c_0) { - t0 = () => { - const items = []; - const a = makeObject(props.a); - const b = makeObject(props.b); - return [a, b]; - }; + t0 = makeObject(props.a); $[0] = props.a; $[1] = t0; } else { t0 = $[1]; } - const c_2 = $[2] !== t0; + const a = t0; + const c_2 = $[2] !== props.b; let t1; if (c_2) { - t1 = t0(); - $[2] = t0; + t1 = makeObject(props.b); + $[2] = props.b; $[3] = t1; } else { t1 = $[3]; } - const [a_0, b_0] = t1; - const c_4 = $[4] !== a_0; - const c_5 = $[5] !== b_0; + const b = t1; + const c_4 = $[4] !== a; + const c_5 = $[5] !== b; let t2; if (c_4 || c_5) { - t2 = [a_0, b_0]; - $[4] = a_0; - $[5] = b_0; + t2 = [a, b]; + $[4] = a; + $[5] = b; $[6] = t2; } else { t2 = $[6]; } - return t2; + const t54 = t2; + const [a_0, b_0] = t54; + const c_7 = $[7] !== a_0; + const c_8 = $[8] !== b_0; + let t3; + if (c_7 || c_8) { + t3 = [a_0, b_0]; + $[7] = a_0; + $[8] = b_0; + $[9] = t3; + } else { + t3 = $[9]; + } + return t3; } ``` diff --git a/compiler/forget/src/__tests__/fixtures/compiler/useMemo-labeled-statement-unconditional-return.expect.md b/compiler/forget/src/__tests__/fixtures/compiler/useMemo-labeled-statement-unconditional-return.expect.md new file mode 100644 index 0000000000..8e9d522789 --- /dev/null +++ b/compiler/forget/src/__tests__/fixtures/compiler/useMemo-labeled-statement-unconditional-return.expect.md @@ -0,0 +1,26 @@ + +## Input + +```javascript +function Component(props) { + const x = useMemo(() => { + label: { + return props.value; + } + }); + return x; +} + +``` + +## Code + +```javascript +function Component(props) { + const t19 = props.value; + const x = t19; + return x; +} + +``` + \ No newline at end of file diff --git a/compiler/forget/src/__tests__/fixtures/compiler/useMemo-labeled-statement-unconditional-return.js b/compiler/forget/src/__tests__/fixtures/compiler/useMemo-labeled-statement-unconditional-return.js new file mode 100644 index 0000000000..a23a7bda64 --- /dev/null +++ b/compiler/forget/src/__tests__/fixtures/compiler/useMemo-labeled-statement-unconditional-return.js @@ -0,0 +1,8 @@ +function Component(props) { + const x = useMemo(() => { + label: { + return props.value; + } + }); + return x; +} diff --git a/compiler/forget/src/__tests__/fixtures/compiler/useMemo-logical.expect.md b/compiler/forget/src/__tests__/fixtures/compiler/useMemo-logical.expect.md new file mode 100644 index 0000000000..e30d844490 --- /dev/null +++ b/compiler/forget/src/__tests__/fixtures/compiler/useMemo-logical.expect.md @@ -0,0 +1,22 @@ + +## Input + +```javascript +function Component(props) { + const x = useMemo(() => props.a && props.b); + return x; +} + +``` + +## Code + +```javascript +function Component(props) { + const t32 = props.a && props.b; + const x = t32; + return x; +} + +``` + \ No newline at end of file diff --git a/compiler/forget/src/__tests__/fixtures/compiler/useMemo-logical.js b/compiler/forget/src/__tests__/fixtures/compiler/useMemo-logical.js new file mode 100644 index 0000000000..1bfbd0e3fa --- /dev/null +++ b/compiler/forget/src/__tests__/fixtures/compiler/useMemo-logical.js @@ -0,0 +1,4 @@ +function Component(props) { + const x = useMemo(() => props.a && props.b); + return x; +} 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 new file mode 100644 index 0000000000..29e32088b8 --- /dev/null +++ b/compiler/forget/src/__tests__/fixtures/compiler/useMemo-multiple-if-else.expect.md @@ -0,0 +1,50 @@ + +## Input + +```javascript +function Component(props) { + const x = useMemo(() => { + let y = []; + if (props.cond) { + y.push(props.a); + } + if (props.cond2) { + return y; + } + y.push(props.b); + return y; + }); + return x; +} + +``` + +## Code + +```javascript +function Component(props) { + const $ = React.unstable_useMemoCache(2); + const c_0 = $[0] !== props; + let t0; + if (c_0) { + const y = []; + if (props.cond) { + y.push(props.a); + } + if (props.cond2) { + t0 = y; + } else { + y.push(props.b); + t0 = y; + } + $[0] = props; + $[1] = t0; + } else { + t0 = $[1]; + } + const x = t0; + return x; +} + +``` + \ No newline at end of file diff --git a/compiler/forget/src/__tests__/fixtures/compiler/useMemo-multiple-if-else.js b/compiler/forget/src/__tests__/fixtures/compiler/useMemo-multiple-if-else.js new file mode 100644 index 0000000000..24e10f21c9 --- /dev/null +++ b/compiler/forget/src/__tests__/fixtures/compiler/useMemo-multiple-if-else.js @@ -0,0 +1,14 @@ +function Component(props) { + const x = useMemo(() => { + let y = []; + if (props.cond) { + y.push(props.a); + } + if (props.cond2) { + return y; + } + y.push(props.b); + return y; + }); + return x; +} diff --git a/compiler/forget/src/__tests__/fixtures/compiler/useMemo-simple.expect.md b/compiler/forget/src/__tests__/fixtures/compiler/useMemo-simple.expect.md index 261051f078..0eaeb73317 100644 --- a/compiler/forget/src/__tests__/fixtures/compiler/useMemo-simple.expect.md +++ b/compiler/forget/src/__tests__/fixtures/compiler/useMemo-simple.expect.md @@ -13,36 +13,28 @@ function component(a) { ```javascript function component(a) { - const $ = React.unstable_useMemoCache(6); + const $ = React.unstable_useMemoCache(4); const c_0 = $[0] !== a; let t0; if (c_0) { - t0 = () => [a]; + t0 = [a]; $[0] = a; $[1] = t0; } else { t0 = $[1]; } - const c_2 = $[2] !== t0; + const t23 = t0; + const x = t23; + const c_2 = $[2] !== x; let t1; if (c_2) { - t1 = t0(); - $[2] = t0; + t1 = ; + $[2] = x; $[3] = t1; } else { t1 = $[3]; } - const x = t1; - const c_4 = $[4] !== x; - let t2; - if (c_4) { - t2 = ; - $[4] = x; - $[5] = t2; - } else { - t2 = $[5]; - } - return t2; + return t1; } ``` 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 new file mode 100644 index 0000000000..35946f0393 --- /dev/null +++ b/compiler/forget/src/__tests__/fixtures/compiler/useMemo-switch-no-fallthrough.expect.md @@ -0,0 +1,39 @@ + +## Input + +```javascript +function Component(props) { + const x = useMemo(() => { + switch (props.key) { + case "key": { + return props.value; + } + default: { + return props.defaultValue; + } + } + }); + return x; +} + +``` + +## Code + +```javascript +function Component(props) { + bb8: switch (props.key) { + case "key": { + const t28 = props.value; + break bb8; + } + default: { + const t28 = props.defaultValue; + } + } + const x = t28; + return x; +} + +``` + \ No newline at end of file diff --git a/compiler/forget/src/__tests__/fixtures/compiler/useMemo-switch-no-fallthrough.js b/compiler/forget/src/__tests__/fixtures/compiler/useMemo-switch-no-fallthrough.js new file mode 100644 index 0000000000..74011509fb --- /dev/null +++ b/compiler/forget/src/__tests__/fixtures/compiler/useMemo-switch-no-fallthrough.js @@ -0,0 +1,13 @@ +function Component(props) { + const x = useMemo(() => { + switch (props.key) { + case "key": { + return props.value; + } + default: { + return props.defaultValue; + } + } + }); + return x; +}