From c03f1c809aeaf69ed26eda39df9dbd7242bc014d Mon Sep 17 00:00:00 2001
From: Joe Savona <joesavona@fb.com>
Date: Tue, 19 Mar 2024 13:28:47 -0700
Subject: [PATCH] Extract shared-runtime helper for validating memoization

Extracts a helper from the repro earlier in the stack into a helper in
shared-runtime. This makes it easy to verify that memoization is actually
working.
---
 ...ble-code-early-return-in-useMemo.expect.md | 110 +++++++-----------
 ...nreachable-code-early-return-in-useMemo.js |  29 +----
 .../snap/src/sprout/shared-runtime.ts         |  24 ++++
 3 files changed, 68 insertions(+), 95 deletions(-)

diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/repro-unreachable-code-early-return-in-useMemo.expect.md b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/repro-unreachable-code-early-return-in-useMemo.expect.md
index 9a7244fd0e..14a5504a6d 100644
--- a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/repro-unreachable-code-early-return-in-useMemo.expect.md
+++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/repro-unreachable-code-early-return-in-useMemo.expect.md
@@ -4,34 +4,10 @@
 ```javascript
 // @enableAssumeHooksFollowRulesOfReact @enableTransitivelyFreezeFunctionExpressions
 import { useMemo, useState } from "react";
-import { Stringify, identity } from "shared-runtime";
+import { ValidateMemoization, identity } from "shared-runtime";
 
 function Component({ value }) {
-  "use no forget";
-  const result = useValue(value);
-  return <Validate inputs={[value]} output={result} />;
-}
-
-function Validate({ inputs, output }) {
-  "use no forget";
-  const [previousInputs, setPreviousInputs] = useState(inputs);
-  const [previousOutput, setPreviousOutput] = useState(output);
-  if (
-    inputs.length !== previousInputs.length ||
-    inputs.some((item, i) => item !== previousInputs[i])
-  ) {
-    // Some input changed, we expect the output to change
-    setPreviousInputs(inputs);
-    setPreviousOutput(output);
-  } else if (output !== previousOutput) {
-    // Else output should be stable
-    throw new Error("Output identity changed but inputs did not");
-  }
-  return <Stringify inputs={inputs} output={output} />;
-}
-
-function useValue(value) {
-  return useMemo(() => {
+  const result = useMemo(() => {
     if (value == null) {
       return null;
     }
@@ -41,6 +17,7 @@ function useValue(value) {
       return null;
     }
   }, [value]);
+  return <ValidateMemoization inputs={[value]} output={result} />;
 }
 
 export const FIXTURE_ENTRYPOINT = {
@@ -69,70 +46,65 @@ import {
   useState,
   unstable_useMemoCache as useMemoCache,
 } from "react";
-import { Stringify, identity } from "shared-runtime";
+import { ValidateMemoization, identity } from "shared-runtime";
 
-function Component({ value }) {
-  "use no forget";
-  const result = useValue(value);
-  return <Validate inputs={[value]} output={result} />;
-}
-
-function Validate({ inputs, output }) {
-  "use no forget";
-  const [previousInputs, setPreviousInputs] = useState(inputs);
-  const [previousOutput, setPreviousOutput] = useState(output);
-  if (
-    inputs.length !== previousInputs.length ||
-    inputs.some((item, i) => item !== previousInputs[i])
-  ) {
-    // Some input changed, we expect the output to change
-    setPreviousInputs(inputs);
-    setPreviousOutput(output);
-  } else if (output !== previousOutput) {
-    // Else output should be stable
-    throw new Error("Output identity changed but inputs did not");
-  }
-  return <Stringify inputs={inputs} output={output} />;
-}
-
-function useValue(value) {
-  const $ = useMemoCache(5);
-  let t0;
+function Component(t0) {
+  const $ = useMemoCache(10);
+  const { value } = t0;
+  let t1;
   bb13: {
     if (value == null) {
       if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
-        t0 = null;
+        t1 = null;
         break bb13;
-        $[0] = t0;
+        $[0] = t1;
       } else {
-        t0 = $[0];
+        t1 = $[0];
       }
     }
     try {
-      let t2;
+      let t3;
       if ($[1] !== value) {
-        t2 = { value };
+        t3 = { value };
         $[1] = value;
-        $[2] = t2;
+        $[2] = t3;
       } else {
-        t2 = $[2];
+        t3 = $[2];
       }
       if ($[3] === Symbol.for("react.memo_cache_sentinel")) {
-        t0 = t2;
-        $[3] = t0;
+        t1 = t3;
+        $[3] = t1;
       } else {
-        t0 = $[3];
+        t1 = $[3];
       }
-    } catch (t1) {
+    } catch (t2) {
       if ($[4] === Symbol.for("react.memo_cache_sentinel")) {
-        t0 = null;
-        $[4] = t0;
+        t1 = null;
+        $[4] = t1;
       } else {
-        t0 = $[4];
+        t1 = $[4];
       }
     }
   }
-  return t0;
+  const result = t1;
+  let t2;
+  if ($[5] !== value) {
+    t2 = [value];
+    $[5] = value;
+    $[6] = t2;
+  } else {
+    t2 = $[6];
+  }
+  let t3;
+  if ($[7] !== t2 || $[8] !== result) {
+    t3 = <ValidateMemoization inputs={t2} output={result} />;
+    $[7] = t2;
+    $[8] = result;
+    $[9] = t3;
+  } else {
+    t3 = $[9];
+  }
+  return t3;
 }
 
 export const FIXTURE_ENTRYPOINT = {
diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/repro-unreachable-code-early-return-in-useMemo.js b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/repro-unreachable-code-early-return-in-useMemo.js
index 33a27aa802..b14c02c2ee 100644
--- a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/repro-unreachable-code-early-return-in-useMemo.js
+++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/repro-unreachable-code-early-return-in-useMemo.js
@@ -1,33 +1,9 @@
 // @enableAssumeHooksFollowRulesOfReact @enableTransitivelyFreezeFunctionExpressions
 import { useMemo, useState } from "react";
-import { Stringify, identity } from "shared-runtime";
+import { ValidateMemoization, identity } from "shared-runtime";
 
 function Component({ value }) {
-  "use no forget";
-  const result = useValue(value);
-  return <Validate inputs={[value]} output={result} />;
-}
-
-function Validate({ inputs, output }) {
-  "use no forget";
-  const [previousInputs, setPreviousInputs] = useState(inputs);
-  const [previousOutput, setPreviousOutput] = useState(output);
-  if (
-    inputs.length !== previousInputs.length ||
-    inputs.some((item, i) => item !== previousInputs[i])
-  ) {
-    // Some input changed, we expect the output to change
-    setPreviousInputs(inputs);
-    setPreviousOutput(output);
-  } else if (output !== previousOutput) {
-    // Else output should be stable
-    throw new Error("Output identity changed but inputs did not");
-  }
-  return <Stringify inputs={inputs} output={output} />;
-}
-
-function useValue(value) {
-  return useMemo(() => {
+  const result = useMemo(() => {
     if (value == null) {
       return null;
     }
@@ -37,6 +13,7 @@ function useValue(value) {
       return null;
     }
   }, [value]);
+  return <ValidateMemoization inputs={[value]} output={result} />;
 }
 
 export const FIXTURE_ENTRYPOINT = {
diff --git a/compiler/packages/snap/src/sprout/shared-runtime.ts b/compiler/packages/snap/src/sprout/shared-runtime.ts
index 6b6b87c2ca..6229fe1211 100644
--- a/compiler/packages/snap/src/sprout/shared-runtime.ts
+++ b/compiler/packages/snap/src/sprout/shared-runtime.ts
@@ -226,6 +226,30 @@ export function Stringify(props: any): React.ReactElement {
   );
 }
 
+export function ValidateMemoization({
+  inputs,
+  output,
+}: {
+  inputs: Array<any>;
+  output: any;
+}) {
+  "use no forget";
+  const [previousInputs, setPreviousInputs] = React.useState(inputs);
+  const [previousOutput, setPreviousOutput] = React.useState(output);
+  if (
+    inputs.length !== previousInputs.length ||
+    inputs.some((item, i) => item !== previousInputs[i])
+  ) {
+    // Some input changed, we expect the output to change
+    setPreviousInputs(inputs);
+    setPreviousOutput(output);
+  } else if (output !== previousOutput) {
+    // Else output should be stable
+    throw new Error("Output identity changed but inputs did not");
+  }
+  return React.createElement(Stringify, { inputs, output });
+}
+
 export function createHookWrapper<TProps, TRet>(
   useMaybeHook: (props: TProps) => TRet
 ): FunctionComponent<TProps> {