From 3a2ff8b51b5ab54c22f55a5f826c53419f718887 Mon Sep 17 00:00:00 2001
From: Joseph Savona <6425824+josephsavona@users.noreply.github.com>
Date: Wed, 18 Jun 2025 16:00:36 -0700
Subject: [PATCH] [compiler] Fix <ValidateMemoization> (#33547)

By accident we were only ever checking the compiled output, but the
intention was in general to be able to compare memoization with/without
forget.

---
[//]: # (BEGIN SAPLING FOOTER)
Stack created with [Sapling](https://sapling-scm.com). Best reviewed
with [ReviewStack](https://reviewstack.dev/facebook/react/pull/33547).
* #33571
* #33558
* __->__ #33547
---
 ...-control-flow-sensitive-mutation.expect.md | 13 ++----
 .../todo-control-flow-sensitive-mutation.tsx  |  9 +---
 ...tivity-createfrom-capture-lambda.expect.md |  4 +-
 ...transitivity-createfrom-capture-lambda.tsx |  2 +-
 ...ity-add-captured-array-to-itself.expect.md |  8 ++--
 ...nsitivity-add-captured-array-to-itself.tsx |  4 +-
 ...tivity-capture-createfrom-lambda.expect.md |  4 +-
 ...transitivity-capture-createfrom-lambda.tsx |  2 +-
 .../transitivity-capture-createfrom.expect.md |  4 +-
 .../transitivity-capture-createfrom.tsx       |  2 +-
 .../transitivity-createfrom-capture.expect.md |  4 +-
 .../transitivity-createfrom-capture.tsx       |  2 +-
 ...ansitivity-phi-assign-or-capture.expect.md |  4 +-
 .../transitivity-phi-assign-or-capture.tsx    |  2 +-
 ...d-identity-function-frozen-input.expect.md | 42 ++++++++++---------
 .../typed-identity-function-frozen-input.js   |  3 +-
 .../compiler/weakmap-constructor.expect.md    | 36 +++++++++++++---
 .../fixtures/compiler/weakmap-constructor.js  | 19 +++++++--
 .../compiler/weakset-constructor.expect.md    | 36 +++++++++++++---
 .../fixtures/compiler/weakset-constructor.js  | 19 +++++++--
 .../snap/src/sprout/shared-runtime.ts         |  4 +-
 21 files changed, 145 insertions(+), 78 deletions(-)

diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/todo-control-flow-sensitive-mutation.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/todo-control-flow-sensitive-mutation.expect.md
index d3555c0b27..a3cf9f638a 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/todo-control-flow-sensitive-mutation.expect.md
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/todo-control-flow-sensitive-mutation.expect.md
@@ -23,14 +23,9 @@ function Component({a, b, c}: {a: number; b: number; c: number}) {
 
   return (
     <>
-      <ValidateMemoization inputs={[a, b, c]} output={x} alwaysCheck={true} />;
+      <ValidateMemoization inputs={[a, b, c]} output={x} />;
       {/* TODO: should only depend on c */}
-      <ValidateMemoization
-        inputs={[a, b, c]}
-        output={x[0]}
-        alwaysCheck={true}
-      />
-      ;
+      <ValidateMemoization inputs={[a, b, c]} output={x[0]} />;
     </>
   );
 }
@@ -98,7 +93,7 @@ function Component(t0) {
   }
   let t3;
   if ($[9] !== t2 || $[10] !== x) {
-    t3 = <ValidateMemoization inputs={t2} output={x} alwaysCheck={true} />;
+    t3 = <ValidateMemoization inputs={t2} output={x} />;
     $[9] = t2;
     $[10] = x;
     $[11] = t3;
@@ -117,7 +112,7 @@ function Component(t0) {
   }
   let t5;
   if ($[16] !== t4 || $[17] !== x[0]) {
-    t5 = <ValidateMemoization inputs={t4} output={x[0]} alwaysCheck={true} />;
+    t5 = <ValidateMemoization inputs={t4} output={x[0]} />;
     $[16] = t4;
     $[17] = x[0];
     $[18] = t5;
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/todo-control-flow-sensitive-mutation.tsx b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/todo-control-flow-sensitive-mutation.tsx
index 80386c46af..61f8c47e45 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/todo-control-flow-sensitive-mutation.tsx
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/todo-control-flow-sensitive-mutation.tsx
@@ -19,14 +19,9 @@ function Component({a, b, c}: {a: number; b: number; c: number}) {
 
   return (
     <>
-      <ValidateMemoization inputs={[a, b, c]} output={x} alwaysCheck={true} />;
+      <ValidateMemoization inputs={[a, b, c]} output={x} />;
       {/* TODO: should only depend on c */}
-      <ValidateMemoization
-        inputs={[a, b, c]}
-        output={x[0]}
-        alwaysCheck={true}
-      />
-      ;
+      <ValidateMemoization inputs={[a, b, c]} output={x[0]} />;
     </>
   );
 }
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/todo-transitivity-createfrom-capture-lambda.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/todo-transitivity-createfrom-capture-lambda.expect.md
index 9210bef3fb..9dba055973 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/todo-transitivity-createfrom-capture-lambda.expect.md
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/todo-transitivity-createfrom-capture-lambda.expect.md
@@ -22,7 +22,7 @@ function Component({a, b}) {
   typedMutate(z, b);
 
   // TODO: this *should* only depend on `a`
-  return <ValidateMemoization inputs={[a, b]} output={x} alwaysCheck={true} />;
+  return <ValidateMemoization inputs={[a, b]} output={x} />;
 }
 
 export const FIXTURE_ENTRYPOINT = {
@@ -86,7 +86,7 @@ function Component(t0) {
   }
   let t3;
   if ($[7] !== t2 || $[8] !== x) {
-    t3 = <ValidateMemoization inputs={t2} output={x} alwaysCheck={true} />;
+    t3 = <ValidateMemoization inputs={t2} output={x} />;
     $[7] = t2;
     $[8] = x;
     $[9] = t3;
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/todo-transitivity-createfrom-capture-lambda.tsx b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/todo-transitivity-createfrom-capture-lambda.tsx
index 9e10bec1b4..d6bd1690f6 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/todo-transitivity-createfrom-capture-lambda.tsx
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/todo-transitivity-createfrom-capture-lambda.tsx
@@ -18,7 +18,7 @@ function Component({a, b}) {
   typedMutate(z, b);
 
   // TODO: this *should* only depend on `a`
-  return <ValidateMemoization inputs={[a, b]} output={x} alwaysCheck={true} />;
+  return <ValidateMemoization inputs={[a, b]} output={x} />;
 }
 
 export const FIXTURE_ENTRYPOINT = {
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/transitivity-add-captured-array-to-itself.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/transitivity-add-captured-array-to-itself.expect.md
index 2813e072e2..fb3c8d0a89 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/transitivity-add-captured-array-to-itself.expect.md
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/transitivity-add-captured-array-to-itself.expect.md
@@ -20,8 +20,8 @@ function Component({a, b}) {
 
   return (
     <>
-      <ValidateMemoization inputs={[a]} output={o} alwaysCheck={true} />;
-      <ValidateMemoization inputs={[a, b]} output={x} alwaysCheck={true} />;
+      <ValidateMemoization inputs={[a]} output={o} />;
+      <ValidateMemoization inputs={[a, b]} output={x} />;
     </>
   );
 }
@@ -92,7 +92,7 @@ function Component(t0) {
   }
   let t5;
   if ($[8] !== o || $[9] !== t4) {
-    t5 = <ValidateMemoization inputs={t4} output={o} alwaysCheck={true} />;
+    t5 = <ValidateMemoization inputs={t4} output={o} />;
     $[8] = o;
     $[9] = t4;
     $[10] = t5;
@@ -110,7 +110,7 @@ function Component(t0) {
   }
   let t7;
   if ($[14] !== t6 || $[15] !== x) {
-    t7 = <ValidateMemoization inputs={t6} output={x} alwaysCheck={true} />;
+    t7 = <ValidateMemoization inputs={t6} output={x} />;
     $[14] = t6;
     $[15] = x;
     $[16] = t7;
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/transitivity-add-captured-array-to-itself.tsx b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/transitivity-add-captured-array-to-itself.tsx
index 7d7ec91bc5..d81c069e33 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/transitivity-add-captured-array-to-itself.tsx
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/transitivity-add-captured-array-to-itself.tsx
@@ -16,8 +16,8 @@ function Component({a, b}) {
 
   return (
     <>
-      <ValidateMemoization inputs={[a]} output={o} alwaysCheck={true} />;
-      <ValidateMemoization inputs={[a, b]} output={x} alwaysCheck={true} />;
+      <ValidateMemoization inputs={[a]} output={o} />;
+      <ValidateMemoization inputs={[a, b]} output={x} />;
     </>
   );
 }
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/transitivity-capture-createfrom-lambda.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/transitivity-capture-createfrom-lambda.expect.md
index 85c356f894..bfb4ede0ad 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/transitivity-capture-createfrom-lambda.expect.md
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/transitivity-capture-createfrom-lambda.expect.md
@@ -21,7 +21,7 @@ function Component({a, b}: {a: number; b: number}) {
   // mutates x
   typedMutate(z, b);
 
-  return <ValidateMemoization inputs={[a, b]} output={x} alwaysCheck={true} />;
+  return <ValidateMemoization inputs={[a, b]} output={x} />;
 }
 
 export const FIXTURE_ENTRYPOINT = {
@@ -85,7 +85,7 @@ function Component(t0) {
   }
   let t3;
   if ($[7] !== t2 || $[8] !== x) {
-    t3 = <ValidateMemoization inputs={t2} output={x} alwaysCheck={true} />;
+    t3 = <ValidateMemoization inputs={t2} output={x} />;
     $[7] = t2;
     $[8] = x;
     $[9] = t3;
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/transitivity-capture-createfrom-lambda.tsx b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/transitivity-capture-createfrom-lambda.tsx
index cee7901a6e..72289eb833 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/transitivity-capture-createfrom-lambda.tsx
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/transitivity-capture-createfrom-lambda.tsx
@@ -17,7 +17,7 @@ function Component({a, b}: {a: number; b: number}) {
   // mutates x
   typedMutate(z, b);
 
-  return <ValidateMemoization inputs={[a, b]} output={x} alwaysCheck={true} />;
+  return <ValidateMemoization inputs={[a, b]} output={x} />;
 }
 
 export const FIXTURE_ENTRYPOINT = {
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/transitivity-capture-createfrom.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/transitivity-capture-createfrom.expect.md
index af15d568ce..bd0a76965d 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/transitivity-capture-createfrom.expect.md
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/transitivity-capture-createfrom.expect.md
@@ -17,7 +17,7 @@ function Component({a, b}: {a: number; b: number}) {
   // mutates x
   typedMutate(z, b);
 
-  return <ValidateMemoization inputs={[a, b]} output={x} alwaysCheck={true} />;
+  return <ValidateMemoization inputs={[a, b]} output={x} />;
 }
 
 export const FIXTURE_ENTRYPOINT = {
@@ -76,7 +76,7 @@ function Component(t0) {
   }
   let t3;
   if ($[7] !== t2 || $[8] !== x) {
-    t3 = <ValidateMemoization inputs={t2} output={x} alwaysCheck={true} />;
+    t3 = <ValidateMemoization inputs={t2} output={x} />;
     $[7] = t2;
     $[8] = x;
     $[9] = t3;
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/transitivity-capture-createfrom.tsx b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/transitivity-capture-createfrom.tsx
index df89c3413e..d06ad11eb5 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/transitivity-capture-createfrom.tsx
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/transitivity-capture-createfrom.tsx
@@ -13,7 +13,7 @@ function Component({a, b}: {a: number; b: number}) {
   // mutates x
   typedMutate(z, b);
 
-  return <ValidateMemoization inputs={[a, b]} output={x} alwaysCheck={true} />;
+  return <ValidateMemoization inputs={[a, b]} output={x} />;
 }
 
 export const FIXTURE_ENTRYPOINT = {
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/transitivity-createfrom-capture.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/transitivity-createfrom-capture.expect.md
index 25dbcdf3ac..742a233e80 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/transitivity-createfrom-capture.expect.md
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/transitivity-createfrom-capture.expect.md
@@ -17,7 +17,7 @@ function Component({a, b}) {
   // does not mutate x, so x should not depend on b
   typedMutate(z, b);
 
-  return <ValidateMemoization inputs={[a]} output={x} alwaysCheck={true} />;
+  return <ValidateMemoization inputs={[a]} output={x} />;
 }
 
 export const FIXTURE_ENTRYPOINT = {
@@ -73,7 +73,7 @@ function Component(t0) {
   }
   let t4;
   if ($[4] !== t3 || $[5] !== x) {
-    t4 = <ValidateMemoization inputs={t3} output={x} alwaysCheck={true} />;
+    t4 = <ValidateMemoization inputs={t3} output={x} />;
     $[4] = t3;
     $[5] = x;
     $[6] = t4;
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/transitivity-createfrom-capture.tsx b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/transitivity-createfrom-capture.tsx
index a8a02564bd..32d65e61e0 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/transitivity-createfrom-capture.tsx
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/transitivity-createfrom-capture.tsx
@@ -13,7 +13,7 @@ function Component({a, b}) {
   // does not mutate x, so x should not depend on b
   typedMutate(z, b);
 
-  return <ValidateMemoization inputs={[a]} output={x} alwaysCheck={true} />;
+  return <ValidateMemoization inputs={[a]} output={x} />;
 }
 
 export const FIXTURE_ENTRYPOINT = {
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/transitivity-phi-assign-or-capture.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/transitivity-phi-assign-or-capture.expect.md
index b5350352cf..c50fe47aed 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/transitivity-phi-assign-or-capture.expect.md
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/transitivity-phi-assign-or-capture.expect.md
@@ -21,7 +21,7 @@ function Component({a, b}) {
   // could mutate x
   typedMutate(z, b);
 
-  return <ValidateMemoization inputs={[a, b]} output={x} alwaysCheck={true} />;
+  return <ValidateMemoization inputs={[a, b]} output={x} />;
 }
 
 export const FIXTURE_ENTRYPOINT = {
@@ -92,7 +92,7 @@ function Component(t0) {
   }
   let t4;
   if ($[9] !== t3 || $[10] !== x) {
-    t4 = <ValidateMemoization inputs={t3} output={x} alwaysCheck={true} />;
+    t4 = <ValidateMemoization inputs={t3} output={x} />;
     $[9] = t3;
     $[10] = x;
     $[11] = t4;
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/transitivity-phi-assign-or-capture.tsx b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/transitivity-phi-assign-or-capture.tsx
index e583bfa83a..90b7597694 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/transitivity-phi-assign-or-capture.tsx
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/transitivity-phi-assign-or-capture.tsx
@@ -17,7 +17,7 @@ function Component({a, b}) {
   // could mutate x
   typedMutate(z, b);
 
-  return <ValidateMemoization inputs={[a, b]} output={x} alwaysCheck={true} />;
+  return <ValidateMemoization inputs={[a, b]} output={x} />;
 }
 
 export const FIXTURE_ENTRYPOINT = {
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/typed-identity-function-frozen-input.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/typed-identity-function-frozen-input.expect.md
index b15248df07..dc8bedaf83 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/typed-identity-function-frozen-input.expect.md
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/typed-identity-function-frozen-input.expect.md
@@ -4,6 +4,7 @@
 ```javascript
 // @enableNewMutationAliasingModel
 
+import {useMemo} from 'react';
 import {
   identity,
   makeObject_Primitives,
@@ -14,7 +15,7 @@ import {
 
 function Component({a, b}) {
   // create a mutable value with input `a`
-  const x = makeObject_Primitives(a);
+  const x = useMemo(() => makeObject_Primitives(a), [a]);
 
   // freeze the value
   useIdentity(x);
@@ -49,6 +50,7 @@ export const FIXTURE_ENTRYPOINT = {
 ```javascript
 import { c as _c } from "react/compiler-runtime"; // @enableNewMutationAliasingModel
 
+import { useMemo } from "react";
 import {
   identity,
   makeObject_Primitives,
@@ -61,13 +63,15 @@ function Component(t0) {
   const $ = _c(7);
   const { a, b } = t0;
   let t1;
+  let t2;
   if ($[0] !== a) {
-    t1 = makeObject_Primitives(a);
+    t2 = makeObject_Primitives(a);
     $[0] = a;
-    $[1] = t1;
+    $[1] = t2;
   } else {
-    t1 = $[1];
+    t2 = $[1];
   }
+  t1 = t2;
   const x = t1;
 
   useIdentity(x);
@@ -75,24 +79,24 @@ function Component(t0) {
   const x2 = typedIdentity(x);
 
   identity(x2, b);
-  let t2;
-  if ($[2] !== a) {
-    t2 = [a];
-    $[2] = a;
-    $[3] = t2;
-  } else {
-    t2 = $[3];
-  }
   let t3;
-  if ($[4] !== t2 || $[5] !== x) {
-    t3 = <ValidateMemoization inputs={t2} output={x} />;
-    $[4] = t2;
-    $[5] = x;
-    $[6] = t3;
+  if ($[2] !== a) {
+    t3 = [a];
+    $[2] = a;
+    $[3] = t3;
   } else {
-    t3 = $[6];
+    t3 = $[3];
   }
-  return t3;
+  let t4;
+  if ($[4] !== t3 || $[5] !== x) {
+    t4 = <ValidateMemoization inputs={t3} output={x} />;
+    $[4] = t3;
+    $[5] = x;
+    $[6] = t4;
+  } else {
+    t4 = $[6];
+  }
+  return t4;
 }
 
 export const FIXTURE_ENTRYPOINT = {
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/typed-identity-function-frozen-input.js b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/typed-identity-function-frozen-input.js
index bcf6ecef02..d0f677ee4d 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/typed-identity-function-frozen-input.js
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/typed-identity-function-frozen-input.js
@@ -1,5 +1,6 @@
 // @enableNewMutationAliasingModel
 
+import {useMemo} from 'react';
 import {
   identity,
   makeObject_Primitives,
@@ -10,7 +11,7 @@ import {
 
 function Component({a, b}) {
   // create a mutable value with input `a`
-  const x = makeObject_Primitives(a);
+  const x = useMemo(() => makeObject_Primitives(a), [a]);
 
   // freeze the value
   useIdentity(x);
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/weakmap-constructor.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/weakmap-constructor.expect.md
index aebaedf6a8..a6def457a4 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/weakmap-constructor.expect.md
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/weakmap-constructor.expect.md
@@ -2,6 +2,7 @@
 ## Input
 
 ```javascript
+import {useMemo} from 'react';
 import {ValidateMemoization} from 'shared-runtime';
 
 function Component({a, b, c}) {
@@ -13,9 +14,21 @@ function Component({a, b, c}) {
 
   return (
     <>
-      <ValidateMemoization inputs={[a, c]} output={map} />
-      <ValidateMemoization inputs={[a, c]} output={mapAlias} />
-      <ValidateMemoization inputs={[b]} output={[hasB]} />
+      <ValidateMemoization
+        inputs={[a, c]}
+        output={map}
+        onlyCheckCompiled={true}
+      />
+      <ValidateMemoization
+        inputs={[a, c]}
+        output={mapAlias}
+        onlyCheckCompiled={true}
+      />
+      <ValidateMemoization
+        inputs={[b]}
+        output={[hasB]}
+        onlyCheckCompiled={true}
+      />
     </>
   );
 }
@@ -44,6 +57,7 @@ export const FIXTURE_ENTRYPOINT = {
 
 ```javascript
 import { c as _c } from "react/compiler-runtime";
+import { useMemo } from "react";
 import { ValidateMemoization } from "shared-runtime";
 
 function Component(t0) {
@@ -76,7 +90,9 @@ function Component(t0) {
   }
   let t2;
   if ($[7] !== map || $[8] !== t1) {
-    t2 = <ValidateMemoization inputs={t1} output={map} />;
+    t2 = (
+      <ValidateMemoization inputs={t1} output={map} onlyCheckCompiled={true} />
+    );
     $[7] = map;
     $[8] = t1;
     $[9] = t2;
@@ -94,7 +110,13 @@ function Component(t0) {
   }
   let t4;
   if ($[13] !== mapAlias || $[14] !== t3) {
-    t4 = <ValidateMemoization inputs={t3} output={mapAlias} />;
+    t4 = (
+      <ValidateMemoization
+        inputs={t3}
+        output={mapAlias}
+        onlyCheckCompiled={true}
+      />
+    );
     $[13] = mapAlias;
     $[14] = t3;
     $[15] = t4;
@@ -119,7 +141,9 @@ function Component(t0) {
   }
   let t7;
   if ($[20] !== t5 || $[21] !== t6) {
-    t7 = <ValidateMemoization inputs={t5} output={t6} />;
+    t7 = (
+      <ValidateMemoization inputs={t5} output={t6} onlyCheckCompiled={true} />
+    );
     $[20] = t5;
     $[21] = t6;
     $[22] = t7;
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/weakmap-constructor.js b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/weakmap-constructor.js
index ecdec6c9e9..d005c9f271 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/weakmap-constructor.js
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/weakmap-constructor.js
@@ -1,3 +1,4 @@
+import {useMemo} from 'react';
 import {ValidateMemoization} from 'shared-runtime';
 
 function Component({a, b, c}) {
@@ -9,9 +10,21 @@ function Component({a, b, c}) {
 
   return (
     <>
-      <ValidateMemoization inputs={[a, c]} output={map} />
-      <ValidateMemoization inputs={[a, c]} output={mapAlias} />
-      <ValidateMemoization inputs={[b]} output={[hasB]} />
+      <ValidateMemoization
+        inputs={[a, c]}
+        output={map}
+        onlyCheckCompiled={true}
+      />
+      <ValidateMemoization
+        inputs={[a, c]}
+        output={mapAlias}
+        onlyCheckCompiled={true}
+      />
+      <ValidateMemoization
+        inputs={[b]}
+        output={[hasB]}
+        onlyCheckCompiled={true}
+      />
     </>
   );
 }
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/weakset-constructor.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/weakset-constructor.expect.md
index 5ebf32d533..94e0c7f055 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/weakset-constructor.expect.md
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/weakset-constructor.expect.md
@@ -2,6 +2,7 @@
 ## Input
 
 ```javascript
+import {useMemo} from 'react';
 import {ValidateMemoization} from 'shared-runtime';
 
 function Component({a, b, c}) {
@@ -13,9 +14,21 @@ function Component({a, b, c}) {
 
   return (
     <>
-      <ValidateMemoization inputs={[a, c]} output={set} />
-      <ValidateMemoization inputs={[a, c]} output={setAlias} />
-      <ValidateMemoization inputs={[b]} output={[hasB]} />
+      <ValidateMemoization
+        inputs={[a, c]}
+        output={set}
+        onlyCheckCompiled={true}
+      />
+      <ValidateMemoization
+        inputs={[a, c]}
+        output={setAlias}
+        onlyCheckCompiled={true}
+      />
+      <ValidateMemoization
+        inputs={[b]}
+        output={[hasB]}
+        onlyCheckCompiled={true}
+      />
     </>
   );
 }
@@ -44,6 +57,7 @@ export const FIXTURE_ENTRYPOINT = {
 
 ```javascript
 import { c as _c } from "react/compiler-runtime";
+import { useMemo } from "react";
 import { ValidateMemoization } from "shared-runtime";
 
 function Component(t0) {
@@ -76,7 +90,9 @@ function Component(t0) {
   }
   let t2;
   if ($[7] !== set || $[8] !== t1) {
-    t2 = <ValidateMemoization inputs={t1} output={set} />;
+    t2 = (
+      <ValidateMemoization inputs={t1} output={set} onlyCheckCompiled={true} />
+    );
     $[7] = set;
     $[8] = t1;
     $[9] = t2;
@@ -94,7 +110,13 @@ function Component(t0) {
   }
   let t4;
   if ($[13] !== setAlias || $[14] !== t3) {
-    t4 = <ValidateMemoization inputs={t3} output={setAlias} />;
+    t4 = (
+      <ValidateMemoization
+        inputs={t3}
+        output={setAlias}
+        onlyCheckCompiled={true}
+      />
+    );
     $[13] = setAlias;
     $[14] = t3;
     $[15] = t4;
@@ -119,7 +141,9 @@ function Component(t0) {
   }
   let t7;
   if ($[20] !== t5 || $[21] !== t6) {
-    t7 = <ValidateMemoization inputs={t5} output={t6} />;
+    t7 = (
+      <ValidateMemoization inputs={t5} output={t6} onlyCheckCompiled={true} />
+    );
     $[20] = t5;
     $[21] = t6;
     $[22] = t7;
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/weakset-constructor.js b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/weakset-constructor.js
index 8c0a7deb18..9114233812 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/weakset-constructor.js
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/weakset-constructor.js
@@ -1,3 +1,4 @@
+import {useMemo} from 'react';
 import {ValidateMemoization} from 'shared-runtime';
 
 function Component({a, b, c}) {
@@ -9,9 +10,21 @@ function Component({a, b, c}) {
 
   return (
     <>
-      <ValidateMemoization inputs={[a, c]} output={set} />
-      <ValidateMemoization inputs={[a, c]} output={setAlias} />
-      <ValidateMemoization inputs={[b]} output={[hasB]} />
+      <ValidateMemoization
+        inputs={[a, c]}
+        output={set}
+        onlyCheckCompiled={true}
+      />
+      <ValidateMemoization
+        inputs={[a, c]}
+        output={setAlias}
+        onlyCheckCompiled={true}
+      />
+      <ValidateMemoization
+        inputs={[b]}
+        output={[hasB]}
+        onlyCheckCompiled={true}
+      />
     </>
   );
 }
diff --git a/compiler/packages/snap/src/sprout/shared-runtime.ts b/compiler/packages/snap/src/sprout/shared-runtime.ts
index 1e0ab108db..f37ca82709 100644
--- a/compiler/packages/snap/src/sprout/shared-runtime.ts
+++ b/compiler/packages/snap/src/sprout/shared-runtime.ts
@@ -269,12 +269,10 @@ export function ValidateMemoization({
   inputs,
   output: rawOutput,
   onlyCheckCompiled = false,
-  alwaysCheck = false,
 }: {
   inputs: Array<any>;
   output: any;
   onlyCheckCompiled?: boolean;
-  alwaysCheck?: boolean;
 }): React.ReactElement {
   'use no forget';
   // Wrap rawOutput as it might be a function, which useState would invoke.
@@ -282,7 +280,7 @@ export function ValidateMemoization({
   const [previousInputs, setPreviousInputs] = React.useState(inputs);
   const [previousOutput, setPreviousOutput] = React.useState(output);
   if (
-    alwaysCheck ||
+    !onlyCheckCompiled ||
     (onlyCheckCompiled &&
       (globalThis as any).__SNAP_EVALUATOR_MODE === 'forget')
   ) {