From cf133ba05f1a6d3b38bc699cb3be61510c6916de Mon Sep 17 00:00:00 2001
From: Joe Savona <joesavona@fb.com>
Date: Mon, 15 Jul 2024 17:32:04 +0900
Subject: [PATCH] [compiler] Handle earlier creation of early return on scopes

I'm experimenting with a new pass that sometimes creates scopes with early returns earlier in the pipeline, but there are a few passes that assume that can't happen. This PR is updating those passes just to be more resilient to help unblock experimentation.

ghstack-source-id: a9e348181ddad1a1e936ef023b5d5ee44aaf3d8c
Pull Request resolved: https://github.com/facebook/react/pull/30333
---
 .../src/ReactiveScopes/PropagateEarlyReturns.ts        | 10 ++++++++--
 .../src/ReactiveScopes/PruneNonEscapingScopes.ts       |  8 ++++++--
 2 files changed, 14 insertions(+), 4 deletions(-)

diff --git a/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PropagateEarlyReturns.ts b/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PropagateEarlyReturns.ts
index ef2c217e2557d..3012314e9e5aa 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PropagateEarlyReturns.ts
+++ b/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PropagateEarlyReturns.ts
@@ -24,8 +24,6 @@ import { EARLY_RETURN_SENTINEL } from "./CodegenReactiveFunction";
 import { ReactiveFunctionTransform, Transformed } from "./visitors";
 
 /**
- * TODO: Actualy propagate early return information, for now we throw a Todo bailout.
- *
  * This pass ensures that reactive blocks honor the control flow behavior of the
  * original code including early return semantics. Specifically, if a reactive
  * scope early returned during the previous execution and the inputs to that block
@@ -135,6 +133,14 @@ class Transform extends ReactiveFunctionTransform<State> {
     scopeBlock: ReactiveScopeBlock,
     parentState: State
   ): void {
+    /**
+     * Exit early if an earlier pass has already created an early return,
+     * which may happen in alternate compiler configurations.
+     */
+    if (scopeBlock.scope.earlyReturnValue !== null) {
+      return;
+    }
+
     const innerState: State = {
       withinReactiveScope: true,
       earlyReturnValue: parentState.earlyReturnValue,
diff --git a/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PruneNonEscapingScopes.ts b/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PruneNonEscapingScopes.ts
index d2b651f8b84ae..d93a8294d17b0 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PruneNonEscapingScopes.ts
+++ b/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PruneNonEscapingScopes.ts
@@ -932,10 +932,14 @@ class PruneScopesTransform extends ReactiveFunctionTransform<
      * is early-returned from within the scope. For now we intentionaly keep
      * these scopes, and let them get pruned later by PruneUnusedScopes
      * _after_ handling the early-return case in PropagateEarlyReturns.
+     *
+     * Also keep the scope if an early return was created by some earlier pass,
+     * which may happen in alternate compiler configurations.
      */
     if (
-      scopeBlock.scope.declarations.size === 0 &&
-      scopeBlock.scope.reassignments.size === 0
+      (scopeBlock.scope.declarations.size === 0 &&
+        scopeBlock.scope.reassignments.size === 0) ||
+      scopeBlock.scope.earlyReturnValue !== null
     ) {
       return { kind: "keep" };
     }