Fix useMemoCache with setState in render

ghstack-source-id: b3e3b82
Pull Request resolved: #30889

Author: josephsavona

packages/react-reconciler/src/ReactFiberHooks.js

@@ -828,7 +828,9 @@ function renderWithHooksAgain<Props, SecondArg>(
     currentHook = null;
     workInProgressHook = null;
 
-    workInProgress.updateQueue = null;
+    if (workInProgress.updateQueue != null) {
+      clearFunctionComponentUpdateQueue(workInProgress.updateQueue);
+    }
 
     if (__DEV__) {
       // Also validate hook order for cascading updates.
@@ -1101,6 +1103,20 @@ if (enableUseMemoCacheHook) {
   };
 }
 
+// NOTE: this function intentionally does not reset memoCache. We reuse updateQueue for the memo
+// cache to avoid increasing the size of fibers that don't need a cache, but we don't want to reset
+// the cache when other properties are reset.
+const clearFunctionComponentUpdateQueue = (
+  updateQueue: FunctionComponentUpdateQueue,
+) => {
+  updateQueue.lastEffect = null;
+  updateQueue.events = null;
+  updateQueue.stores = null;
+  if (updateQueue.memoCache != null) {
+    updateQueue.memoCache.index = 0;
+  }
+};
+
 function useThenable<T>(thenable: Thenable<T>): T {
   // Track the position of the thenable within this fiber.
   const index = thenableIndexCounter;

packages/react-reconciler/src/__tests__/useMemoCache-test.js

@@ -667,7 +667,7 @@ describe('useMemoCache()', () => {
     }
 
     // Baseline / source code
-    function useUserMemo(value) {
+    function useManualMemo(value) {
       return useMemo(() => [value], [value]);
     }
 
@@ -683,24 +683,22 @@ describe('useMemoCache()', () => {
     }
 
     /**
-     * Test case: note that the initial render never completes
+     * Test with useMemoCache
      */
     let root = ReactNoop.createRoot();
-    const IncorrectInfiniteComponent = makeComponent(useCompilerMemo);
-    root.render(<IncorrectInfiniteComponent value={2} />);
-    await waitForThrow(
-      'Too many re-renders. React limits the number of renders to prevent ' +
-        'an infinite loop.',
-    );
+    const CompilerMemoComponent = makeComponent(useCompilerMemo);
+    await act(() => {
+      root.render(<CompilerMemoComponent value={2} />);
+    });
+    expect(root).toMatchRenderedOutput(<div>2</div>);
 
     /**
-     * Baseline test: initial render is expected to complete after a retry
-     * (triggered by the setState)
+     * Test with useMemo
      */
     root = ReactNoop.createRoot();
-    const CorrectComponent = makeComponent(useUserMemo);
+    const HookMemoComponent = makeComponent(useManualMemo);
     await act(() => {
-      root.render(<CorrectComponent value={2} />);
+      root.render(<HookMemoComponent value={2} />);
     });
     expect(root).toMatchRenderedOutput(<div>2</div>);
   });