Reuse hooks when replaying a suspended component

acdlite · acdlite · commit 33e3d2878e9e · 2022-11-17T14:48:42.000-05:00
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.
diff --git a/packages/react-reconciler/src/ReactFiberBeginWork.new.js b/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,
diff --git a/packages/react-reconciler/src/ReactFiberBeginWork.old.js b/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,
diff --git a/packages/react-reconciler/src/ReactFiberHooks.new.js b/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;
 }
 
diff --git a/packages/react-reconciler/src/ReactFiberHooks.old.js b/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;
 }
 
diff --git a/packages/react-reconciler/src/ReactFiberWorkLoop.new.js b/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
diff --git a/packages/react-reconciler/src/ReactFiberWorkLoop.old.js b/packages/react-reconciler/src/ReactFiberWorkLoop.old.js
diff --git a/packages/react-reconciler/src/__tests__/ReactThenable-test.js b/packages/react-reconciler/src/__tests__/ReactThenable-test.js
