Reuse hooks when replaying a suspended component

packages/react-dom/src/__tests__/ReactDOMServerIntegrationHooks-test.js

@@ -430,26 +430,6 @@ describe('ReactDOMServerHooks', () => {
       expect(domNode.textContent).toEqual('hi');
     });
 
-    itThrowsWhenRendering(
-      'with a warning for useRef inside useReducer',
-      async render => {
-        function App() {
-          const [value, dispatch] = useReducer((state, action) => {
-            useRef(0);
-            return state + 1;
-          }, 0);
-          if (value === 0) {
-            dispatch();
-          }
-          return value;
-        }
-
-        const domNode = await render(<App />, 1);
-        expect(domNode.textContent).toEqual('1');
-      },
-      'Rendered more hooks than during the previous render',
-    );
-
     itRenders('with a warning for useRef inside useState', async render => {
       function App() {
         const [value] = useState(() => {
@@ -686,6 +666,32 @@ describe('ReactDOMServerHooks', () => {
     );
   });
 
+  describe('invalid hooks', () => {
+    it('warns when calling useRef inside useReducer', async () => {
+      function App() {
+        const [value, dispatch] = useReducer((state, action) => {
+          useRef(0);
+          return state + 1;
+        }, 0);
+        if (value === 0) {
+          dispatch();
+        }
+        return value;
+      }
+
+      let error;
+      try {
+        await serverRender(<App />);
+      } catch (x) {
+        error = x;
+      }
+      expect(error).not.toBe(undefined);
+      expect(error.message).toContain(
+        'Rendered more hooks than during the previous render',
+      );
+    });
+  });
+
   itRenders(
     'can use the same context multiple times in the same function',
     async render => {

packages/react-reconciler/src/ReactFiberBeginWork.new.js

@@ -38,6 +38,7 @@ import type {
 import type {UpdateQueue} from './ReactFiberClassUpdateQueue.new';
 import type {RootState} from './ReactFiberRoot.new';
 import type {TracingMarkerInstance} from './ReactFiberTracingMarkerComponent.new';
+
 import checkPropTypes from 'shared/checkPropTypes';
 import {
   markComponentRenderStarted,
@@ -203,6 +204,7 @@ import {
   renderWithHooks,
   checkDidRenderIdHook,
   bailoutHooks,
+  replaySuspendedComponentWithHooks,
 } from './ReactFiberHooks.new';
 import {stopProfilerTimerIfRunning} from './ReactProfilerTimer.new';
 import {
@@ -1159,6 +1161,54 @@ function updateFunctionComponent(
   return workInProgress.child;
 }
 
+export function replayFunctionComponent(
+  current: Fiber | null,
+  workInProgress: Fiber,
+  nextProps: any,
+  Component: any,
+  renderLanes: Lanes,
+): Fiber | null {
+  // This function is used to replay a component that previously suspended,
+  // after its data resolves. It's a simplified version of
+  // updateFunctionComponent that reuses the hooks from the previous attempt.
+
+  let context;
+  if (!disableLegacyContext) {
+    const unmaskedContext = getUnmaskedContext(workInProgress, Component, true);
+    context = getMaskedContext(workInProgress, unmaskedContext);
+  }
+
+  prepareToReadContext(workInProgress, renderLanes);
+  if (enableSchedulingProfiler) {
+    markComponentRenderStarted(workInProgress);
+  }
+  const nextChildren = replaySuspendedComponentWithHooks(
+    current,
+    workInProgress,
+    Component,
+    nextProps,
+    context,
+  );
+  const hasId = checkDidRenderIdHook();
+  if (enableSchedulingProfiler) {
+    markComponentRenderStopped();
+  }
+
+  if (current !== null && !didReceiveUpdate) {
+    bailoutHooks(current, workInProgress, renderLanes);
+    return bailoutOnAlreadyFinishedWork(current, workInProgress, renderLanes);
+  }
+
+  if (getIsHydrating() && hasId) {
+    pushMaterializedTreeId(workInProgress);
+  }
+
+  // React DevTools reads this flag.
+  workInProgress.flags |= PerformedWork;
+  reconcileChildren(current, workInProgress, nextChildren, renderLanes);
+  return workInProgress.child;
+}
+
 function updateClassComponent(
   current: Fiber | null,
   workInProgress: Fiber,

packages/react-reconciler/src/ReactFiberBeginWork.old.js

@@ -38,6 +38,7 @@ import type {
 import type {UpdateQueue} from './ReactFiberClassUpdateQueue.old';
 import type {RootState} from './ReactFiberRoot.old';
 import type {TracingMarkerInstance} from './ReactFiberTracingMarkerComponent.old';
+
 import checkPropTypes from 'shared/checkPropTypes';
 import {
   markComponentRenderStarted,
@@ -203,6 +204,7 @@ import {
   renderWithHooks,
   checkDidRenderIdHook,
   bailoutHooks,
+  replaySuspendedComponentWithHooks,
 } from './ReactFiberHooks.old';
 import {stopProfilerTimerIfRunning} from './ReactProfilerTimer.old';
 import {
@@ -1159,6 +1161,54 @@ function updateFunctionComponent(
   return workInProgress.child;
 }
 
+export function replayFunctionComponent(
+  current: Fiber | null,
+  workInProgress: Fiber,
+  nextProps: any,
+  Component: any,
+  renderLanes: Lanes,
+): Fiber | null {
+  // This function is used to replay a component that previously suspended,
+  // after its data resolves. It's a simplified version of
+  // updateFunctionComponent that reuses the hooks from the previous attempt.
+
+  let context;
+  if (!disableLegacyContext) {
+    const unmaskedContext = getUnmaskedContext(workInProgress, Component, true);
+    context = getMaskedContext(workInProgress, unmaskedContext);
+  }
+
+  prepareToReadContext(workInProgress, renderLanes);
+  if (enableSchedulingProfiler) {
+    markComponentRenderStarted(workInProgress);
+  }
+  const nextChildren = replaySuspendedComponentWithHooks(
+    current,
+    workInProgress,
+    Component,
+    nextProps,
+    context,
+  );
+  const hasId = checkDidRenderIdHook();
+  if (enableSchedulingProfiler) {
+    markComponentRenderStopped();
+  }
+
+  if (current !== null && !didReceiveUpdate) {
+    bailoutHooks(current, workInProgress, renderLanes);
+    return bailoutOnAlreadyFinishedWork(current, workInProgress, renderLanes);
+  }
+
+  if (getIsHydrating() && hasId) {
+    pushMaterializedTreeId(workInProgress);
+  }
+
+  // React DevTools reads this flag.
+  workInProgress.flags |= PerformedWork;
+  reconcileChildren(current, workInProgress, nextChildren, renderLanes);
+  return workInProgress.child;
+}
+
 function updateClassComponent(
   current: Fiber | null,
   workInProgress: Fiber,