[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

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
+      <ValidateMemoization inputs={[a, b, c]} output={x[0]} />;
-        inputs={[a, b, c]}
-        output={x[0]}
-        alwaysCheck={true}
-      />
-      ;
     </>
   );
 }
@@ -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;

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
+      <ValidateMemoization inputs={[a, b, c]} output={x[0]} />;
-        inputs={[a, b, c]}
-        output={x[0]}
-        alwaysCheck={true}
-      />
-      ;
     </>
   );
 }

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;

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 = {

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]} output={o} />;
-      <ValidateMemoization inputs={[a, b]} output={x} alwaysCheck={true} />;
+      <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;

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]} output={o} />;
-      <ValidateMemoization inputs={[a, b]} output={x} alwaysCheck={true} />;
+      <ValidateMemoization inputs={[a, b]} output={x} />;
     </>
   );
 }

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;

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 = {

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;

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 = {

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;

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 = {

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;

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 = {

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) {
+  if ($[2] !== a) {
-    t3 = <ValidateMemoization inputs={t2} output={x} />;
+    t3 = [a];
-    $[4] = t2;
+    $[2] = a;
-    $[5] = x;
+    $[3] = t3;
-    $[6] = 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 = {

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

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
-      <ValidateMemoization inputs={[a, c]} output={mapAlias} />
+        inputs={[a, c]}
-      <ValidateMemoization inputs={[b]} output={[hasB]} />
+        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;

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
-      <ValidateMemoization inputs={[a, c]} output={mapAlias} />
+        inputs={[a, c]}
-      <ValidateMemoization inputs={[b]} output={[hasB]} />
+        output={map}
+        onlyCheckCompiled={true}
+      />
+      <ValidateMemoization
+        inputs={[a, c]}
+        output={mapAlias}
+        onlyCheckCompiled={true}
+      />
+      <ValidateMemoization
+        inputs={[b]}
+        output={[hasB]}
+        onlyCheckCompiled={true}
+      />
     </>
   );
 }

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
-      <ValidateMemoization inputs={[a, c]} output={setAlias} />
+        inputs={[a, c]}
-      <ValidateMemoization inputs={[b]} output={[hasB]} />
+        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;

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
-      <ValidateMemoization inputs={[a, c]} output={setAlias} />
+        inputs={[a, c]}
-      <ValidateMemoization inputs={[b]} output={[hasB]} />
+        output={set}
+        onlyCheckCompiled={true}
+      />
+      <ValidateMemoization
+        inputs={[a, c]}
+        output={setAlias}
+        onlyCheckCompiled={true}
+      />
+      <ValidateMemoization
+        inputs={[b]}
+        output={[hasB]}
+        onlyCheckCompiled={true}
+      />
     </>
   );
 }

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')
   ) {