[compiler] Handle more cases with for in try/catch

Author: josephsavona

compiler/packages/babel-plugin-react-compiler/src/HIR/HIR.ts

@@ -612,7 +612,7 @@ export type TryTerminal = {
 export type MaybeThrowTerminal = {
   kind: 'maybe-throw';
   continuation: BlockId;
-  handler: BlockId;
+  handler: BlockId | null;
   id: InstructionId;
   loc: SourceLocation;
   fallthrough?: never;

compiler/packages/babel-plugin-react-compiler/src/HIR/PrintHIR.ts

@@ -291,7 +291,9 @@ export function printTerminal(terminal: Terminal): Array<string> | string {
       break;
     }
     case 'maybe-throw': {
-      value = `[${terminal.id}] MaybeThrow continuation=bb${terminal.continuation} handler=bb${terminal.handler}`;
+      const handlerStr =
+        terminal.handler !== null ? `bb${terminal.handler}` : '(none)';
+      value = `[${terminal.id}] MaybeThrow continuation=bb${terminal.continuation} handler=${handlerStr}`;
       if (terminal.effects != null) {
         value += `\n    ${terminal.effects.map(printAliasingEffect).join('\n    ')}`;
       }

compiler/packages/babel-plugin-react-compiler/src/HIR/visitors.ts

@@ -909,7 +909,8 @@ export function mapTerminalSuccessors(
     }
     case 'maybe-throw': {
       const continuation = fn(terminal.continuation);
-      const handler = fn(terminal.handler);
+      const handler =
+        terminal.handler !== null ? fn(terminal.handler) : null;
       return {
         kind: 'maybe-throw',
         continuation,
@@ -1083,7 +1084,9 @@ export function* eachTerminalSuccessor(terminal: Terminal): Iterable<BlockId> {
     }
     case 'maybe-throw': {
       yield terminal.continuation;
-      yield terminal.handler;
+      if (terminal.handler !== null) {
+        yield terminal.handler;
+      }
       break;
     }
     case 'try': {

compiler/packages/babel-plugin-react-compiler/src/Inference/InferMutationAliasingEffects.ts

@@ -508,7 +508,7 @@ function inferBlock(
   const terminal = block.terminal;
   if (terminal.kind === 'try' && terminal.handlerBinding != null) {
     context.catchHandlers.set(terminal.handler, terminal.handlerBinding);
-  } else if (terminal.kind === 'maybe-throw') {
+  } else if (terminal.kind === 'maybe-throw' && terminal.handler !== null) {
     const handlerParam = context.catchHandlers.get(terminal.handler);
     if (handlerParam != null) {
       CompilerError.invariant(state.kind(handlerParam) != null, {

compiler/packages/babel-plugin-react-compiler/src/Optimization/PruneMaybeThrows.ts

@@ -9,7 +9,6 @@ import {CompilerError} from '..';
 import {
   BlockId,
   GeneratedSource,
-  GotoVariant,
   HIRFunction,
   Instruction,
   assertConsistentIdentifiers,
@@ -25,9 +24,15 @@ import {
 } from '../HIR/HIRBuilder';
 import {printPlace} from '../HIR/PrintHIR';
 
-/*
- * 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
+/**
+ * This pass updates `maybe-throw` terminals for blocks that can provably *never* throw,
+ * nulling out the handler to indicate that control will always continue. Note that
+ * rewriting to a `goto` disrupts the structure of the HIR, making it more difficult to
+ * reconstruct an ast during BuildReactiveFunction. Preserving the maybe-throw makes the
+ * continuations clear, while nulling out the handler tells us that control cannot flow
+ * to the handler.
+ *
+ * For now the analysis 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 {
@@ -82,13 +87,7 @@ function pruneMaybeThrowsImpl(fn: HIRFunction): Map<BlockId, BlockId> | null {
     if (!canThrow) {
       const source = terminalMapping.get(block.id) ?? block.id;
       terminalMapping.set(terminal.continuation, source);
-      block.terminal = {
-        kind: 'goto',
-        block: terminal.continuation,
-        variant: GotoVariant.Break,
-        id: terminal.id,
-        loc: terminal.loc,
-      };
+      terminal.handler = null;
     }
   }
   return terminalMapping.size > 0 ? terminalMapping : null;

compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/BuildReactiveFunction.ts

@@ -168,10 +168,12 @@ class Driver {
       innerValue = innerValue.value;
     }
 
-    // Only add the final instruction if the innermost value is not just a LoadLocal
-    // of the same place we're storing to (which would be a no-op).
-    // This happens when MaybeThrow blocks cause the sequence to already contain
-    // all the necessary instructions.
+    /*
+     * Only add the final instruction if the innermost value is not just a LoadLocal
+     * of the same place we're storing to (which would be a no-op).
+     * This happens when MaybeThrow blocks cause the sequence to already contain
+     * all the necessary instructions.
+     */
     const isLoadOfSamePlace =
       innerValue.kind === 'LoadLocal' &&
       innerValue.place.identifier.id === result.place.identifier.id;
@@ -496,7 +498,10 @@ class Driver {
         const init = this.visitValueBlock(terminal.init, terminal.loc);
         const initValue = this.valueBlockResultToSequence(init, terminal.loc);
 
-        const testValue = this.visitValueBlock(terminal.test, terminal.loc).value;
+        const testValue = this.visitValueBlock(
+          terminal.test,
+          terminal.loc,
+        ).value;
 
         const updateValue =
           terminal.update !== null

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/bug-for-loop-in-try.expect.md

@@ -0,0 +1,95 @@
+
+## Input
+
+```javascript
+function Foo({items}) {
+  const results = [];
+  try {
+    for (let i = 0; i < items.length; i++) {
+      results.push(items[i]);
+    }
+  } catch (e) {
+    return <span>Error</span>;
+  }
+  return <span>{results.join(', ')}</span>;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Foo,
+  params: [{items: ['a', 'b', 'c']}],
+  sequentialRenders: [
+    {items: ['a', 'b', 'c']},
+    {items: ['a', 'b', 'c']},
+    {items: ['x', 'y']},
+    {items: []},
+    {items: ['single']},
+  ],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime";
+function Foo(t0) {
+  const $ = _c(5);
+  const { items } = t0;
+  let results;
+  let t1;
+  if ($[0] !== items) {
+    t1 = Symbol.for("react.early_return_sentinel");
+    bb0: {
+      results = [];
+      try {
+        for (let i = 0; i < items.length; i++) {
+          results.push(items[i]);
+        }
+      } catch (t2) {
+        t1 = <span>Error</span>;
+        break bb0;
+      }
+    }
+    $[0] = items;
+    $[1] = results;
+    $[2] = t1;
+  } else {
+    results = $[1];
+    t1 = $[2];
+  }
+  if (t1 !== Symbol.for("react.early_return_sentinel")) {
+    return t1;
+  }
+
+  const t2 = results.join(", ");
+  let t3;
+  if ($[3] !== t2) {
+    t3 = <span>{t2}</span>;
+    $[3] = t2;
+    $[4] = t3;
+  } else {
+    t3 = $[4];
+  }
+  return t3;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Foo,
+  params: [{ items: ["a", "b", "c"] }],
+  sequentialRenders: [
+    { items: ["a", "b", "c"] },
+    { items: ["a", "b", "c"] },
+    { items: ["x", "y"] },
+    { items: [] },
+    { items: ["single"] },
+  ],
+};
+
+```
+      
+### Eval output
+(kind: ok) <span>a, b, c</span>
+<span>a, b, c</span>
+<span>x, y</span>
+<span></span>
+<span>single</span>

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/bug-for-loop-in-try.js

@@ -0,0 +1,23 @@
+function Foo({items}) {
+  const results = [];
+  try {
+    for (let i = 0; i < items.length; i++) {
+      results.push(items[i]);
+    }
+  } catch (e) {
+    return <span>Error</span>;
+  }
+  return <span>{results.join(', ')}</span>;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Foo,
+  params: [{items: ['a', 'b', 'c']}],
+  sequentialRenders: [
+    {items: ['a', 'b', 'c']},
+    {items: ['a', 'b', 'c']},
+    {items: ['x', 'y']},
+    {items: []},
+    {items: ['single']},
+  ],
+};

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/bug-invariant-expected-break-target.expect.md

@@ -0,0 +1,122 @@
+
+## Input
+
+```javascript
+import {useMemo} from 'react';
+
+function useFoo(text) {
+  return useMemo(() => {
+    try {
+      let formattedText = '';
+      try {
+        formattedText = format(text);
+      } catch {
+        formattedText = text;
+      }
+      return formattedText || '';
+    } catch (e) {
+      return '';
+    }
+  }, [text]);
+}
+
+function format(text) {
+  return text.toUpperCase();
+}
+
+function Foo({text}) {
+  const result = useFoo(text);
+  return <span>{result}</span>;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Foo,
+  params: [{text: 'hello'}],
+  sequentialRenders: [
+    {text: 'hello'},
+    {text: 'hello'},
+    {text: 'world'},
+    {text: ''},
+  ],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime";
+import { useMemo } from "react";
+
+function useFoo(text) {
+  const $ = _c(2);
+  let t0;
+  try {
+    let formattedText;
+    try {
+      let t2;
+      if ($[0] !== text) {
+        t2 = format(text);
+        $[0] = text;
+        $[1] = t2;
+      } else {
+        t2 = $[1];
+      }
+      formattedText = t2;
+    } catch {
+      formattedText = text;
+    }
+
+    t0 = formattedText || "";
+  } catch (t1) {
+    t0 = "";
+  }
+  return t0;
+}
+
+function format(text) {
+  const $ = _c(2);
+  let t0;
+  if ($[0] !== text) {
+    t0 = text.toUpperCase();
+    $[0] = text;
+    $[1] = t0;
+  } else {
+    t0 = $[1];
+  }
+  return t0;
+}
+
+function Foo(t0) {
+  const $ = _c(2);
+  const { text } = t0;
+  const result = useFoo(text);
+  let t1;
+  if ($[0] !== result) {
+    t1 = <span>{result}</span>;
+    $[0] = result;
+    $[1] = t1;
+  } else {
+    t1 = $[1];
+  }
+  return t1;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Foo,
+  params: [{ text: "hello" }],
+  sequentialRenders: [
+    { text: "hello" },
+    { text: "hello" },
+    { text: "world" },
+    { text: "" },
+  ],
+};
+
+```
+      
+### Eval output
+(kind: ok) <span>HELLO</span>
+<span>HELLO</span>
+<span>WORLD</span>
+<span></span>