Reuse hooks when replaying a suspended component

When a component suspends, under some conditions, we can wait for the
data to resolve and replay the component without unwinding the stack or
showing a fallback in the interim. When we do this, we reuse the
promises that were unwrapped during the previous attempts, so that if
they aren't memoized, the result can still be used.

We should do the same for all hooks. That way, if you _do_ memoize an
async function call with useMemo, it won't be called again during the
replay. This effectively gives you a local version of the functionality
provided by `cache`, using the normal memoization patterns that have
long existed in React.

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

@@ -38,6 +38,8 @@ import type {
 import type {UpdateQueue} from './ReactFiberClassUpdateQueue.new';
 import type {RootState} from './ReactFiberRoot.new';
 import type {TracingMarkerInstance} from './ReactFiberTracingMarkerComponent.new';
+import type {ThenableState} from './ReactFiberThenable.new';
+
 import checkPropTypes from 'shared/checkPropTypes';
 import {
   markComponentRenderStarted,
@@ -203,6 +205,7 @@ import {
   renderWithHooks,
   checkDidRenderIdHook,
   bailoutHooks,
+  replaySuspendedComponentWithHooks,
 } from './ReactFiberHooks.new';
 import {stopProfilerTimerIfRunning} from './ReactProfilerTimer.new';
 import {
@@ -1159,6 +1162,56 @@ function updateFunctionComponent(
   return workInProgress.child;
 }
 
+export function replayFunctionComponent(
+  current: Fiber | null,
+  workInProgress: Fiber,
+  nextProps: any,
+  Component: any,
+  prevThenableState: ThenableState,
+  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,
+    prevThenableState,
+  );
+  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,8 @@ import type {
 import type {UpdateQueue} from './ReactFiberClassUpdateQueue.old';
 import type {RootState} from './ReactFiberRoot.old';
 import type {TracingMarkerInstance} from './ReactFiberTracingMarkerComponent.old';
+import type {ThenableState} from './ReactFiberThenable.old';
+
 import checkPropTypes from 'shared/checkPropTypes';
 import {
   markComponentRenderStarted,
@@ -203,6 +205,7 @@ import {
   renderWithHooks,
   checkDidRenderIdHook,
   bailoutHooks,
+  replaySuspendedComponentWithHooks,
 } from './ReactFiberHooks.old';
 import {stopProfilerTimerIfRunning} from './ReactProfilerTimer.old';
 import {
@@ -1159,6 +1162,56 @@ function updateFunctionComponent(
   return workInProgress.child;
 }
 
+export function replayFunctionComponent(
+  current: Fiber | null,
+  workInProgress: Fiber,
+  nextProps: any,
+  Component: any,
+  prevThenableState: ThenableState,
+  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,
+    prevThenableState,
+  );
+  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/ReactFiberHooks.new.js

@@ -545,6 +545,12 @@ export function renderWithHooks<Props, SecondArg>(
     }
   }
 
+  finishRenderingHooks(current, workInProgress);
+
+  return children;
+}
+
+function finishRenderingHooks(current: Fiber | null, workInProgress: Fiber) {
   // We can assume the previous dispatcher is always this one, since we set it
   // at the beginning of the render phase and there's no re-entrance.
   ReactCurrentDispatcher.current = ContextOnlyDispatcher;
@@ -638,7 +644,41 @@ export function renderWithHooks<Props, SecondArg>(
       }
     }
   }
+}
 
+export function replaySuspendedComponentWithHooks<Props, SecondArg>(
+  current: Fiber | null,
+  workInProgress: Fiber,
+  Component: (p: Props, arg: SecondArg) => any,
+  props: Props,
+  secondArg: SecondArg,
+  prevThenableState: ThenableState | null,
+): any {
+  // This function is used to replay a component that previously suspended,
+  // after its data resolves.
+  //
+  // It's a simplified version of renderWithHooks, but it doesn't need to do
+  // most of the set up work because they weren't reset when we suspended; they
+  // only get reset when the component either completes (finishRenderingHooks)
+  // or unwinds (resetHooksOnUnwind).
+  if (__DEV__) {
+    hookTypesDev =
+      current !== null
+        ? ((current._debugHookTypes: any): Array<HookType>)
+        : null;
+    hookTypesUpdateIndexDev = -1;
+    // Used for hot reloading:
+    ignorePreviousDependencies =
+      current !== null && current.type !== workInProgress.type;
+  }
+  const children = renderWithHooksAgain(
+    workInProgress,
+    Component,
+    props,
+    secondArg,
+    prevThenableState,
+  );
+  finishRenderingHooks(current, workInProgress);
   return children;
 }
 

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

@@ -545,6 +545,12 @@ export function renderWithHooks<Props, SecondArg>(
     }
   }
 
+  finishRenderingHooks(current, workInProgress);
+
+  return children;
+}
+
+function finishRenderingHooks(current: Fiber | null, workInProgress: Fiber) {
   // We can assume the previous dispatcher is always this one, since we set it
   // at the beginning of the render phase and there's no re-entrance.
   ReactCurrentDispatcher.current = ContextOnlyDispatcher;
@@ -638,7 +644,41 @@ export function renderWithHooks<Props, SecondArg>(
       }
     }
   }
+}
 
+export function replaySuspendedComponentWithHooks<Props, SecondArg>(
+  current: Fiber | null,
+  workInProgress: Fiber,
+  Component: (p: Props, arg: SecondArg) => any,
+  props: Props,
+  secondArg: SecondArg,
+  prevThenableState: ThenableState | null,
+): any {
+  // This function is used to replay a component that previously suspended,
+  // after its data resolves.
+  //
+  // It's a simplified version of renderWithHooks, but it doesn't need to do
+  // most of the set up work because they weren't reset when we suspended; they
+  // only get reset when the component either completes (finishRenderingHooks)
+  // or unwinds (resetHooksOnUnwind).
+  if (__DEV__) {
+    hookTypesDev =
+      current !== null
+        ? ((current._debugHookTypes: any): Array<HookType>)
+        : null;
+    hookTypesUpdateIndexDev = -1;
+    // Used for hot reloading:
+    ignorePreviousDependencies =
+      current !== null && current.type !== workInProgress.type;
+  }
+  const children = renderWithHooksAgain(
+    workInProgress,
+    Component,
+    props,
+    secondArg,
+    prevThenableState,
+  );
+  finishRenderingHooks(current, workInProgress);
   return children;
 }
 

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

@@ -181,6 +181,7 @@ import {requestCurrentTransition, NoTransition} from './ReactFiberTransition';
 import {
   SelectiveHydrationException,
   beginWork as originalBeginWork,
+  replayFunctionComponent,
 } from './ReactFiberBeginWork.new';
 import {completeWork} from './ReactFiberCompleteWork.new';
 import {unwindWork, unwindInterruptedWork} from './ReactFiberUnwindWork.new';
@@ -282,6 +283,7 @@ import {
   getSuspenseHandler,
   isBadSuspenseFallback,
 } from './ReactFiberSuspenseContext.new';
+import {resolveDefaultProps} from './ReactFiberLazyComponent.new';
 
 const ceil = Math.ceil;
 
@@ -2353,22 +2355,79 @@ function replaySuspendedUnitOfWork(
   // This is a fork of performUnitOfWork specifcally for replaying a fiber that
   // just suspended.
   //
-  // Instead of unwinding the stack and potentially showing a fallback, unwind
-  // only the last stack frame, reset the fiber, and try rendering it again.
   const current = unitOfWork.alternate;
-  resetSuspendedWorkLoopOnUnwind();
-  unwindInterruptedWork(current, unitOfWork, workInProgressRootRenderLanes);
-  unitOfWork = workInProgress = resetWorkInProgress(unitOfWork, renderLanes);
-
   setCurrentDebugFiberInDEV(unitOfWork);
 
   let next;
-  if (enableProfilerTimer && (unitOfWork.mode & ProfileMode) !== NoMode) {
+  setCurrentDebugFiberInDEV(unitOfWork);
+  const isProfilingMode =
+    enableProfilerTimer && (unitOfWork.mode & ProfileMode) !== NoMode;
+  if (isProfilingMode) {
     startProfilerTimer(unitOfWork);
-    next = beginWork(current, unitOfWork, renderLanes);
+  }
+  switch (unitOfWork.tag) {
+    case IndeterminateComponent: {
+      // Because it suspended with `use`, we can assume it's a
+      // function component.
+      unitOfWork.tag = FunctionComponent;
+      // Fallthrough to the next branch.
+    }
+    // eslint-disable-next-line no-fallthrough
+    case FunctionComponent:
+    case ForwardRef: {
+      // Resolve `defaultProps`. This logic is copied from `beginWork`.
+      // TODO: Consider moving this switch statement into that module. Also,
+      // could maybe use this as an opportunity to say `use` doesn't work with
+      // `defaultProps` :)
+      const Component = unitOfWork.type;
+      const unresolvedProps = unitOfWork.pendingProps;
+      const resolvedProps =
+        unitOfWork.elementType === Component
+          ? unresolvedProps
+          : resolveDefaultProps(Component, unresolvedProps);
+      next = replayFunctionComponent(
+        current,
+        unitOfWork,
+        resolvedProps,
+        Component,
+        thenableState,
+        workInProgressRootRenderLanes,
+      );
+      break;
+    }
+    case SimpleMemoComponent: {
+      const Component = unitOfWork.type;
+      const nextProps = unitOfWork.pendingProps;
+      next = replayFunctionComponent(
+        current,
+        unitOfWork,
+        nextProps,
+        Component,
+        thenableState,
+        workInProgressRootRenderLanes,
+      );
+      break;
+    }
+    default: {
+      if (__DEV__) {
+        console.error(
+          'Unexpected type of work: %s, Currently only function ' +
+            'components are replayed after suspending. This is a bug in React.',
+          unitOfWork.tag,
+        );
+      }
+      resetSuspendedWorkLoopOnUnwind();
+      unwindInterruptedWork(current, unitOfWork, workInProgressRootRenderLanes);
+      unitOfWork = workInProgress = resetWorkInProgress(
+        unitOfWork,
+        renderLanes,
+      );
+      next = beginWork(current, unitOfWork, renderLanes);
+      break;
+    }
+  }
+  if (isProfilingMode) {
     stopProfilerTimerIfRunningAndRecordDelta(unitOfWork, true);
-  } else {
-    next = beginWork(current, unitOfWork, renderLanes);
   }
 
   // The begin phase finished successfully without suspending. Reset the state

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

@@ -181,6 +181,7 @@ import {requestCurrentTransition, NoTransition} from './ReactFiberTransition';
 import {
   SelectiveHydrationException,
   beginWork as originalBeginWork,
+  replayFunctionComponent,
 } from './ReactFiberBeginWork.old';
 import {completeWork} from './ReactFiberCompleteWork.old';
 import {unwindWork, unwindInterruptedWork} from './ReactFiberUnwindWork.old';
@@ -282,6 +283,7 @@ import {
   getSuspenseHandler,
   isBadSuspenseFallback,
 } from './ReactFiberSuspenseContext.old';
+import {resolveDefaultProps} from './ReactFiberLazyComponent.old';
 
 const ceil = Math.ceil;
 
@@ -2353,22 +2355,79 @@ function replaySuspendedUnitOfWork(
   // This is a fork of performUnitOfWork specifcally for replaying a fiber that
   // just suspended.
   //
-  // Instead of unwinding the stack and potentially showing a fallback, unwind
-  // only the last stack frame, reset the fiber, and try rendering it again.
   const current = unitOfWork.alternate;
-  resetSuspendedWorkLoopOnUnwind();
-  unwindInterruptedWork(current, unitOfWork, workInProgressRootRenderLanes);
-  unitOfWork = workInProgress = resetWorkInProgress(unitOfWork, renderLanes);
-
   setCurrentDebugFiberInDEV(unitOfWork);
 
   let next;
-  if (enableProfilerTimer && (unitOfWork.mode & ProfileMode) !== NoMode) {
+  setCurrentDebugFiberInDEV(unitOfWork);
+  const isProfilingMode =
+    enableProfilerTimer && (unitOfWork.mode & ProfileMode) !== NoMode;
+  if (isProfilingMode) {
     startProfilerTimer(unitOfWork);
-    next = beginWork(current, unitOfWork, renderLanes);
+  }
+  switch (unitOfWork.tag) {
+    case IndeterminateComponent: {
+      // Because it suspended with `use`, we can assume it's a
+      // function component.
+      unitOfWork.tag = FunctionComponent;
+      // Fallthrough to the next branch.
+    }
+    // eslint-disable-next-line no-fallthrough
+    case FunctionComponent:
+    case ForwardRef: {
+      // Resolve `defaultProps`. This logic is copied from `beginWork`.
+      // TODO: Consider moving this switch statement into that module. Also,
+      // could maybe use this as an opportunity to say `use` doesn't work with
+      // `defaultProps` :)
+      const Component = unitOfWork.type;
+      const unresolvedProps = unitOfWork.pendingProps;
+      const resolvedProps =
+        unitOfWork.elementType === Component
+          ? unresolvedProps
+          : resolveDefaultProps(Component, unresolvedProps);
+      next = replayFunctionComponent(
+        current,
+        unitOfWork,
+        resolvedProps,
+        Component,
+        thenableState,
+        workInProgressRootRenderLanes,
+      );
+      break;
+    }
+    case SimpleMemoComponent: {
+      const Component = unitOfWork.type;
+      const nextProps = unitOfWork.pendingProps;
+      next = replayFunctionComponent(
+        current,
+        unitOfWork,
+        nextProps,
+        Component,
+        thenableState,
+        workInProgressRootRenderLanes,
+      );
+      break;
+    }
+    default: {
+      if (__DEV__) {
+        console.error(
+          'Unexpected type of work: %s, Currently only function ' +
+            'components are replayed after suspending. This is a bug in React.',
+          unitOfWork.tag,
+        );
+      }
+      resetSuspendedWorkLoopOnUnwind();
+      unwindInterruptedWork(current, unitOfWork, workInProgressRootRenderLanes);
+      unitOfWork = workInProgress = resetWorkInProgress(
+        unitOfWork,
+        renderLanes,
+      );
+      next = beginWork(current, unitOfWork, renderLanes);
+      break;
+    }
+  }
+  if (isProfilingMode) {
     stopProfilerTimerIfRunningAndRecordDelta(unitOfWork, true);
-  } else {
-    next = beginWork(current, unitOfWork, renderLanes);
   }
 
   // The begin phase finished successfully without suspending. Reset the state

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

@@ -6,6 +6,7 @@ let Scheduler;
 let act;
 let use;
 let useState;
+let useMemo;
 let Suspense;
 let startTransition;
 let pendingTextRequests;
@@ -20,6 +21,7 @@ describe('ReactThenable', () => {
     act = require('jest-react').act;
     use = React.use;
     useState = React.useState;
+    useMemo = React.useMemo;
     Suspense = React.Suspense;
     startTransition = React.startTransition;
 
@@ -789,4 +791,89 @@ describe('ReactThenable', () => {
     ]);
     expect(root).toMatchRenderedOutput('(empty)');
   });
+
+  test('when replaying a suspended component, reuses the hooks computed during the previous attempt', async () => {
+    function ExcitingText({text}) {
+      // This computes the uppercased version of some text. Pretend it's an
+      // expensive operation that we want to reuse.
+      const uppercaseText = useMemo(() => {
+        Scheduler.unstable_yieldValue('Compute uppercase: ' + text);
+        return text.toUpperCase();
+      }, [text]);
+
+      // This adds an exclamation point to the text. Pretend it's an async
+      // operation that is sent to a service for processing.
+      const exclamatoryText = use(getAsyncText(uppercaseText + '!'));
+
+      // This surrounds the text with sparkle emojis. The purpose in this test
+      // is to show that you can suspend in the middle of a sequence of hooks
+      // without breaking anything.
+      const sparklingText = useMemo(() => {
+        Scheduler.unstable_yieldValue('Add sparkles: ' + exclamatoryText);
+        return `✨ ${exclamatoryText} ✨`;
+      }, [exclamatoryText]);
+
+      return <Text text={sparklingText} />;
+    }
+
+    const root = ReactNoop.createRoot();
+    await act(async () => {
+      startTransition(() => {
+        root.render(<ExcitingText text="Hello" />);
+      });
+    });
+    // Suspends while we wait for the async service to respond.
+    expect(Scheduler).toHaveYielded([
+      'Compute uppercase: Hello',
+      'Async text requested [HELLO!]',
+    ]);
+    expect(root).toMatchRenderedOutput(null);
+
+    // The data is received.
+    await act(async () => {
+      resolveTextRequests('HELLO!');
+    });
+    expect(Scheduler).toHaveYielded([
+      // We shouldn't run the uppercase computation again, because we can reuse
+      // the computation from the previous attempt.
+      // 'Compute uppercase: Hello',
+
+      'Async text requested [HELLO!]',
+      'Add sparkles: HELLO!',
+      '✨ HELLO! ✨',
+    ]);
+  });
+
+  // @gate enableUseHook
+  test(
+    'wrap an async function with useMemo to skip running the function ' +
+      'twice when loading new data',
+    async () => {
+      function App({text}) {
+        const promiseForText = useMemo(async () => getAsyncText(text), [text]);
+        const asyncText = use(promiseForText);
+        return <Text text={asyncText} />;
+      }
+
+      const root = ReactNoop.createRoot();
+      await act(async () => {
+        startTransition(() => {
+          root.render(<App text="Hello" />);
+        });
+      });
+      expect(Scheduler).toHaveYielded(['Async text requested [Hello]']);
+      expect(root).toMatchRenderedOutput(null);
+
+      await act(async () => {
+        resolveTextRequests('Hello');
+      });
+      expect(Scheduler).toHaveYielded([
+        // We shouldn't request async text again, because the async function
+        // was memoized
+        // 'Async text requested [Hello]'
+
+        'Hello',
+      ]);
+    },
+  );
 });