[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];
}
```

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),
+          );
+        }
+      }
     }
 
     /*

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 {

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

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;

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>

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

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;