fix[devtools/useMemoCache]: implement a working copy of useMemoCache (#27659)

In #27472 I've removed broken
`useMemoCache` implementation and replaced it with a stub. It actually
produces errors when trying to inspect components, which are compiled
with Forget.

The main difference from the implementation in
#26696 is that we are using
corresponding `Fiber` here, which has patched `updateQueue` with
`memoCache`. Previously we would check it on a hook object, which
doesn't have `updateQueue`.

Tested on pages, which are using Forget and by inspecting elements,
which are transpiled with Forget.

Author: hoxyq

packages/react-debug-tools/src/ReactDebugHooks.js

@@ -26,6 +26,7 @@ import {
   ContextProvider,
   ForwardRef,
 } from 'react-reconciler/src/ReactWorkTags';
+import {REACT_MEMO_CACHE_SENTINEL} from 'shared/ReactSymbols';
 
 type CurrentDispatcherRef = typeof ReactSharedInternals.ReactCurrentDispatcher;
 
@@ -93,7 +94,9 @@ function getPrimitiveStackCache(): Map<string, Array<any>> {
   return primitiveStackCache;
 }
 
+let currentFiber: null | Fiber = null;
 let currentHook: null | Hook = null;
+
 function nextHook(): null | Hook {
   const hook = currentHook;
   if (hook !== null) {
@@ -319,9 +322,31 @@ function useId(): string {
 
 // useMemoCache is an implementation detail of Forget's memoization
 // it should not be called directly in user-generated code
-// we keep it as a stub for dispatcher
 function useMemoCache(size: number): Array<any> {
-  return [];
+  const fiber = currentFiber;
+  // Don't throw, in case this is called from getPrimitiveStackCache
+  if (fiber == null) {
+    return [];
+  }
+
+  // $FlowFixMe[incompatible-use]: updateQueue is mixed
+  const memoCache = fiber.updateQueue?.memoCache;
+  if (memoCache == null) {
+    return [];
+  }
+
+  let data = memoCache.data[memoCache.index];
+  if (data === undefined) {
+    data = memoCache.data[memoCache.index] = new Array(size);
+    for (let i = 0; i < size; i++) {
+      data[i] = REACT_MEMO_CACHE_SENTINEL;
+    }
+  }
+
+  // We don't write anything to hookLog on purpose, so this hook remains invisible to users.
+
+  memoCache.index++;
+  return data;
 }
 
 const Dispatcher: DispatcherType = {
@@ -699,9 +724,11 @@ export function inspectHooks<Props>(
   }
 
   const previousDispatcher = currentDispatcher.current;
-  let readHookLog;
   currentDispatcher.current = DispatcherProxy;
+
+  let readHookLog;
   let ancestorStackError;
+
   try {
     ancestorStackError = new Error();
     renderFunction(props);
@@ -798,19 +825,25 @@ export function inspectHooksOfFiber(
       'Unknown Fiber. Needs to be a function component to inspect hooks.',
     );
   }
+
   // Warm up the cache so that it doesn't consume the currentHook.
   getPrimitiveStackCache();
+
+  // Set up the current hook so that we can step through and read the
+  // current state from them.
+  currentHook = (fiber.memoizedState: Hook);
+  currentFiber = fiber;
+
   const type = fiber.type;
   let props = fiber.memoizedProps;
   if (type !== fiber.elementType) {
     props = resolveDefaultProps(type, props);
   }
-  // Set up the current hook so that we can step through and read the
-  // current state from them.
-  currentHook = (fiber.memoizedState: Hook);
-  const contextMap = new Map<ReactContext<$FlowFixMe>, $FlowFixMe>();
+
+  const contextMap = new Map<ReactContext<any>, any>();
   try {
     setupContexts(contextMap, fiber);
+
     if (fiber.tag === ForwardRef) {
       return inspectHooksOfForwardRef(
         type.render,
@@ -820,9 +853,12 @@ export function inspectHooksOfFiber(
         includeHooksSource,
       );
     }
+
     return inspectHooks(type, props, currentDispatcher, includeHooksSource);
   } finally {
+    currentFiber = null;
     currentHook = null;
+
     restoreContexts(contextMap);
   }
 }

packages/react-debug-tools/src/__tests__/ReactHooksInspectionIntegration-test.js

@@ -633,27 +633,64 @@ describe('ReactHooksInspectionIntegration', () => {
     });
   });
 
-  // @gate enableUseMemoCacheHook
-  it('useMemoCache should not be inspectable', () => {
-    function Foo() {
-      const $ = useMemoCache(1);
-      let t0;
-
-      if ($[0] === Symbol.for('react.memo_cache_sentinel')) {
-        t0 = <div>{1}</div>;
-        $[0] = t0;
-      } else {
-        t0 = $[0];
+  describe('useMemoCache', () => {
+    // @gate enableUseMemoCacheHook
+    it('should not be inspectable', () => {
+      function Foo() {
+        const $ = useMemoCache(1);
+        let t0;
+
+        if ($[0] === Symbol.for('react.memo_cache_sentinel')) {
+          t0 = <div>{1}</div>;
+          $[0] = t0;
+        } else {
+          t0 = $[0];
+        }
+
+        return t0;
       }
 
-      return t0;
-    }
+      const renderer = ReactTestRenderer.create(<Foo />);
+      const childFiber = renderer.root.findByType(Foo)._currentFiber();
+      const tree = ReactDebugTools.inspectHooksOfFiber(childFiber);
 
-    const renderer = ReactTestRenderer.create(<Foo />);
-    const childFiber = renderer.root.findByType(Foo)._currentFiber();
-    const tree = ReactDebugTools.inspectHooksOfFiber(childFiber);
+      expect(tree.length).toEqual(0);
+    });
+
+    // @gate enableUseMemoCacheHook
+    it('should work in combination with other hooks', () => {
+      function useSomething() {
+        const [something] = React.useState(null);
+        const changeOtherSomething = React.useCallback(() => {}, [something]);
+
+        return [something, changeOtherSomething];
+      }
 
-    expect(tree.length).toEqual(0);
+      function Foo() {
+        const $ = useMemoCache(10);
+
+        useSomething();
+        React.useState(1);
+        React.useEffect(() => {});
+
+        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(3);
+    });
   });
 
   describe('useDebugValue', () => {