From bd36112553ba770f461e0c6821d40bf92aafaad5 Mon Sep 17 00:00:00 2001
From: Lauren Tan <poteto@users.noreply.github.com>
Date: Fri, 21 Apr 2023 15:48:41 -0400
Subject: [PATCH] [DevTools] Add support for useMemoCache

useMemoCache wasn't previously supported in the DevTools, so any attempt
to inspect a component using the hook would result in a
`dispatcher.useMemoCache is not a function (it is undefined)` error.
---
 .../react-debug-tools/src/ReactDebugHooks.js  | 48 +++++++++++++++++++
 .../ReactHooksInspectionIntegration-test.js   | 36 ++++++++++++++
 2 files changed, 84 insertions(+)

diff --git a/packages/react-debug-tools/src/ReactDebugHooks.js b/packages/react-debug-tools/src/ReactDebugHooks.js
index 362e5f36d2bb1..4829722558d08 100644
--- a/packages/react-debug-tools/src/ReactDebugHooks.js
+++ b/packages/react-debug-tools/src/ReactDebugHooks.js
@@ -30,6 +30,8 @@ import {
   ForwardRef,
 } from 'react-reconciler/src/ReactWorkTags';
 
+const MEMO_CACHE_SENTINEL = Symbol.for('react.memo_cache_sentinel');
+
 type CurrentDispatcherRef = typeof ReactSharedInternals.ReactCurrentDispatcher;
 
 // Used to track hooks called during a render
@@ -51,9 +53,19 @@ type Dispatch<A> = A => void;
 
 let primitiveStackCache: null | Map<string, Array<any>> = null;
 
+type MemoCache = {
+  data: Array<Array<any>>,
+  index: number,
+};
+
+type FunctionComponentUpdateQueue = {
+  memoCache?: MemoCache | null,
+};
+
 type Hook = {
   memoizedState: any,
   next: Hook | null,
+  updateQueue: FunctionComponentUpdateQueue | null,
 };
 
 function getPrimitiveStackCache(): Map<string, Array<any>> {
@@ -79,6 +91,10 @@ function getPrimitiveStackCache(): Map<string, Array<any>> {
       Dispatcher.useDebugValue(null);
       Dispatcher.useCallback(() => {});
       Dispatcher.useMemo(() => null);
+      if (typeof Dispatcher.useMemoCache === 'function') {
+        // This type check is for Flow only.
+        Dispatcher.useMemoCache(0);
+      }
     } finally {
       readHookLog = hookLog;
       hookLog = [];
@@ -333,6 +349,37 @@ function useId(): string {
   return id;
 }
 
+function useMemoCache(size: number): Array<any> {
+  const hook = nextHook();
+  let memoCache: MemoCache;
+  if (
+    hook !== null &&
+    hook.updateQueue !== null &&
+    hook.updateQueue.memoCache != null
+  ) {
+    memoCache = hook.updateQueue.memoCache;
+  } else {
+    memoCache = {
+      data: [],
+      index: 0,
+    };
+  }
+
+  let data = memoCache.data[memoCache.index];
+  if (data === undefined) {
+    data = new Array(size);
+    for (let i = 0; i < size; i++) {
+      data[i] = MEMO_CACHE_SENTINEL;
+    }
+  }
+  hookLog.push({
+    primitive: 'MemoCache',
+    stackError: new Error(),
+    value: data,
+  });
+  return data;
+}
+
 const Dispatcher: DispatcherType = {
   use,
   readContext,
@@ -345,6 +392,7 @@ const Dispatcher: DispatcherType = {
   useLayoutEffect,
   useInsertionEffect,
   useMemo,
+  useMemoCache,
   useReducer,
   useRef,
   useState,
diff --git a/packages/react-debug-tools/src/__tests__/ReactHooksInspectionIntegration-test.js b/packages/react-debug-tools/src/__tests__/ReactHooksInspectionIntegration-test.js
index 5e4860ce7045b..e2e90b4428b81 100644
--- a/packages/react-debug-tools/src/__tests__/ReactHooksInspectionIntegration-test.js
+++ b/packages/react-debug-tools/src/__tests__/ReactHooksInspectionIntegration-test.js
@@ -633,6 +633,42 @@ describe('ReactHooksInspectionIntegration', () => {
     });
   });
 
+  // @gate enableUseMemoCacheHook
+  it('should support useMemoCache hook', () => {
+    function Foo() {
+      const $ = React.unstable_useMemoCache(1);
+      let t0;
+
+      if ($[0] === Symbol.for('react.memo_cache_sentinel')) {
+        t0 = <div>{1}</div>;
+        $[0] = t0;
+      } else {
+        t0 = $[0];
+      }
+
+      return t0;
+    }
+
+    const renderer = ReactTestRenderer.create(<Foo />);
+    const childFiber = renderer.root.findByType(Foo)._currentFiber();
+    const tree = ReactDebugTools.inspectHooksOfFiber(childFiber);
+
+    expect(tree.length).toEqual(1);
+    expect(tree[0]).toMatchInlineSnapshot(`
+      {
+        "id": 0,
+        "isStateEditable": false,
+        "name": "MemoCache",
+        "subHooks": [],
+        "value": [
+          <div>
+            1
+          </div>,
+        ],
+      }
+    `);
+  });
+
   describe('useDebugValue', () => {
     it('should support inspectable values for multiple custom hooks', () => {
       function useLabeledValue(label) {