[compiler][wip] Support finally clauses

First pass at adding support for `finally` clauses. I took a simple approach where we reify the JS semantics around return shadowing in the HIR and output. Ie, we create a temporary variable that will be the return value, and change returns in the try/catch to assign to that value and then break to the finalizer. In the finalizer returns work as normal, but we put a final extra "if temporary set, return it" statement. It would be more ideal to fully restore the return shadowing, but given that finally clauses are relatively rare this seems like a reasonable compromise.

I need to do more analysis to make sure the approach is correct.

Author: josephsavona

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

@@ -1329,82 +1329,102 @@ function lowerStatement(
       const continuationBlock = builder.reserve('block');
 
       const handlerPath = stmt.get('handler');
-      if (!hasNode(handlerPath)) {
+      const finalizerPath = stmt.get('finalizer');
+
+      // Must have at least a catch or finally clause
+      if (!hasNode(handlerPath) && !hasNode(finalizerPath)) {
         builder.errors.push({
-          reason: `(BuildHIR::lowerStatement) Handle TryStatement without a catch clause`,
+          reason: `(BuildHIR::lowerStatement) Handle TryStatement without a catch or finally clause`,
           category: ErrorCategory.Todo,
           loc: stmt.node.loc ?? null,
           suggestions: null,
         });
         return;
       }
-      if (hasNode(stmt.get('finalizer'))) {
-        builder.errors.push({
-          reason: `(BuildHIR::lowerStatement) Handle TryStatement with a finalizer ('finally') clause`,
-          category: ErrorCategory.Todo,
-          loc: stmt.node.loc ?? null,
-          suggestions: null,
-        });
-      }
 
-      const handlerBindingPath = handlerPath.get('param');
+      // Build handler (catch) block if present
+      let handler: BlockId | null = null;
       let handlerBinding: {
         place: Place;
         path: NodePath<t.Identifier | t.ArrayPattern | t.ObjectPattern>;
       } | null = null;
-      if (hasNode(handlerBindingPath)) {
-        const place: Place = {
-          kind: 'Identifier',
-          identifier: builder.makeTemporary(
-            handlerBindingPath.node.loc ?? GeneratedSource,
-          ),
-          effect: Effect.Unknown,
-          reactive: false,
-          loc: handlerBindingPath.node.loc ?? GeneratedSource,
-        };
-        promoteTemporary(place.identifier);
-        lowerValueToTemporary(builder, {
-          kind: 'DeclareLocal',
-          lvalue: {
-            kind: InstructionKind.Catch,
-            place: {...place},
-          },
-          type: null,
-          loc: handlerBindingPath.node.loc ?? GeneratedSource,
-        });
 
-        handlerBinding = {
-          path: handlerBindingPath,
-          place,
-        };
-      }
+      if (hasNode(handlerPath)) {
+        const handlerBindingPath = handlerPath.get('param');
+        if (hasNode(handlerBindingPath)) {
+          const place: Place = {
+            kind: 'Identifier',
+            identifier: builder.makeTemporary(
+              handlerBindingPath.node.loc ?? GeneratedSource,
+            ),
+            effect: Effect.Unknown,
+            reactive: false,
+            loc: handlerBindingPath.node.loc ?? GeneratedSource,
+          };
+          promoteTemporary(place.identifier);
+          lowerValueToTemporary(builder, {
+            kind: 'DeclareLocal',
+            lvalue: {
+              kind: InstructionKind.Catch,
+              place: {...place},
+            },
+            type: null,
+            loc: handlerBindingPath.node.loc ?? GeneratedSource,
+          });
 
-      const handler = builder.enter('catch', _blockId => {
-        if (handlerBinding !== null) {
-          lowerAssignment(
-            builder,
-            handlerBinding.path.node.loc ?? GeneratedSource,
-            InstructionKind.Catch,
-            handlerBinding.path,
-            {...handlerBinding.place},
-            'Assignment',
-          );
+          handlerBinding = {
+            path: handlerBindingPath,
+            place,
+          };
         }
-        lowerStatement(builder, handlerPath.get('body'));
-        return {
-          kind: 'goto',
-          block: continuationBlock.id,
-          variant: GotoVariant.Break,
-          id: makeInstructionId(0),
-          loc: handlerPath.node.loc ?? GeneratedSource,
-        };
-      });
+
+        handler = builder.enter('catch', _blockId => {
+          if (handlerBinding !== null) {
+            lowerAssignment(
+              builder,
+              handlerBinding.path.node.loc ?? GeneratedSource,
+              InstructionKind.Catch,
+              handlerBinding.path,
+              {...handlerBinding.place},
+              'Assignment',
+            );
+          }
+          lowerStatement(builder, handlerPath.get('body'));
+          return {
+            kind: 'goto',
+            block: continuationBlock.id,
+            variant: GotoVariant.Break,
+            id: makeInstructionId(0),
+            loc: handlerPath.node.loc ?? GeneratedSource,
+          };
+        });
+      }
+
+      // Build finalizer (finally) block if present
+      let finalizer: BlockId | null = null;
+      if (hasNode(finalizerPath)) {
+        finalizer = builder.enter('block', _blockId => {
+          lowerStatement(builder, finalizerPath);
+          return {
+            kind: 'goto',
+            block: continuationBlock.id,
+            variant: GotoVariant.Break,
+            id: makeInstructionId(0),
+            loc: finalizerPath.node.loc ?? GeneratedSource,
+          };
+        });
+      }
 
       const block = builder.enter('block', _blockId => {
         const block = stmt.get('block');
-        builder.enterTryCatch(handler, () => {
+        // If there's a handler, exceptions go there; otherwise they propagate normally
+        if (handler !== null) {
+          builder.enterTryCatch(handler, () => {
+            lowerStatement(builder, block);
+          });
+        } else {
           lowerStatement(builder, block);
-        });
+        }
         return {
           kind: 'goto',
           block: continuationBlock.id,
@@ -1421,6 +1441,7 @@ function lowerStatement(
           handlerBinding:
             handlerBinding !== null ? {...handlerBinding.place} : null,
           handler,
+          finalizer,
           fallthrough: continuationBlock.id,
           id: makeInstructionId(0),
           loc: stmt.node.loc ?? GeneratedSource,

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

@@ -273,7 +273,8 @@ export type ReactiveTryTerminal = {
   kind: 'try';
   block: ReactiveBlock;
   handlerBinding: Place | null;
-  handler: ReactiveBlock;
+  handler: ReactiveBlock | null;
+  finalizer: ReactiveBlock | null;
   id: InstructionId;
   loc: SourceLocation;
 };
@@ -602,8 +603,8 @@ export type TryTerminal = {
   kind: 'try';
   block: BlockId;
   handlerBinding: Place | null;
-  handler: BlockId;
-  // TODO: support `finally`
+  handler: BlockId | null;
+  finalizer: BlockId | null;
   fallthrough: BlockId;
   id: InstructionId;
   loc: SourceLocation;

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

@@ -954,6 +954,7 @@ export function removeUnnecessaryTryCatch(fn: HIR): void {
   for (const [, block] of fn.blocks) {
     if (
       block.terminal.kind === 'try' &&
+      block.terminal.handler !== null &&
       !fn.blocks.has(block.terminal.handler)
     ) {
       const handlerId = block.terminal.handler;

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

@@ -310,12 +310,14 @@ export function printTerminal(terminal: Terminal): Array<string> | string {
       break;
     }
     case 'try': {
-      value = `[${terminal.id}] Try block=bb${terminal.block} handler=bb${
-        terminal.handler
+      value = `[${terminal.id}] Try block=bb${terminal.block}${
+        terminal.handler !== null ? ` handler=bb${terminal.handler}` : ''
       }${
         terminal.handlerBinding !== null
           ? ` handlerBinding=(${printPlace(terminal.handlerBinding)})`
           : ''
+      }${
+        terminal.finalizer !== null ? ` finalizer=bb${terminal.finalizer}` : ''
       } fallthrough=${
         terminal.fallthrough != null ? `bb${terminal.fallthrough}` : ''
       }`;

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

@@ -921,13 +921,17 @@ export function mapTerminalSuccessors(
     }
     case 'try': {
       const block = fn(terminal.block);
-      const handler = fn(terminal.handler);
+      const handler =
+        terminal.handler !== null ? fn(terminal.handler) : null;
+      const finalizer =
+        terminal.finalizer !== null ? fn(terminal.finalizer) : null;
       const fallthrough = fn(terminal.fallthrough);
       return {
         kind: 'try',
         block,
         handlerBinding: terminal.handlerBinding,
         handler,
+        finalizer,
         fallthrough,
         id: makeInstructionId(0),
         loc: terminal.loc,
@@ -1088,6 +1092,12 @@ export function* eachTerminalSuccessor(terminal: Terminal): Iterable<BlockId> {
     }
     case 'try': {
       yield terminal.block;
+      if (terminal.handler !== null) {
+        yield terminal.handler;
+      }
+      if (terminal.finalizer !== null) {
+        yield terminal.finalizer;
+      }
       break;
     }
     case 'scope':

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

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

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

@@ -844,14 +844,21 @@ class Driver {
           const scheduleId = this.cx.schedule(fallthroughId, 'if');
           scheduleIds.push(scheduleId);
         }
-        this.cx.scheduleCatchHandler(terminal.handler);
+        if (terminal.handler !== null) {
+          this.cx.scheduleCatchHandler(terminal.handler);
+        }
 
         const block = this.traverseBlock(
           this.cx.ir.blocks.get(terminal.block)!,
         );
-        const handler = this.traverseBlock(
-          this.cx.ir.blocks.get(terminal.handler)!,
-        );
+        const handler =
+          terminal.handler !== null
+            ? this.traverseBlock(this.cx.ir.blocks.get(terminal.handler)!)
+            : null;
+        const finalizer =
+          terminal.finalizer !== null
+            ? this.traverseBlock(this.cx.ir.blocks.get(terminal.finalizer)!)
+            : null;
 
         this.cx.unscheduleAll(scheduleIds);
         blockValue.push({
@@ -864,6 +871,7 @@ class Driver {
             block,
             handlerBinding: terminal.handlerBinding,
             handler,
+            finalizer,
             id: terminal.id,
           },
         });

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

@@ -1307,15 +1307,26 @@ function codegenTerminal(
       return codegenBlock(cx, terminal.block);
     }
     case 'try': {
-      let catchParam = null;
-      if (terminal.handlerBinding !== null) {
-        catchParam = convertIdentifier(terminal.handlerBinding.identifier);
-        cx.temp.set(terminal.handlerBinding.identifier.declarationId, null);
+      let catchClause: t.CatchClause | null = null;
+      if (terminal.handler !== null) {
+        let catchParam = null;
+        if (terminal.handlerBinding !== null) {
+          catchParam = convertIdentifier(terminal.handlerBinding.identifier);
+          cx.temp.set(terminal.handlerBinding.identifier.declarationId, null);
+        }
+        catchClause = t.catchClause(catchParam, codegenBlock(cx, terminal.handler));
       }
+
+      const finalizer =
+        terminal.finalizer !== null
+          ? codegenBlock(cx, terminal.finalizer)
+          : null;
+
       return createTryStatement(
         terminal.loc,
         codegenBlock(cx, terminal.block),
-        t.catchClause(catchParam, codegenBlock(cx, terminal.handler)),
+        catchClause,
+        finalizer,
       );
     }
     default: {

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

@@ -395,14 +395,24 @@ function writeTerminal(writer: Writer, terminal: ReactiveTerminal): void {
     case 'try': {
       writer.writeLine(`[${terminal.id}] try {`);
       writeReactiveInstructions(writer, terminal.block);
-      writer.write(`} catch `);
-      if (terminal.handlerBinding !== null) {
-        writer.writeLine(`(${printPlace(terminal.handlerBinding)}) {`);
-      } else {
-        writer.writeLine(`{`);
+      if (terminal.handler !== null) {
+        writer.write(`} catch `);
+        if (terminal.handlerBinding !== null) {
+          writer.writeLine(`(${printPlace(terminal.handlerBinding)}) {`);
+        } else {
+          writer.writeLine(`{`);
+        }
+        writeReactiveInstructions(writer, terminal.handler);
+        writer.writeLine('}');
+      }
+      if (terminal.finalizer !== null) {
+        writer.writeLine(`} finally {`);
+        writeReactiveInstructions(writer, terminal.finalizer);
+        writer.writeLine('}');
+      }
+      if (terminal.handler === null && terminal.finalizer === null) {
+        writer.writeLine('}');
       }
-      writeReactiveInstructions(writer, terminal.handler);
-      writer.writeLine('}');
       break;
     }
     default:

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

@@ -169,7 +169,12 @@ export class ReactiveFunctionVisitor<TState = void> {
       }
       case 'try': {
         this.visitBlock(terminal.block, state);
-        this.visitBlock(terminal.handler, state);
+        if (terminal.handler !== null) {
+          this.visitBlock(terminal.handler, state);
+        }
+        if (terminal.finalizer !== null) {
+          this.visitBlock(terminal.finalizer, state);
+        }
         break;
       }
       default: {
@@ -559,7 +564,12 @@ export class ReactiveFunctionTransform<
         if (terminal.handlerBinding !== null) {
           this.visitPlace(terminal.id, terminal.handlerBinding, state);
         }
-        this.visitBlock(terminal.handler, state);
+        if (terminal.handler !== null) {
+          this.visitBlock(terminal.handler, state);
+        }
+        if (terminal.finalizer !== null) {
+          this.visitBlock(terminal.finalizer, state);
+        }
         break;
       }
       default: {
@@ -653,7 +663,12 @@ export function mapTerminalBlocks(
     }
     case 'try': {
       terminal.block = fn(terminal.block);
-      terminal.handler = fn(terminal.handler);
+      if (terminal.handler !== null) {
+        terminal.handler = fn(terminal.handler);
+      }
+      if (terminal.finalizer !== null) {
+        terminal.finalizer = fn(terminal.finalizer);
+      }
       break;
     }
     default: {