From 26442def5b6da8dabe335146ede017acc39cd1b7 Mon Sep 17 00:00:00 2001
From: Joe Savona <joesavona@fb.com>
Date: Thu, 28 Aug 2025 22:26:29 -0700
Subject: [PATCH] [compiler] Fix for scopes with unreachable fallthroughs

Fixes #34108. If a scope ends with with a conditional where some/all branches exit via labeled break, we currently compile in a way that works but bypasses memoization. We end up with a shape like


```js
let t0;
label: {
 if (changed) {
   ...
   if (cond) {
     t0 = ...;
     break label;
   }
   // we don't save the output if the break happens!
   t0 = ...;
   $[0] = t0;
 } else {
   t0 = $[0];
}
```

The fix here is to update AlignReactiveScopesToBlockScopes to take account of breaks that don't go to the natural fallthrough. In this case, we take any active scopes and extend them to start at least as early as the label, and extend at least to the label fallthrough. Thus we produce the correct:

```js
let t0;
if (changed) {
  label: {
    ...
    if (cond) {
      t0 = ...;
      break label;
    }
    t0 = ...;
  }
  // now the break jumps here, and we cache the value
  $[0] = t0;
} else {
  t0 = $[0];
}
```
---
 .../AlignReactiveScopesToBlockScopesHIR.ts    |  35 ++++++
 ...copes-reactive-scope-overlaps-if.expect.md |  16 +--
 .../mutation-within-jsx-and-break.expect.md   |  19 +--
 .../useMemo-multiple-if-else.expect.md        |  37 +++---
 ...seMemo-if-else-both-early-return.expect.md | 118 ++++++++++++++++++
 ...repro-useMemo-if-else-both-early-return.js |  35 ++++++
 .../useMemo-multiple-if-else.expect.md        |  37 +++---
 7 files changed, 241 insertions(+), 56 deletions(-)
 create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro-useMemo-if-else-both-early-return.expect.md
 create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro-useMemo-if-else-both-early-return.js

diff --git a/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/AlignReactiveScopesToBlockScopesHIR.ts b/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/AlignReactiveScopesToBlockScopesHIR.ts
index 2b4e890a40d..e440340bd29 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/AlignReactiveScopesToBlockScopesHIR.ts
+++ b/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/AlignReactiveScopesToBlockScopesHIR.ts
@@ -175,6 +175,41 @@ export function alignReactiveScopesToBlockScopesHIR(fn: HIRFunction): void {
       if (node != null) {
         valueBlockNodes.set(fallthrough, node);
       }
+    } else if (terminal.kind === 'goto') {
+      /**
+       * If we encounter a goto that is not to the natural fallthrough of the current
+       * block (not the topmost fallthrough on the stack), then this is a goto to a
+       * label. Any scopes that extend beyond the goto must be extended to include
+       * the labeled range, so that the break statement doesn't accidentally jump
+       * out of the scope. We do this by extending the start and end of the scope's
+       * range to the label and its fallthrough respectively.
+       */
+      const start = activeBlockFallthroughRanges.find(
+        range => range.fallthrough === terminal.block,
+      );
+      if (start != null && start !== activeBlockFallthroughRanges.at(-1)) {
+        const fallthroughBlock = fn.body.blocks.get(start.fallthrough)!;
+        const firstId =
+          fallthroughBlock.instructions[0]?.id ?? fallthroughBlock.terminal.id;
+        for (const scope of activeScopes) {
+          /**
+           * activeScopes is only filtered at block start points, so some of the
+           * scopes may not actually be active anymore, ie we've past their end
+           * instruction. Only extend ranges for scopes that are actually active.
+           *
+           * TODO: consider pruning activeScopes per instruction
+           */
+          if (scope.range.end <= terminal.id) {
+            continue;
+          }
+          scope.range.start = makeInstructionId(
+            Math.min(start.range.start, scope.range.start),
+          );
+          scope.range.end = makeInstructionId(
+            Math.max(firstId, scope.range.end),
+          );
+        }
+      }
     }
 
     /*
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/align-scopes-reactive-scope-overlaps-if.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/align-scopes-reactive-scope-overlaps-if.expect.md
index 7136b3a173f..03939d16d66 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/align-scopes-reactive-scope-overlaps-if.expect.md
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/align-scopes-reactive-scope-overlaps-if.expect.md
@@ -46,14 +46,16 @@ function useFoo(t0) {
     t1 = $[0];
   }
   let items = t1;
-  bb0: if ($[1] !== cond) {
-    if (cond) {
-      items = [];
-    } else {
-      break bb0;
-    }
+  if ($[1] !== cond) {
+    bb0: {
+      if (cond) {
+        items = [];
+      } else {
+        break bb0;
+      }
 
-    items.push(2);
+      items.push(2);
+    }
     $[1] = cond;
     $[2] = items;
   } else {
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/mutation-within-jsx-and-break.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/mutation-within-jsx-and-break.expect.md
index 1d0f40e29f5..17a8524eeee 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/mutation-within-jsx-and-break.expect.md
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/mutation-within-jsx-and-break.expect.md
@@ -49,12 +49,12 @@ import {
 } from "shared-runtime";
 
 function useFoo(t0) {
-  const $ = _c(3);
+  const $ = _c(4);
   const { data } = t0;
   let obj;
   let myDiv = null;
-  bb0: if (data.cond) {
-    if ($[0] !== data.cond1) {
+  if ($[0] !== data.cond || $[1] !== data.cond1) {
+    bb0: if (data.cond) {
       obj = makeObject_Primitives();
       if (data.cond1) {
         myDiv = <Stringify value={mutateAndReturn(obj)} />;
@@ -62,13 +62,14 @@ function useFoo(t0) {
       }
 
       mutate(obj);
-      $[0] = data.cond1;
-      $[1] = obj;
-      $[2] = myDiv;
-    } else {
-      obj = $[1];
-      myDiv = $[2];
     }
+    $[0] = data.cond;
+    $[1] = data.cond1;
+    $[2] = obj;
+    $[3] = myDiv;
+  } else {
+    obj = $[2];
+    myDiv = $[3];
   }
   return myDiv;
 }
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/propagate-scope-deps-hir-fork/useMemo-multiple-if-else.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/propagate-scope-deps-hir-fork/useMemo-multiple-if-else.expect.md
index a75b592b832..4ce8ef58028 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/propagate-scope-deps-hir-fork/useMemo-multiple-if-else.expect.md
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/propagate-scope-deps-hir-fork/useMemo-multiple-if-else.expect.md
@@ -34,17 +34,16 @@ import { c as _c } from "react/compiler-runtime"; // @enablePropagateDepsInHIR
 import { useMemo } from "react";
 
 function Component(props) {
-  const $ = _c(6);
+  const $ = _c(5);
   let t0;
-  bb0: {
-    let y;
-    if (
-      $[0] !== props.a ||
-      $[1] !== props.b ||
-      $[2] !== props.cond ||
-      $[3] !== props.cond2
-    ) {
-      y = [];
+  if (
+    $[0] !== props.a ||
+    $[1] !== props.b ||
+    $[2] !== props.cond ||
+    $[3] !== props.cond2
+  ) {
+    bb0: {
+      const y = [];
       if (props.cond) {
         y.push(props.a);
       }
@@ -54,17 +53,15 @@ function Component(props) {
       }
 
       y.push(props.b);
-      $[0] = props.a;
-      $[1] = props.b;
-      $[2] = props.cond;
-      $[3] = props.cond2;
-      $[4] = y;
-      $[5] = t0;
-    } else {
-      y = $[4];
-      t0 = $[5];
+      t0 = y;
     }
-    t0 = y;
+    $[0] = props.a;
+    $[1] = props.b;
+    $[2] = props.cond;
+    $[3] = props.cond2;
+    $[4] = t0;
+  } else {
+    t0 = $[4];
   }
   const x = t0;
   return x;
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro-useMemo-if-else-both-early-return.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro-useMemo-if-else-both-early-return.expect.md
new file mode 100644
index 00000000000..acb3b72e3a4
--- /dev/null
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro-useMemo-if-else-both-early-return.expect.md
@@ -0,0 +1,118 @@
+
+## Input
+
+```javascript
+import {useMemo} from 'react';
+import {
+  makeObject_Primitives,
+  mutate,
+  Stringify,
+  ValidateMemoization,
+} from 'shared-runtime';
+
+function Component({cond}) {
+  const memoized = useMemo(() => {
+    const value = makeObject_Primitives();
+    if (cond) {
+      return value;
+    } else {
+      mutate(value);
+      return value;
+    }
+  }, [cond]);
+  return <ValidateMemoization inputs={[cond]} output={memoized} />;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{cond: false}],
+  sequentialRenders: [
+    {cond: false},
+    {cond: false},
+    {cond: true},
+    {cond: true},
+    {cond: false},
+    {cond: true},
+    {cond: false},
+    {cond: true},
+  ],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime";
+import { useMemo } from "react";
+import {
+  makeObject_Primitives,
+  mutate,
+  Stringify,
+  ValidateMemoization,
+} from "shared-runtime";
+
+function Component(t0) {
+  const $ = _c(7);
+  const { cond } = t0;
+  let t1;
+  if ($[0] !== cond) {
+    const value = makeObject_Primitives();
+    if (cond) {
+      t1 = value;
+    } else {
+      mutate(value);
+      t1 = value;
+    }
+    $[0] = cond;
+    $[1] = t1;
+  } else {
+    t1 = $[1];
+  }
+  const memoized = t1;
+  let t2;
+  if ($[2] !== cond) {
+    t2 = [cond];
+    $[2] = cond;
+    $[3] = t2;
+  } else {
+    t2 = $[3];
+  }
+  let t3;
+  if ($[4] !== memoized || $[5] !== t2) {
+    t3 = <ValidateMemoization inputs={t2} output={memoized} />;
+    $[4] = memoized;
+    $[5] = t2;
+    $[6] = t3;
+  } else {
+    t3 = $[6];
+  }
+  return t3;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ cond: false }],
+  sequentialRenders: [
+    { cond: false },
+    { cond: false },
+    { cond: true },
+    { cond: true },
+    { cond: false },
+    { cond: true },
+    { cond: false },
+    { cond: true },
+  ],
+};
+
+```
+      
+### Eval output
+(kind: ok) <div>{"inputs":[false],"output":{"a":0,"b":"value1","c":true,"wat0":"joe"}}</div>
+<div>{"inputs":[false],"output":{"a":0,"b":"value1","c":true,"wat0":"joe"}}</div>
+<div>{"inputs":[true],"output":{"a":0,"b":"value1","c":true}}</div>
+<div>{"inputs":[true],"output":{"a":0,"b":"value1","c":true}}</div>
+<div>{"inputs":[false],"output":{"a":0,"b":"value1","c":true,"wat0":"joe"}}</div>
+<div>{"inputs":[true],"output":{"a":0,"b":"value1","c":true}}</div>
+<div>{"inputs":[false],"output":{"a":0,"b":"value1","c":true,"wat0":"joe"}}</div>
+<div>{"inputs":[true],"output":{"a":0,"b":"value1","c":true}}</div>
\ No newline at end of file
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro-useMemo-if-else-both-early-return.js b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro-useMemo-if-else-both-early-return.js
new file mode 100644
index 00000000000..d3816610771
--- /dev/null
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro-useMemo-if-else-both-early-return.js
@@ -0,0 +1,35 @@
+import {useMemo} from 'react';
+import {
+  makeObject_Primitives,
+  mutate,
+  Stringify,
+  ValidateMemoization,
+} from 'shared-runtime';
+
+function Component({cond}) {
+  const memoized = useMemo(() => {
+    const value = makeObject_Primitives();
+    if (cond) {
+      return value;
+    } else {
+      mutate(value);
+      return value;
+    }
+  }, [cond]);
+  return <ValidateMemoization inputs={[cond]} output={memoized} />;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{cond: false}],
+  sequentialRenders: [
+    {cond: false},
+    {cond: false},
+    {cond: true},
+    {cond: true},
+    {cond: false},
+    {cond: true},
+    {cond: false},
+    {cond: true},
+  ],
+};
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useMemo-multiple-if-else.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useMemo-multiple-if-else.expect.md
index 23b05b54822..7f4b8af9650 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useMemo-multiple-if-else.expect.md
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useMemo-multiple-if-else.expect.md
@@ -33,17 +33,16 @@ import { c as _c } from "react/compiler-runtime";
 import { useMemo } from "react";
 
 function Component(props) {
-  const $ = _c(6);
+  const $ = _c(5);
   let t0;
-  bb0: {
-    let y;
-    if (
-      $[0] !== props.a ||
-      $[1] !== props.b ||
-      $[2] !== props.cond ||
-      $[3] !== props.cond2
-    ) {
-      y = [];
+  if (
+    $[0] !== props.a ||
+    $[1] !== props.b ||
+    $[2] !== props.cond ||
+    $[3] !== props.cond2
+  ) {
+    bb0: {
+      const y = [];
       if (props.cond) {
         y.push(props.a);
       }
@@ -53,17 +52,15 @@ function Component(props) {
       }
 
       y.push(props.b);
-      $[0] = props.a;
-      $[1] = props.b;
-      $[2] = props.cond;
-      $[3] = props.cond2;
-      $[4] = y;
-      $[5] = t0;
-    } else {
-      y = $[4];
-      t0 = $[5];
+      t0 = y;
     }
-    t0 = y;
+    $[0] = props.a;
+    $[1] = props.b;
+    $[2] = props.cond;
+    $[3] = props.cond2;
+    $[4] = t0;
+  } else {
+    t0 = $[4];
   }
   const x = t0;
   return x;