Track thenable state in work loop

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

@@ -102,6 +102,7 @@ import {
   requestEventTime,
   markSkippedUpdateLanes,
   isInvalidExecutionContextForEventFunction,
+  getSuspendedThenableState,
 } from './ReactFiberWorkLoop.new';
 
 import getComponentNameFromFiber from 'react-reconciler/src/getComponentNameFromFiber';
@@ -134,6 +135,7 @@ import {
 import {getTreeId} from './ReactFiberTreeContext.new';
 import {now} from './Scheduler';
 import {
+  prepareThenableState,
   trackUsedThenable,
   getPreviouslyUsedThenableAtIndex,
 } from './ReactFiberThenable.new';
@@ -465,6 +467,9 @@ export function renderWithHooks<Props, SecondArg>(
         : HooksDispatcherOnUpdate;
   }
 
+  // If this is a replay, restore the thenable state from the previous attempt.
+  const prevThenableState = getSuspendedThenableState();
+  prepareThenableState(prevThenableState);
   let children = Component(props, secondArg);
 
   // Check if there was a render phase update
@@ -506,6 +511,7 @@ export function renderWithHooks<Props, SecondArg>(
         ? HooksDispatcherOnRerenderInDEV
         : HooksDispatcherOnRerender;
 
+      prepareThenableState(prevThenableState);
       children = Component(props, secondArg);
     } while (didScheduleRenderPhaseUpdateDuringThisPass);
   }

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

@@ -102,6 +102,7 @@ import {
   requestEventTime,
   markSkippedUpdateLanes,
   isInvalidExecutionContextForEventFunction,
+  getSuspendedThenableState,
 } from './ReactFiberWorkLoop.old';
 
 import getComponentNameFromFiber from 'react-reconciler/src/getComponentNameFromFiber';
@@ -134,6 +135,7 @@ import {
 import {getTreeId} from './ReactFiberTreeContext.old';
 import {now} from './Scheduler';
 import {
+  prepareThenableState,
   trackUsedThenable,
   getPreviouslyUsedThenableAtIndex,
 } from './ReactFiberThenable.old';
@@ -465,6 +467,9 @@ export function renderWithHooks<Props, SecondArg>(
         : HooksDispatcherOnUpdate;
   }
 
+  // If this is a replay, restore the thenable state from the previous attempt.
+  const prevThenableState = getSuspendedThenableState();
+  prepareThenableState(prevThenableState);
   let children = Component(props, secondArg);
 
   // Check if there was a render phase update
@@ -506,6 +511,7 @@ export function renderWithHooks<Props, SecondArg>(
         ? HooksDispatcherOnRerenderInDEV
         : HooksDispatcherOnRerender;
 
+      prepareThenableState(prevThenableState);
       children = Component(props, secondArg);
     } while (didScheduleRenderPhaseUpdateDuringThisPass);
   }

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

@@ -17,34 +17,62 @@ import type {
 import ReactSharedInternals from 'shared/ReactSharedInternals';
 const {ReactCurrentActQueue} = ReactSharedInternals;
 
-let suspendedThenable: Thenable<any> | null = null;
-let usedThenables: Array<Thenable<any> | void> | null = null;
+// TODO: Sparse arrays are bad for performance.
+export opaque type ThenableState = Array<Thenable<any> | void>;
 
-export function isTrackingSuspendedThenable(): boolean {
-  return suspendedThenable !== null;
+let thenableState: ThenableState | null = null;
+
+export function createThenableState(): ThenableState {
+  // The ThenableState is created the first time a component suspends. If it
+  // suspends again, we'll reuse the same state.
+  return [];
+}
+
+export function prepareThenableState(prevThenableState: ThenableState | null) {
+  // This function is called before every function that might suspend
+  // with `use`. Right now, that's only Hooks, but in the future we'll use the
+  // same mechanism for unwrapping promises during reconciliation.
+  thenableState = prevThenableState;
+}
+
+export function getThenableStateAfterSuspending(): ThenableState | null {
+  // Called by the work loop so it can stash the thenable state. It will use
+  // the state to replay the component when the promise resolves.
+  if (
+    thenableState !== null &&
+    // If we only `use`-ed resolved promises, then there is no suspended state
+    // TODO: The only reason we do this is to distinguish between throwing a
+    // promise (old Suspense pattern) versus `use`-ing one. A better solution is
+    // for `use` to throw a special, opaque value instead of a promise.
+    !isThenableStateResolved(thenableState)
+  ) {
+    const state = thenableState;
+    thenableState = null;
+    return state;
+  }
+  return null;
 }
 
-export function suspendedThenableDidResolve(): boolean {
-  if (suspendedThenable !== null) {
-    const status = suspendedThenable.status;
+export function isThenableStateResolved(thenables: ThenableState): boolean {
+  const lastThenable = thenables[thenables.length - 1];
+  if (lastThenable !== undefined) {
+    const status = lastThenable.status;
     return status === 'fulfilled' || status === 'rejected';
   }
-  return false;
+  return true;
 }
 
 export function trackUsedThenable<T>(thenable: Thenable<T>, index: number) {
   if (__DEV__ && ReactCurrentActQueue.current !== null) {
     ReactCurrentActQueue.didUsePromise = true;
   }
 
-  if (usedThenables === null) {
-    usedThenables = [thenable];
+  if (thenableState === null) {
+    thenableState = [thenable];
   } else {
-    usedThenables[index] = thenable;
+    thenableState[index] = thenable;
   }
 
-  suspendedThenable = thenable;
-
   // We use an expando to track the status and result of a thenable so that we
   // can synchronously unwrap the value. Think of this as an extension of the
   // Promise API, or a custom interface that is a superset of Thenable.
@@ -59,7 +87,6 @@ export function trackUsedThenable<T>(thenable: Thenable<T>, index: number) {
       // this thenable, because if we keep trying it will likely infinite loop
       // without ever resolving.
       // TODO: Log a warning?
-      suspendedThenable = null;
       break;
     default: {
       if (typeof thenable.status === 'string') {
@@ -91,19 +118,11 @@ export function trackUsedThenable<T>(thenable: Thenable<T>, index: number) {
   }
 }
 
-export function resetWakeableStateAfterEachAttempt() {
-  suspendedThenable = null;
-}
-
-export function resetThenableStateOnCompletion() {
-  usedThenables = null;
-}
-
 export function getPreviouslyUsedThenableAtIndex<T>(
   index: number,
 ): Thenable<T> | null {
-  if (usedThenables !== null) {
-    const thenable = usedThenables[index];
+  if (thenableState !== null) {
+    const thenable = thenableState[index];
     if (thenable !== undefined) {
       return thenable;
     }

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

@@ -17,34 +17,62 @@ import type {
 import ReactSharedInternals from 'shared/ReactSharedInternals';
 const {ReactCurrentActQueue} = ReactSharedInternals;
 
-let suspendedThenable: Thenable<any> | null = null;
-let usedThenables: Array<Thenable<any> | void> | null = null;
+// TODO: Sparse arrays are bad for performance.
+export opaque type ThenableState = Array<Thenable<any> | void>;
 
-export function isTrackingSuspendedThenable(): boolean {
-  return suspendedThenable !== null;
+let thenableState: ThenableState | null = null;
+
+export function createThenableState(): ThenableState {
+  // The ThenableState is created the first time a component suspends. If it
+  // suspends again, we'll reuse the same state.
+  return [];
+}
+
+export function prepareThenableState(prevThenableState: ThenableState | null) {
+  // This function is called before every function that might suspend
+  // with `use`. Right now, that's only Hooks, but in the future we'll use the
+  // same mechanism for unwrapping promises during reconciliation.
+  thenableState = prevThenableState;
+}
+
+export function getThenableStateAfterSuspending(): ThenableState | null {
+  // Called by the work loop so it can stash the thenable state. It will use
+  // the state to replay the component when the promise resolves.
+  if (
+    thenableState !== null &&
+    // If we only `use`-ed resolved promises, then there is no suspended state
+    // TODO: The only reason we do this is to distinguish between throwing a
+    // promise (old Suspense pattern) versus `use`-ing one. A better solution is
+    // for `use` to throw a special, opaque value instead of a promise.
+    !isThenableStateResolved(thenableState)
+  ) {
+    const state = thenableState;
+    thenableState = null;
+    return state;
+  }
+  return null;
 }
 
-export function suspendedThenableDidResolve(): boolean {
-  if (suspendedThenable !== null) {
-    const status = suspendedThenable.status;
+export function isThenableStateResolved(thenables: ThenableState): boolean {
+  const lastThenable = thenables[thenables.length - 1];
+  if (lastThenable !== undefined) {
+    const status = lastThenable.status;
     return status === 'fulfilled' || status === 'rejected';
   }
-  return false;
+  return true;
 }
 
 export function trackUsedThenable<T>(thenable: Thenable<T>, index: number) {
   if (__DEV__ && ReactCurrentActQueue.current !== null) {
     ReactCurrentActQueue.didUsePromise = true;
   }
 
-  if (usedThenables === null) {
-    usedThenables = [thenable];
+  if (thenableState === null) {
+    thenableState = [thenable];
   } else {
-    usedThenables[index] = thenable;
+    thenableState[index] = thenable;
   }
 
-  suspendedThenable = thenable;
-
   // We use an expando to track the status and result of a thenable so that we
   // can synchronously unwrap the value. Think of this as an extension of the
   // Promise API, or a custom interface that is a superset of Thenable.
@@ -59,7 +87,6 @@ export function trackUsedThenable<T>(thenable: Thenable<T>, index: number) {
       // this thenable, because if we keep trying it will likely infinite loop
       // without ever resolving.
       // TODO: Log a warning?
-      suspendedThenable = null;
       break;
     default: {
       if (typeof thenable.status === 'string') {
@@ -91,19 +118,11 @@ export function trackUsedThenable<T>(thenable: Thenable<T>, index: number) {
   }
 }
 
-export function resetWakeableStateAfterEachAttempt() {
-  suspendedThenable = null;
-}
-
-export function resetThenableStateOnCompletion() {
-  usedThenables = null;
-}
-
 export function getPreviouslyUsedThenableAtIndex<T>(
   index: number,
 ): Thenable<T> | null {
-  if (usedThenables !== null) {
-    const thenable = usedThenables[index];
+  if (thenableState !== null) {
+    const thenable = thenableState[index];
     if (thenable !== undefined) {
       return thenable;
     }

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

@@ -22,6 +22,7 @@ import type {
   TransitionAbort,
 } from './ReactFiberTracingMarkerComponent.new';
 import type {OffscreenInstance} from './ReactFiberOffscreenComponent';
+import type {ThenableState} from './ReactFiberThenable.new';
 
 import {
   warnAboutDeprecatedLifecycles,
@@ -265,10 +266,8 @@ import {
 } from './ReactFiberAct.new';
 import {processTransitionCallbacks} from './ReactFiberTracingMarkerComponent.new';
 import {
-  resetWakeableStateAfterEachAttempt,
-  resetThenableStateOnCompletion,
-  suspendedThenableDidResolve,
-  isTrackingSuspendedThenable,
+  getThenableStateAfterSuspending,
+  isThenableStateResolved,
 } from './ReactFiberThenable.new';
 import {schedulePostPaintCallback} from './ReactPostPaintCallback';
 
@@ -315,6 +314,7 @@ let workInProgressRootRenderLanes: Lanes = NoLanes;
 // immediately instead of unwinding the stack.
 let workInProgressIsSuspended: boolean = false;
 let workInProgressThrownValue: mixed = null;
+let workInProgressSuspendedThenableState: ThenableState | null = null;
 
 // Whether a ping listener was attached during this render. This is slightly
 // different that whether something suspended, because we don't add multiple
@@ -1686,15 +1686,14 @@ function prepareFreshStack(root: FiberRoot, lanes: Lanes): Fiber {
       );
       interruptedWork = interruptedWork.return;
     }
-    resetWakeableStateAfterEachAttempt();
-    resetThenableStateOnCompletion();
   }
   workInProgressRoot = root;
   const rootWorkInProgress = createWorkInProgress(root.current, null);
   workInProgress = rootWorkInProgress;
   workInProgressRootRenderLanes = renderLanes = lanes;
   workInProgressIsSuspended = false;
   workInProgressThrownValue = null;
+  workInProgressSuspendedThenableState = null;
   workInProgressRootDidAttachPingListener = false;
   workInProgressRootExitStatus = RootInProgress;
   workInProgressRootFatalError = null;
@@ -1729,6 +1728,7 @@ function handleThrow(root, thrownValue): void {
   // as suspending the execution of the work loop.
   workInProgressIsSuspended = true;
   workInProgressThrownValue = thrownValue;
+  workInProgressSuspendedThenableState = getThenableStateAfterSuspending();
 
   const erroredWork = workInProgress;
   if (erroredWork === null) {
@@ -2014,7 +2014,7 @@ function renderRootConcurrent(root: FiberRoot, lanes: Lanes) {
       break;
     } catch (thrownValue) {
       handleThrow(root, thrownValue);
-      if (isTrackingSuspendedThenable()) {
+      if (workInProgressSuspendedThenableState !== null) {
         // If this fiber just suspended, it's possible the data is already
         // cached. Yield to the main thread to give it a chance to ping. If
         // it does, we can retry immediately without unwinding the stack.
@@ -2117,13 +2117,14 @@ function resumeSuspendedUnitOfWork(
   // instead of unwinding the stack. It's a separate function to keep the
   // additional logic out of the work loop's hot path.
 
-  const wasPinged = suspendedThenableDidResolve();
-  resetWakeableStateAfterEachAttempt();
+  const wasPinged =
+    workInProgressSuspendedThenableState !== null &&
+    isThenableStateResolved(workInProgressSuspendedThenableState);
 
   if (!wasPinged) {
     // The thenable wasn't pinged. Return to the normal work loop. This will
     // unwind the stack, and potentially result in showing a fallback.
-    resetThenableStateOnCompletion();
+    workInProgressSuspendedThenableState = null;
 
     const returnFiber = unitOfWork.return;
     if (returnFiber === null || workInProgressRoot === null) {
@@ -2188,7 +2189,7 @@ function resumeSuspendedUnitOfWork(
   // The begin phase finished successfully without suspending. Reset the state
   // used to track the fiber while it was suspended. Then return to the normal
   // work loop.
-  resetThenableStateOnCompletion();
+  workInProgressSuspendedThenableState = null;
 
   resetCurrentDebugFiberInDEV();
   unitOfWork.memoizedProps = unitOfWork.pendingProps;
@@ -2202,6 +2203,10 @@ function resumeSuspendedUnitOfWork(
   ReactCurrentOwner.current = null;
 }
 
+export function getSuspendedThenableState(): ThenableState | null {
+  return workInProgressSuspendedThenableState;
+}
+
 function completeUnitOfWork(unitOfWork: Fiber): void {
   // Attempt to complete the current unit of work, then move to the next
   // sibling. If there are no more siblings, return to the parent fiber.