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/ReactiveScopes/PruneNonEscapingScopes.ts b/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PruneNonEscapingScopes.ts
index 3afb00b71a8..1dcaf0b798e 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PruneNonEscapingScopes.ts
+++ b/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PruneNonEscapingScopes.ts
@@ -411,7 +411,9 @@ class CollectDependenciesVisitor extends ReactiveFunctionVisitor<
this.state = state;
this.options = {
memoizeJsxElements: !this.env.config.enableForest,
- forceMemoizePrimitives: this.env.config.enableForest,
+ forceMemoizePrimitives:
+ this.env.config.enableForest ||
+ this.env.config.enablePreserveExistingMemoizationGuarantees,
};
}
@@ -534,9 +536,23 @@ class CollectDependenciesVisitor extends ReactiveFunctionVisitor<
case 'JSXText':
case 'BinaryExpression':
case 'UnaryExpression': {
- const level = options.forceMemoizePrimitives
- ? MemoizationLevel.Memoized
- : MemoizationLevel.Never;
+ if (options.forceMemoizePrimitives) {
+ /**
+ * Because these instructions produce primitives we usually don't consider
+ * them as escape points: they are known to copy, not return references.
+ * However if we're forcing memoization of primitives then we mark these
+ * instructions as needing memoization and walk their rvalues to ensure
+ * any scopes transitively reachable from the rvalues are considered for
+ * memoization. Note: we may still prune primitive-producing scopes if
+ * they don't ultimately escape at all.
+ */
+ const level = MemoizationLevel.Memoized;
+ return {
+ lvalues: lvalue !== null ? [{place: lvalue, level}] : [],
+ rvalues: [...eachReactiveValueOperand(value)],
+ };
+ }
+ const level = MemoizationLevel.Never;
return {
// All of these instructions return a primitive value and never need to be memoized
lvalues: lvalue !== null ? [{place: lvalue, level}] : [],
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/dont-memoize-primitive-function-call-non-escaping-useMemo.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/dont-memoize-primitive-function-call-non-escaping-useMemo.expect.md
new file mode 100644
index 00000000000..93b08128a0a
--- /dev/null
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/dont-memoize-primitive-function-call-non-escaping-useMemo.expect.md
@@ -0,0 +1,77 @@
+
+## Input
+
+```javascript
+// @compilationMode:"infer" @enablePreserveExistingMemoizationGuarantees @validatePreserveExistingMemoizationGuarantees
+import {useMemo} from 'react';
+import {makeObject_Primitives, ValidateMemoization} from 'shared-runtime';
+
+function Component(props) {
+ const result = useMemo(
+ () => makeObject(props.value).value + 1,
+ [props.value]
+ );
+ console.log(result);
+ return 'ok';
+}
+
+function makeObject(value) {
+ console.log(value);
+ return {value};
+}
+
+export const TODO_FIXTURE_ENTRYPOINT = {
+ fn: Component,
+ params: [{value: 42}],
+ sequentialRenders: [
+ {value: 42},
+ {value: 42},
+ {value: 3.14},
+ {value: 3.14},
+ {value: 42},
+ {value: 3.14},
+ {value: 42},
+ {value: 3.14},
+ ],
+};
+
+```
+
+## Code
+
+```javascript
+// @compilationMode:"infer" @enablePreserveExistingMemoizationGuarantees @validatePreserveExistingMemoizationGuarantees
+import { useMemo } from "react";
+import { makeObject_Primitives, ValidateMemoization } from "shared-runtime";
+
+function Component(props) {
+ const result = makeObject(props.value).value + 1;
+
+ console.log(result);
+ return "ok";
+}
+
+function makeObject(value) {
+ console.log(value);
+ return { value };
+}
+
+export const TODO_FIXTURE_ENTRYPOINT = {
+ fn: Component,
+ params: [{ value: 42 }],
+ sequentialRenders: [
+ { value: 42 },
+ { value: 42 },
+ { value: 3.14 },
+ { value: 3.14 },
+ { value: 42 },
+ { value: 3.14 },
+ { value: 42 },
+ { value: 3.14 },
+ ],
+};
+
+```
+
+### Eval output
+(kind: exception) Fixture not implemented
\ No newline at end of file
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/dont-memoize-primitive-function-call-non-escaping-useMemo.js b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/dont-memoize-primitive-function-call-non-escaping-useMemo.js
new file mode 100644
index 00000000000..2ee24917c53
--- /dev/null
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/dont-memoize-primitive-function-call-non-escaping-useMemo.js
@@ -0,0 +1,32 @@
+// @compilationMode:"infer" @enablePreserveExistingMemoizationGuarantees @validatePreserveExistingMemoizationGuarantees
+import {useMemo} from 'react';
+import {makeObject_Primitives, ValidateMemoization} from 'shared-runtime';
+
+function Component(props) {
+ const result = useMemo(
+ () => makeObject(props.value).value + 1,
+ [props.value]
+ );
+ console.log(result);
+ return 'ok';
+}
+
+function makeObject(value) {
+ console.log(value);
+ return {value};
+}
+
+export const TODO_FIXTURE_ENTRYPOINT = {
+ fn: Component,
+ params: [{value: 42}],
+ sequentialRenders: [
+ {value: 42},
+ {value: 42},
+ {value: 3.14},
+ {value: 3.14},
+ {value: 42},
+ {value: 3.14},
+ {value: 42},
+ {value: 3.14},
+ ],
+};
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/dont-memoize-primitive-function-call-non-escaping.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/dont-memoize-primitive-function-call-non-escaping.expect.md
new file mode 100644
index 00000000000..e2f6c9e6c2c
--- /dev/null
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/dont-memoize-primitive-function-call-non-escaping.expect.md
@@ -0,0 +1,81 @@
+
+## Input
+
+```javascript
+// @compilationMode:"infer" @enablePreserveExistingMemoizationGuarantees @validatePreserveExistingMemoizationGuarantees
+import {useMemo} from 'react';
+import {makeObject_Primitives, ValidateMemoization} from 'shared-runtime';
+
+function Component(props) {
+ const result = makeObject(props.value).value + 1;
+ console.log(result);
+ return 'ok';
+}
+
+function makeObject(value) {
+ console.log(value);
+ return {value};
+}
+
+export const FIXTURE_ENTRYPOINT = {
+ fn: Component,
+ params: [{value: 42}],
+ sequentialRenders: [
+ {value: 42},
+ {value: 42},
+ {value: 3.14},
+ {value: 3.14},
+ {value: 42},
+ {value: 3.14},
+ {value: 42},
+ {value: 3.14},
+ ],
+};
+
+```
+
+## Code
+
+```javascript
+// @compilationMode:"infer" @enablePreserveExistingMemoizationGuarantees @validatePreserveExistingMemoizationGuarantees
+import { useMemo } from "react";
+import { makeObject_Primitives, ValidateMemoization } from "shared-runtime";
+
+function Component(props) {
+ const result = makeObject(props.value).value + 1;
+ console.log(result);
+ return "ok";
+}
+
+function makeObject(value) {
+ console.log(value);
+ return { value };
+}
+
+export const FIXTURE_ENTRYPOINT = {
+ fn: Component,
+ params: [{ value: 42 }],
+ sequentialRenders: [
+ { value: 42 },
+ { value: 42 },
+ { value: 3.14 },
+ { value: 3.14 },
+ { value: 42 },
+ { value: 3.14 },
+ { value: 42 },
+ { value: 3.14 },
+ ],
+};
+
+```
+
+### Eval output
+(kind: ok) "ok"
+"ok"
+"ok"
+"ok"
+"ok"
+"ok"
+"ok"
+"ok"
+logs: [42,43,42,43,3.14,4.140000000000001,3.14,4.140000000000001,42,43,3.14,4.140000000000001,42,43,3.14,4.140000000000001]
\ No newline at end of file
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/dont-memoize-primitive-function-call-non-escaping.js b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/dont-memoize-primitive-function-call-non-escaping.js
new file mode 100644
index 00000000000..b4d8d344441
--- /dev/null
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/dont-memoize-primitive-function-call-non-escaping.js
@@ -0,0 +1,29 @@
+// @compilationMode:"infer" @enablePreserveExistingMemoizationGuarantees @validatePreserveExistingMemoizationGuarantees
+import {useMemo} from 'react';
+import {makeObject_Primitives, ValidateMemoization} from 'shared-runtime';
+
+function Component(props) {
+ const result = makeObject(props.value).value + 1;
+ console.log(result);
+ return 'ok';
+}
+
+function makeObject(value) {
+ console.log(value);
+ return {value};
+}
+
+export const FIXTURE_ENTRYPOINT = {
+ fn: Component,
+ params: [{value: 42}],
+ sequentialRenders: [
+ {value: 42},
+ {value: 42},
+ {value: 3.14},
+ {value: 3.14},
+ {value: 42},
+ {value: 3.14},
+ {value: 42},
+ {value: 3.14},
+ ],
+};
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/memoize-primitive-function-calls.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/memoize-primitive-function-calls.expect.md
new file mode 100644
index 00000000000..70e70a26b46
--- /dev/null
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/memoize-primitive-function-calls.expect.md
@@ -0,0 +1,107 @@
+
+## Input
+
+```javascript
+// @compilationMode:"infer" @enablePreserveExistingMemoizationGuarantees @validatePreserveExistingMemoizationGuarantees
+import {useMemo} from 'react';
+import {makeObject_Primitives, ValidateMemoization} from 'shared-runtime';
+
+function Component(props) {
+ const result = useMemo(() => {
+ return makeObject(props.value).value + 1;
+ }, [props.value]);
+ return