Fix useMemoCache with setState in render

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 {
@@ -630,18 +628,9 @@ describe('useMemoCache()', () => {
 
     // Finish loading the data.
     await act(() => updatedChunkB.resolve('B2'));
-    if (gate(flags => flags.enableNoCloningMemoCache)) {
-      // We did not have process the first chunk again. We reused the
-      // computation from the earlier attempt.
-      assertLog([]);
-    } else {
-      // Because we clone/reset the memo cache after every aborted attempt, we
-      // must process the first chunk again.
-      //
-      // That's three total times we've processed the first chunk, compared to
-      // just once when enableNoCloningMemoCache is on.
-      assertLog(['Some expensive processing... [A2]']);
-    }
+    // We did not have process the first chunk again. We reused the
+    // computation from the earlier attempt.
+    assertLog([]);
     expect(root).toMatchRenderedOutput(
       <>
         <div>Input: hi!</div>
@@ -667,7 +656,7 @@ describe('useMemoCache()', () => {
     }
 
     // Baseline / source code
-    function useUserMemo(value) {
+    function useManualMemo(value) {
       return useMemo(() => [value], [value]);
     }
 
@@ -683,24 +672,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>);
   });