From ff0b05848b514e347d93fe8145bae24a4d21dc40 Mon Sep 17 00:00:00 2001 From: Joe Savona Date: Thu, 7 Sep 2023 16:32:53 -0700 Subject: [PATCH] TryStatement: optimization pass, handle early returns Adds an optimization pass to prune unnecessary maybe-throw terminals, when the block can be proven not to throw. For now we're _very_ conservative about what instructions we consider not to throw. There isn't too much of an advantage in pruning further, either. This PR also updates BuildReactiveFunction to handle the possibility of early returns within try or catch blocks, making sure we don't hit the invariant of emitting the same block twice. --- .../src/Entrypoint/Pipeline.ts | 9 ++- .../src/HIR/BuildHIR.ts | 2 +- .../babel-plugin-react-forget/src/HIR/HIR.ts | 1 + .../src/Optimization/PruneMaybeThrows.ts | 62 +++++++++++++++++++ .../src/Optimization/index.ts | 1 + .../ReactiveScopes/BuildReactiveFunction.ts | 22 +++++-- .../compiler/try-catch-with-return.expect.md | 49 +++++++++++++++ .../compiler/try-catch-with-return.js | 14 +++++ .../packages/sprout/src/SproutTodoFilter.ts | 1 + 9 files changed, 154 insertions(+), 7 deletions(-) create mode 100644 compiler/packages/babel-plugin-react-forget/src/Optimization/PruneMaybeThrows.ts create mode 100644 compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/try-catch-with-return.expect.md create mode 100644 compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/try-catch-with-return.js diff --git a/compiler/packages/babel-plugin-react-forget/src/Entrypoint/Pipeline.ts b/compiler/packages/babel-plugin-react-forget/src/Entrypoint/Pipeline.ts index 68179ba8d1..14e6506e69 100644 --- a/compiler/packages/babel-plugin-react-forget/src/Entrypoint/Pipeline.ts +++ b/compiler/packages/babel-plugin-react-forget/src/Entrypoint/Pipeline.ts @@ -26,7 +26,11 @@ import { inferReferenceEffects, inlineUseMemo, } from "../Inference"; -import { constantPropagation, deadCodeElimination } from "../Optimization"; +import { + constantPropagation, + deadCodeElimination, + pruneMaybeThrows, +} from "../Optimization"; import { CodegenFunction, alignReactiveScopesToBlockScopes, @@ -78,6 +82,9 @@ export function* run( const hir = lower(func, env).unwrap(); yield log({ kind: "hir", name: "HIR", value: hir }); + pruneMaybeThrows(hir); + yield log({ kind: "hir", name: "PruneMaybeThrows", value: hir }); + if (config?.inlineUseMemo) { inlineUseMemo(hir); yield log({ kind: "hir", name: "RewriteUseMemo", value: hir }); diff --git a/compiler/packages/babel-plugin-react-forget/src/HIR/BuildHIR.ts b/compiler/packages/babel-plugin-react-forget/src/HIR/BuildHIR.ts index d52d6b835a..2de4e3f054 100644 --- a/compiler/packages/babel-plugin-react-forget/src/HIR/BuildHIR.ts +++ b/compiler/packages/babel-plugin-react-forget/src/HIR/BuildHIR.ts @@ -955,7 +955,7 @@ function lowerStatement( return { kind: "goto", block: continuationBlock.id, - variant: GotoVariant.Break, + variant: GotoVariant.Try, id: makeInstructionId(0), loc: block.node.loc ?? GeneratedSource, }; 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 373fe1e08e..9e1a51b014 100644 --- a/compiler/packages/babel-plugin-react-forget/src/HIR/HIR.ts +++ b/compiler/packages/babel-plugin-react-forget/src/HIR/HIR.ts @@ -339,6 +339,7 @@ export type GotoTerminal = { export enum GotoVariant { Break = "Break", Continue = "Continue", + Try = "Try", } export type IfTerminal = { diff --git a/compiler/packages/babel-plugin-react-forget/src/Optimization/PruneMaybeThrows.ts b/compiler/packages/babel-plugin-react-forget/src/Optimization/PruneMaybeThrows.ts new file mode 100644 index 0000000000..abfe0dbd03 --- /dev/null +++ b/compiler/packages/babel-plugin-react-forget/src/Optimization/PruneMaybeThrows.ts @@ -0,0 +1,62 @@ +/** + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * This source code is licensed under the MIT license found in the + * LICENSE file in the root directory of this source tree. + */ + +import { + GotoVariant, + HIRFunction, + Instruction, + mergeConsecutiveBlocks, +} from "../HIR"; + +/** + * This pass prunes `maybe-throw` terminals for blocks that can provably *never* throw. + * For now this is very conservative, and only affects blocks with primitives or + * array/object literals. Even a variable reference could throw bc of the TDZ. + */ +export function pruneMaybeThrows(fn: HIRFunction): void { + const didPrune = pruneMaybeThrowsImpl(fn); + if (didPrune) { + mergeConsecutiveBlocks(fn); + } +} + +function pruneMaybeThrowsImpl(fn: HIRFunction): boolean { + let hasChanges = false; + for (const [_, block] of fn.body.blocks) { + const terminal = block.terminal; + if (terminal.kind !== "maybe-throw") { + continue; + } + const canThrow = block.instructions.some((instr) => + instructionMayThrow(instr) + ); + if (!canThrow) { + hasChanges = true; + block.terminal = { + kind: "goto", + block: terminal.continuation, + variant: GotoVariant.Break, + id: terminal.id, + loc: terminal.loc, + }; + } + } + return hasChanges; +} + +function instructionMayThrow(instr: Instruction): boolean { + switch (instr.value.kind) { + case "Primitive": + case "ArrayExpression": + case "ObjectExpression": { + return false; + } + default: { + return true; + } + } +} diff --git a/compiler/packages/babel-plugin-react-forget/src/Optimization/index.ts b/compiler/packages/babel-plugin-react-forget/src/Optimization/index.ts index 440157b4f0..4d2cffdbdf 100644 --- a/compiler/packages/babel-plugin-react-forget/src/Optimization/index.ts +++ b/compiler/packages/babel-plugin-react-forget/src/Optimization/index.ts @@ -7,3 +7,4 @@ export { constantPropagation } from "./ConstantPropagation"; export { deadCodeElimination } from "./DeadCodeElimination"; +export { pruneMaybeThrows } from "./PruneMaybeThrows"; diff --git a/compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/BuildReactiveFunction.ts b/compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/BuildReactiveFunction.ts index fc5942bc68..626233a837 100644 --- a/compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/BuildReactiveFunction.ts +++ b/compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/BuildReactiveFunction.ts @@ -627,6 +627,9 @@ class Driver { } break; } + case GotoVariant.Try: { + break; + } default: { assertExhaustive( terminal.variant, @@ -639,10 +642,12 @@ class Driver { case "maybe-throw": { // ReactiveFunction does not explicit model maybe-throw semantics, // so these terminals flatten away - this.visitBlock( - this.cx.ir.blocks.get(terminal.continuation)!, - blockValue - ); + if (!this.cx.isScheduled(terminal.continuation)) { + this.visitBlock( + this.cx.ir.blocks.get(terminal.continuation)!, + blockValue + ); + } break; } case "try": { @@ -655,6 +660,7 @@ class Driver { const scheduleId = this.cx.schedule(fallthroughId, "if"); scheduleIds.push(scheduleId); } + this.cx.scheduleCatchHandler(terminal.handler); const block = this.traverseBlock( this.cx.ir.blocks.get(terminal.block)! @@ -1091,6 +1097,8 @@ class Context { */ #scheduled: Set = new Set(); + #catchHandlers: Set = new Set(); + /** * Represents which control flow operations are currently in scope, with the innermost * scope last. Roughly speaking, the last ControlFlowTarget on the stack indicates where @@ -1108,6 +1116,10 @@ class Context { return this.ir.blocks.get(id)!; } + scheduleCatchHandler(block: BlockId): void { + this.#catchHandlers.add(block); + } + /** * Record that the given block will be emitted (eg by the codegen of a parent node) * so that child nodes can avoid re-emitting it. @@ -1194,7 +1206,7 @@ class Context { * Check if the given @param block is scheduled or not. */ isScheduled(block: BlockId): boolean { - return this.#scheduled.has(block); + return this.#scheduled.has(block) || this.#catchHandlers.has(block); } /** diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/try-catch-with-return.expect.md b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/try-catch-with-return.expect.md new file mode 100644 index 0000000000..c4efbbd002 --- /dev/null +++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/try-catch-with-return.expect.md @@ -0,0 +1,49 @@ + +## Input + +```javascript +// @debug +function Component(props) { + let x = []; + try { + const y = foo(); + if (y == null) { + return; + } + x.push(bar(y)); + } catch { + return null; + } + return x; +} + +``` + +## Code + +```javascript +import { unstable_useMemoCache as useMemoCache } from "react"; // @debug +function Component(props) { + const $ = useMemoCache(1); + let x; + if ($[0] === Symbol.for("react.memo_cache_sentinel")) { + x = []; + try { + const y = foo(); + if (y == null) { + return; + } + + x.push(bar(y)); + } catch { + return null; + } + $[0] = x; + } else { + x = $[0]; + } + return x; +} + +``` + \ No newline at end of file diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/try-catch-with-return.js b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/try-catch-with-return.js new file mode 100644 index 0000000000..5014f6fefa --- /dev/null +++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/try-catch-with-return.js @@ -0,0 +1,14 @@ +// @debug +function Component(props) { + let x = []; + try { + const y = foo(); + if (y == null) { + return; + } + x.push(bar(y)); + } catch { + return null; + } + return x; +} diff --git a/compiler/packages/sprout/src/SproutTodoFilter.ts b/compiler/packages/sprout/src/SproutTodoFilter.ts index 75988da6fe..7398ebbcd8 100644 --- a/compiler/packages/sprout/src/SproutTodoFilter.ts +++ b/compiler/packages/sprout/src/SproutTodoFilter.ts @@ -415,6 +415,7 @@ const skipFilter = new Set([ "try-catch-mutate-outer-value", "try-catch-within-mutable-range", "try-catch", + "try-catch-with-return", // TODO: 🌲 "forest-basic",