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.

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 });

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,
         };

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 = {

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;
+    }
+  }
+}

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";

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<BlockId> = new Set();
 
+  #catchHandlers: Set<BlockId> = 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);
   }
 
   /**

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;
+}
+
+```
+      

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;
+}

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",