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 = ;
@@ -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 ;
+}
+
+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 = ;
+ $[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)
{"inputs":[false],"output":{"a":0,"b":"value1","c":true,"wat0":"joe"}}
+{"inputs":[false],"output":{"a":0,"b":"value1","c":true,"wat0":"joe"}}
+{"inputs":[true],"output":{"a":0,"b":"value1","c":true}}
+{"inputs":[true],"output":{"a":0,"b":"value1","c":true}}
+{"inputs":[false],"output":{"a":0,"b":"value1","c":true,"wat0":"joe"}}
+{"inputs":[true],"output":{"a":0,"b":"value1","c":true}}
+{"inputs":[false],"output":{"a":0,"b":"value1","c":true,"wat0":"joe"}}
+{"inputs":[true],"output":{"a":0,"b":"value1","c":true}}
\ 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 ;
+}
+
+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;