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.
Joe Savona committed
Sep 7, 2023 at 16:32 UTC
ff0b05848b514e347d93fe8145bae24a4d21dc40
9 files changed
+154
-7
compiler/packages/babel-plugin-react-forget/src/Entrypoint/Pipeline.ts
+8
-1
@@ -26,7 +26,11 @@ import {
26
inferReferenceEffects,
27
inlineUseMemo,
28
} from "../Inference";
29
-import { constantPropagation, deadCodeElimination } from "../Optimization";
29
+import {
30
+ constantPropagation,
31
+ deadCodeElimination,
32
+ pruneMaybeThrows,
33
+} from "../Optimization";
34
import {
35
CodegenFunction,
36
alignReactiveScopesToBlockScopes,
@@ -78,6 +82,9 @@ export function* run(
82
const hir = lower(func, env).unwrap();
83
yield log({ kind: "hir", name: "HIR", value: hir });
84
85
+ pruneMaybeThrows(hir);
86
+ yield log({ kind: "hir", name: "PruneMaybeThrows", value: hir });
87
+
88
if (config?.inlineUseMemo) {
89
inlineUseMemo(hir);
90
yield log({ kind: "hir", name: "RewriteUseMemo", value: hir });
compiler/packages/babel-plugin-react-forget/src/HIR/BuildHIR.ts
+1
-1
@@ -955,7 +955,7 @@ function lowerStatement(
955
return {
956
kind: "goto",
957
block: continuationBlock.id,
958
- variant: GotoVariant.Break,
958
+ variant: GotoVariant.Try,
959
id: makeInstructionId(0),
960
loc: block.node.loc ?? GeneratedSource,
961
};
compiler/packages/babel-plugin-react-forget/src/HIR/HIR.ts
+1
@@ -339,6 +339,7 @@ export type GotoTerminal = {
339
export enum GotoVariant {
340
Break = "Break",
341
Continue = "Continue",
342
+ Try = "Try",
343
}
344
345
export type IfTerminal = {
compiler/packages/babel-plugin-react-forget/src/Optimization/PruneMaybeThrows.ts
new
+62
@@ -0,0 +1,62 @@
1
+/**
2
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
3
+ *
4
+ * This source code is licensed under the MIT license found in the
5
+ * LICENSE file in the root directory of this source tree.
6
+ */
7
+
8
+import {
9
+ GotoVariant,
10
+ HIRFunction,
11
+ Instruction,
12
+ mergeConsecutiveBlocks,
13
+} from "../HIR";
14
+
15
+/**
16
+ * This pass prunes `maybe-throw` terminals for blocks that can provably *never* throw.
17
+ * For now this is very conservative, and only affects blocks with primitives or
18
+ * array/object literals. Even a variable reference could throw bc of the TDZ.
19
+ */
20
+export function pruneMaybeThrows(fn: HIRFunction): void {
21
+ const didPrune = pruneMaybeThrowsImpl(fn);
22
+ if (didPrune) {
23
+ mergeConsecutiveBlocks(fn);
24
+ }
25
+}
26
+
27
+function pruneMaybeThrowsImpl(fn: HIRFunction): boolean {
28
+ let hasChanges = false;
29
+ for (const [_, block] of fn.body.blocks) {
30
+ const terminal = block.terminal;
31
+ if (terminal.kind !== "maybe-throw") {
32
+ continue;
33
+ }
34
+ const canThrow = block.instructions.some((instr) =>
35
+ instructionMayThrow(instr)
36
+ );
37
+ if (!canThrow) {
38
+ hasChanges = true;
39
+ block.terminal = {
40
+ kind: "goto",
41
+ block: terminal.continuation,
42
+ variant: GotoVariant.Break,
43
+ id: terminal.id,
44
+ loc: terminal.loc,
45
+ };
46
+ }
47
+ }
48
+ return hasChanges;
49
+}
50
+
51
+function instructionMayThrow(instr: Instruction): boolean {
52
+ switch (instr.value.kind) {
53
+ case "Primitive":
54
+ case "ArrayExpression":
55
+ case "ObjectExpression": {
56
+ return false;
57
+ }
58
+ default: {
59
+ return true;
60
+ }
61
+ }
62
+}
compiler/packages/babel-plugin-react-forget/src/Optimization/index.ts
+1
@@ -7,3 +7,4 @@
7
8
export { constantPropagation } from "./ConstantPropagation";
9
export { deadCodeElimination } from "./DeadCodeElimination";
10
+export { pruneMaybeThrows } from "./PruneMaybeThrows";
compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/BuildReactiveFunction.ts
+17
-5
@@ -627,6 +627,9 @@ class Driver {
627
}
628
break;
629
}
630
+ case GotoVariant.Try: {
631
+ break;
632
+ }
633
default: {
634
assertExhaustive(
635
terminal.variant,
@@ -639,10 +642,12 @@ class Driver {
642
case "maybe-throw": {
643
// ReactiveFunction does not explicit model maybe-throw semantics,
644
// so these terminals flatten away
642
- this.visitBlock(
643
- this.cx.ir.blocks.get(terminal.continuation)!,
644
- blockValue
645
- );
645
+ if (!this.cx.isScheduled(terminal.continuation)) {
646
+ this.visitBlock(
647
+ this.cx.ir.blocks.get(terminal.continuation)!,
648
+ blockValue
649
+ );
650
+ }
651
break;
652
}
653
case "try": {
@@ -655,6 +660,7 @@ class Driver {
660
const scheduleId = this.cx.schedule(fallthroughId, "if");
661
scheduleIds.push(scheduleId);
662
}
663
+ this.cx.scheduleCatchHandler(terminal.handler);
664
665
const block = this.traverseBlock(
666
this.cx.ir.blocks.get(terminal.block)!
@@ -1091,6 +1097,8 @@ class Context {
1097
*/
1098
#scheduled: Set<BlockId> = new Set();
1099
1100
+ #catchHandlers: Set<BlockId> = new Set();
1101
+
1102
/**
1103
* Represents which control flow operations are currently in scope, with the innermost
1104
* scope last. Roughly speaking, the last ControlFlowTarget on the stack indicates where
@@ -1108,6 +1116,10 @@ class Context {
1116
return this.ir.blocks.get(id)!;
1117
}
1118
1119
+ scheduleCatchHandler(block: BlockId): void {
1120
+ this.#catchHandlers.add(block);
1121
+ }
1122
+
1123
/**
1124
* Record that the given block will be emitted (eg by the codegen of a parent node)
1125
* so that child nodes can avoid re-emitting it.
@@ -1194,7 +1206,7 @@ class Context {
1206
* Check if the given @param block is scheduled or not.
1207
*/
1208
isScheduled(block: BlockId): boolean {
1197
- return this.#scheduled.has(block);
1209
+ return this.#scheduled.has(block) || this.#catchHandlers.has(block);
1210
}
1211
1212
/**
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/try-catch-with-return.expect.md
new
+49
@@ -0,0 +1,49 @@
1
+
2
+## Input
3
+
4
+```javascript
5
+// @debug
6
+function Component(props) {
7
+ let x = [];
8
+ try {
9
+ const y = foo();
10
+ if (y == null) {
11
+ return;
12
+ }
13
+ x.push(bar(y));
14
+ } catch {
15
+ return null;
16
+ }
17
+ return x;
18
+}
19
+
20
+```
21
+
22
+## Code
23
+
24
+```javascript
25
+import { unstable_useMemoCache as useMemoCache } from "react"; // @debug
26
+function Component(props) {
27
+ const $ = useMemoCache(1);
28
+ let x;
29
+ if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
30
+ x = [];
31
+ try {
32
+ const y = foo();
33
+ if (y == null) {
34
+ return;
35
+ }
36
+
37
+ x.push(bar(y));
38
+ } catch {
39
+ return null;
40
+ }
41
+ $[0] = x;
42
+ } else {
43
+ x = $[0];
44
+ }
45
+ return x;
46
+}
47
+
48
+```
49
+
\ No newline at end of file
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/try-catch-with-return.js
new
+14
@@ -0,0 +1,14 @@
1
+// @debug
2
+function Component(props) {
3
+ let x = [];
4
+ try {
5
+ const y = foo();
6
+ if (y == null) {
7
+ return;
8
+ }
9
+ x.push(bar(y));
10
+ } catch {
11
+ return null;
12
+ }
13
+ return x;
14
+}
compiler/packages/sprout/src/SproutTodoFilter.ts
+1
@@ -415,6 +415,7 @@ const skipFilter = new Set([
415
"try-catch-mutate-outer-value",
416
"try-catch-within-mutable-range",
417
"try-catch",
418
+ "try-catch-with-return",
419
420
// TODO: 🌲
421
"forest-basic",