Fix useMemoCache with setState in render

Fixes the bug that @alexmckenley and @mofeiZ found where setState-in-render can reset useMemoCache and cause an infinite loop. The bug was that renderWithHooksAgain() was not resetting hook state when rerendering (so useMemo values were preserved) but was resetting the updateQueue. This meant that the entire memo cache was cleared on a setState-in-render.

The fix here is to call a new helper function to clear the update queue. It nulls out other properties, but for memoCache it just sets the index back to zero.

ghstack-source-id: 4c6350a60286b92113ee5922645406fa94b7b34a
Pull Request resolved: #30889

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) {
+      resetFunctionComponentUpdateQueue(workInProgress.updateQueue);
+    }
 
     if (__DEV__) {
       // Also validate hook order for cascading updates.
@@ -1101,6 +1103,22 @@ if (enableUseMemoCacheHook) {
   };
 }
 
+function resetFunctionComponentUpdateQueue(
+  updateQueue: FunctionComponentUpdateQueue,
+): void {
+  updateQueue.lastEffect = null;
+  updateQueue.events = null;
+  updateQueue.stores = null;
+  if (enableUseMemoCacheHook) {
+    if (updateQueue.memoCache != null) {
+      // NOTE: this function intentionally does not reset memoCache data. 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.
+      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

@@ -16,7 +16,6 @@ let assertLog;
 let useMemo;
 let useState;
 let useMemoCache;
-let waitForThrow;
 let MemoCacheSentinel;
 let ErrorBoundary;
 
@@ -32,7 +31,6 @@ describe('useMemoCache()', () => {
     useMemo = React.useMemo;
     useMemoCache = require('react/compiler-runtime').c;
     useState = React.useState;
-    waitForThrow = require('internal-test-utils').waitForThrow;
     MemoCacheSentinel = Symbol.for('react.memo_cache_sentinel');
 
     class _ErrorBoundary extends React.Component {
@@ -667,7 +665,7 @@ describe('useMemoCache()', () => {
     }
 
     // Baseline / source code
-    function useUserMemo(value) {
+    function useManualMemo(value) {
       return useMemo(() => [value], [value]);
     }
 
@@ -683,24 +681,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>);
   });