Track thenable state in work loop

This is a refactor to track the array of thenables that is preserved
across replays in the work loop instead of the Thenable module.

The reason is that I'm about to add additional state to the Thenable
module that is specific to a particular attempt — like the current
index — and is reset between replays. So it's helpful to keep the two
kinds of state separate so it's clearer which state gets reset when.

The array of thenables is not reset until the work-in-progress either
completes or unwinds.

This also makes the structure more similar to Fizz and Flight.

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,19 +17,49 @@ 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 suspendedThenableDidResolve(): boolean {
-  if (suspendedThenable !== null) {
-    const status = suspendedThenable.status;
+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 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) {
@@ -37,14 +67,12 @@ export function trackUsedThenable<T>(thenable: Thenable<T>, index: number) {
     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,19 +17,49 @@ 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 suspendedThenableDidResolve(): boolean {
-  if (suspendedThenable !== null) {
-    const status = suspendedThenable.status;
+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 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) {
@@ -37,14 +67,12 @@ export function trackUsedThenable<T>(thenable: Thenable<T>, index: number) {
     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,8 +1686,6 @@ function prepareFreshStack(root: FiberRoot, lanes: Lanes): Fiber {
       );
       interruptedWork = interruptedWork.return;
     }
-    resetWakeableStateAfterEachAttempt();
-    resetThenableStateOnCompletion();
   }
   workInProgressRoot = root;
   const rootWorkInProgress = createWorkInProgress(root.current, null);
@@ -1695,6 +1693,7 @@ function prepareFreshStack(root: FiberRoot, lanes: Lanes): Fiber {
   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.

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

@@ -22,6 +22,7 @@ import type {
   TransitionAbort,
 } from './ReactFiberTracingMarkerComponent.old';
 import type {OffscreenInstance} from './ReactFiberOffscreenComponent';
+import type {ThenableState} from './ReactFiberThenable.old';
 
 import {
   warnAboutDeprecatedLifecycles,
@@ -265,10 +266,8 @@ import {
 } from './ReactFiberAct.old';
 import {processTransitionCallbacks} from './ReactFiberTracingMarkerComponent.old';
 import {
-  resetWakeableStateAfterEachAttempt,
-  resetThenableStateOnCompletion,
-  suspendedThenableDidResolve,
-  isTrackingSuspendedThenable,
+  getThenableStateAfterSuspending,
+  isThenableStateResolved,
 } from './ReactFiberThenable.old';
 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,8 +1686,6 @@ function prepareFreshStack(root: FiberRoot, lanes: Lanes): Fiber {
       );
       interruptedWork = interruptedWork.return;
     }
-    resetWakeableStateAfterEachAttempt();
-    resetThenableStateOnCompletion();
   }
   workInProgressRoot = root;
   const rootWorkInProgress = createWorkInProgress(root.current, null);
@@ -1695,6 +1693,7 @@ function prepareFreshStack(root: FiberRoot, lanes: Lanes): Fiber {
   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.

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

@@ -8,7 +8,7 @@ let use;
 let Suspense;
 let startTransition;
 
-describe('ReactWakeable', () => {
+describe('ReactThenable', () => {
   beforeEach(() => {
     jest.resetModules();
 
@@ -243,6 +243,76 @@ describe('ReactWakeable', () => {
     expect(Scheduler).toHaveYielded(['Oops!', 'Oops!']);
   });
 
+  // @gate enableUseHook
+  test('use(promise) in multiple components', async () => {
+    // This tests that the state for tracking promises is reset per component.
+    const promiseA = Promise.resolve('A');
+    const promiseB = Promise.resolve('B');
+    const promiseC = Promise.resolve('C');
+    const promiseD = Promise.resolve('D');
+
+    function Child({prefix}) {
+      return <Text text={prefix + use(promiseC) + use(promiseD)} />;
+    }
+
+    function Parent() {
+      return <Child prefix={use(promiseA) + use(promiseB)} />;
+    }
+
+    function App() {
+      return (
+        <Suspense fallback={<Text text="Loading..." />}>
+          <Parent />
+        </Suspense>
+      );
+    }
+
+    const root = ReactNoop.createRoot();
+    await act(async () => {
+      startTransition(() => {
+        root.render(<App />);
+      });
+    });
+    expect(Scheduler).toHaveYielded(['ABCD']);
+    expect(root).toMatchRenderedOutput('ABCD');
+  });
+
+  // @gate enableUseHook
+  test('use(promise) in multiple sibling components', async () => {
+    // This tests that the state for tracking promises is reset per component.
+
+    const promiseA = {then: () => {}, status: 'pending', value: null};
+    const promiseB = {then: () => {}, status: 'pending', value: null};
+    const promiseC = {then: () => {}, status: 'fulfilled', value: 'C'};
+    const promiseD = {then: () => {}, status: 'fulfilled', value: 'D'};
+
+    function Sibling1({prefix}) {
+      return <Text text={use(promiseA) + use(promiseB)} />;
+    }
+
+    function Sibling2() {
+      return <Text text={use(promiseC) + use(promiseD)} />;
+    }
+
+    function App() {
+      return (
+        <Suspense fallback={<Text text="Loading..." />}>
+          <Sibling1 />
+          <Sibling2 />
+        </Suspense>
+      );
+    }
+
+    const root = ReactNoop.createRoot();
+    await act(async () => {
+      startTransition(() => {
+        root.render(<App />);
+      });
+    });
+    expect(Scheduler).toHaveYielded(['CD', 'Loading...']);
+    expect(root).toMatchRenderedOutput('Loading...');
+  });
+
   // @gate enableUseHook
   test('erroring in the same component as an uncached promise does not result in an infinite loop', async () => {
     class ErrorBoundary extends React.Component {