Track ThenableState alongside other hooks

Now that hook state is preserved while the work loop is suspended, we don't need to track the thenable state in the work loop. We can track it alongside the rest of the hook state. This is a nice simplification and also aligns better with how it works in Fizz and Flight. The promises will still be cleared when the component finishes rendering (either complete or unwind). In the future, we could stash the promises on the fiber and reuse them during an update. However, this would only work for `use` calls that occur before an prop/state/context is processed, because `use` calls can only be assumed to execute in the same order if no other props/state/context have changed. So it might not be worth doing until we have finer grained memoization.
facebook · Nov 17, 2022 · f284d9f · f284d9f
1 parent 6b4c031
commit f284d9f
Show file tree

Hide file tree

Showing 9 changed files with 142 additions and 221 deletions.
diff --git a/packages/react-reconciler/src/ReactFiberBeginWork.new.js b/packages/react-reconciler/src/ReactFiberBeginWork.new.js
@@ -38,7 +38,6 @@ 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 {
@@ -1167,7 +1166,6 @@ export function replayFunctionComponent(
   workInProgress: Fiber,
   nextProps: any,
   Component: any,
-  prevThenableState: ThenableState,
   renderLanes: Lanes,
 ): Fiber | null {
   // This function is used to replay a component that previously suspended,
@@ -1190,7 +1188,6 @@ export function replayFunctionComponent(
     Component,
     nextProps,
     context,
-    prevThenableState,
   );
   const hasId = checkDidRenderIdHook();
   if (enableSchedulingProfiler) {

diff --git a/packages/react-reconciler/src/ReactFiberBeginWork.old.js b/packages/react-reconciler/src/ReactFiberBeginWork.old.js
@@ -38,7 +38,6 @@ 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 {
@@ -1167,7 +1166,6 @@ export function replayFunctionComponent(
   workInProgress: Fiber,
   nextProps: any,
   Component: any,
-  prevThenableState: ThenableState,
   renderLanes: Lanes,
 ): Fiber | null {
   // This function is used to replay a component that previously suspended,
@@ -1190,7 +1188,6 @@ export function replayFunctionComponent(
     Component,
     nextProps,
     context,
-    prevThenableState,
   );
   const hasId = checkDidRenderIdHook();
   if (enableSchedulingProfiler) {

diff --git a/packages/react-reconciler/src/ReactFiberHooks.new.js b/packages/react-reconciler/src/ReactFiberHooks.new.js
@@ -105,7 +105,6 @@ import {
   requestEventTime,
   markSkippedUpdateLanes,
   isInvalidExecutionContextForEventFunction,
-  getSuspendedThenableState,
 } from './ReactFiberWorkLoop.new';
 
 import getComponentNameFromFiber from 'react-reconciler/src/getComponentNameFromFiber';
@@ -141,9 +140,9 @@ import {
 import {getTreeId} from './ReactFiberTreeContext.new';
 import {now} from './Scheduler';
 import {
-  prepareThenableState,
   trackUsedThenable,
   checkIfUseWrappedInTryCatch,
+  createThenableState,
 } from './ReactFiberThenable.new';
 import type {ThenableState} from './ReactFiberThenable.new';
 
@@ -247,6 +246,7 @@ let shouldDoubleInvokeUserFnsInHooksDEV: boolean = false;
 let localIdCounter: number = 0;
 // Counts number of `use`-d thenables
 let thenableIndexCounter: number = 0;
+let thenableState: ThenableState | null = null;
 
 // Used for ids that are generated completely client-side (i.e. not during
 // hydration). This counter is global, so client ids are not stable across
@@ -449,6 +449,7 @@ export function renderWithHooks<Props, SecondArg>(
   // didScheduleRenderPhaseUpdate = false;
   // localIdCounter = 0;
   // thenableIndexCounter = 0;
+  // thenableState = null;
 
   // TODO Warn if no hooks are used at all during mount, then some are used during update.
   // Currently we will identify the update render as a mount because memoizedState === null.
@@ -477,10 +478,6 @@ export function renderWithHooks<Props, SecondArg>(
         : HooksDispatcherOnUpdate;
   }
 
-  // If this is a replay, restore the thenable state from the previous attempt.
-  const prevThenableState = getSuspendedThenableState();
-  prepareThenableState(prevThenableState);
-
   // In Strict Mode, during development, user functions are double invoked to
   // help detect side effects. The logic for how this is implemented for in
   // hook components is a bit complex so let's break it down.
@@ -525,7 +522,6 @@ export function renderWithHooks<Props, SecondArg>(
       Component,
       props,
       secondArg,
-      prevThenableState,
     );
   }
 
@@ -538,7 +534,6 @@ export function renderWithHooks<Props, SecondArg>(
         Component,
         props,
         secondArg,
-        prevThenableState,
       );
     } finally {
       setIsStrictModeForDevtools(false);
@@ -600,7 +595,9 @@ function finishRenderingHooks(current: Fiber | null, workInProgress: Fiber) {
   didScheduleRenderPhaseUpdate = false;
   // This is reset by checkDidRenderIdHook
   // localIdCounter = 0;
+
   thenableIndexCounter = 0;
+  thenableState = null;
 
   if (didRenderTooFewHooks) {
     throw new Error(
@@ -652,7 +649,6 @@ export function replaySuspendedComponentWithHooks<Props, SecondArg>(
   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.
@@ -676,7 +672,6 @@ export function replaySuspendedComponentWithHooks<Props, SecondArg>(
     Component,
     props,
     secondArg,
-    prevThenableState,
   );
   finishRenderingHooks(current, workInProgress);
   return children;
@@ -687,7 +682,6 @@ function renderWithHooksAgain<Props, SecondArg>(
   Component: (p: Props, arg: SecondArg) => any,
   props: Props,
   secondArg: SecondArg,
-  prevThenableState: ThenableState | null,
 ) {
   // This is used to perform another render pass. It's used when setState is
   // called during render, and for double invoking components in Strict Mode
@@ -735,7 +729,6 @@ function renderWithHooksAgain<Props, SecondArg>(
       ? HooksDispatcherOnRerenderInDEV
       : HooksDispatcherOnRerender;
 
-    prepareThenableState(prevThenableState);
     children = Component(props, secondArg);
   } while (didScheduleRenderPhaseUpdateDuringThisPass);
   return children;
@@ -821,6 +814,7 @@ export function resetHooksOnUnwind(): void {
   didScheduleRenderPhaseUpdateDuringThisPass = false;
   localIdCounter = 0;
   thenableIndexCounter = 0;
+  thenableState = null;
 }
 
 function mountWorkInProgressHook(): Hook {
@@ -954,7 +948,11 @@ function use<T>(usable: Usable<T>): T {
       // Track the position of the thenable within this fiber.
       const index = thenableIndexCounter;
       thenableIndexCounter += 1;
-      return trackUsedThenable(thenable, index);
+
+      if (thenableState === null) {
+        thenableState = createThenableState();
+      }
+      return trackUsedThenable(thenableState, thenable, index);
     } else if (
       usable.$$typeof === REACT_CONTEXT_TYPE ||
       usable.$$typeof === REACT_SERVER_CONTEXT_TYPE

diff --git a/packages/react-reconciler/src/ReactFiberHooks.old.js b/packages/react-reconciler/src/ReactFiberHooks.old.js
@@ -105,7 +105,6 @@ import {
   requestEventTime,
   markSkippedUpdateLanes,
   isInvalidExecutionContextForEventFunction,
-  getSuspendedThenableState,
 } from './ReactFiberWorkLoop.old';
 
 import getComponentNameFromFiber from 'react-reconciler/src/getComponentNameFromFiber';
@@ -141,9 +140,9 @@ import {
 import {getTreeId} from './ReactFiberTreeContext.old';
 import {now} from './Scheduler';
 import {
-  prepareThenableState,
   trackUsedThenable,
   checkIfUseWrappedInTryCatch,
+  createThenableState,
 } from './ReactFiberThenable.old';
 import type {ThenableState} from './ReactFiberThenable.old';
 
@@ -247,6 +246,7 @@ let shouldDoubleInvokeUserFnsInHooksDEV: boolean = false;
 let localIdCounter: number = 0;
 // Counts number of `use`-d thenables
 let thenableIndexCounter: number = 0;
+let thenableState: ThenableState | null = null;
 
 // Used for ids that are generated completely client-side (i.e. not during
 // hydration). This counter is global, so client ids are not stable across
@@ -449,6 +449,7 @@ export function renderWithHooks<Props, SecondArg>(
   // didScheduleRenderPhaseUpdate = false;
   // localIdCounter = 0;
   // thenableIndexCounter = 0;
+  // thenableState = null;
 
   // TODO Warn if no hooks are used at all during mount, then some are used during update.
   // Currently we will identify the update render as a mount because memoizedState === null.
@@ -477,10 +478,6 @@ export function renderWithHooks<Props, SecondArg>(
         : HooksDispatcherOnUpdate;
   }
 
-  // If this is a replay, restore the thenable state from the previous attempt.
-  const prevThenableState = getSuspendedThenableState();
-  prepareThenableState(prevThenableState);
-
   // In Strict Mode, during development, user functions are double invoked to
   // help detect side effects. The logic for how this is implemented for in
   // hook components is a bit complex so let's break it down.
@@ -525,7 +522,6 @@ export function renderWithHooks<Props, SecondArg>(
       Component,
       props,
       secondArg,
-      prevThenableState,
     );
   }
 
@@ -538,7 +534,6 @@ export function renderWithHooks<Props, SecondArg>(
         Component,
         props,
         secondArg,
-        prevThenableState,
       );
     } finally {
       setIsStrictModeForDevtools(false);
@@ -600,7 +595,9 @@ function finishRenderingHooks(current: Fiber | null, workInProgress: Fiber) {
   didScheduleRenderPhaseUpdate = false;
   // This is reset by checkDidRenderIdHook
   // localIdCounter = 0;
+
   thenableIndexCounter = 0;
+  thenableState = null;
 
   if (didRenderTooFewHooks) {
     throw new Error(
@@ -652,7 +649,6 @@ export function replaySuspendedComponentWithHooks<Props, SecondArg>(
   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.
@@ -676,7 +672,6 @@ export function replaySuspendedComponentWithHooks<Props, SecondArg>(
     Component,
     props,
     secondArg,
-    prevThenableState,
   );
   finishRenderingHooks(current, workInProgress);
   return children;
@@ -687,7 +682,6 @@ function renderWithHooksAgain<Props, SecondArg>(
   Component: (p: Props, arg: SecondArg) => any,
   props: Props,
   secondArg: SecondArg,
-  prevThenableState: ThenableState | null,
 ) {
   // This is used to perform another render pass. It's used when setState is
   // called during render, and for double invoking components in Strict Mode
@@ -735,7 +729,6 @@ function renderWithHooksAgain<Props, SecondArg>(
       ? HooksDispatcherOnRerenderInDEV
       : HooksDispatcherOnRerender;
 
-    prepareThenableState(prevThenableState);
     children = Component(props, secondArg);
   } while (didScheduleRenderPhaseUpdateDuringThisPass);
   return children;
@@ -821,6 +814,7 @@ export function resetHooksOnUnwind(): void {
   didScheduleRenderPhaseUpdateDuringThisPass = false;
   localIdCounter = 0;
   thenableIndexCounter = 0;
+  thenableState = null;
 }
 
 function mountWorkInProgressHook(): Hook {
@@ -954,7 +948,11 @@ function use<T>(usable: Usable<T>): T {
       // Track the position of the thenable within this fiber.
       const index = thenableIndexCounter;
       thenableIndexCounter += 1;
-      return trackUsedThenable(thenable, index);
+
+      if (thenableState === null) {
+        thenableState = createThenableState();
+      }
+      return trackUsedThenable(thenableState, thenable, index);
     } else if (
       usable.$$typeof === REACT_CONTEXT_TYPE ||
       usable.$$typeof === REACT_SERVER_CONTEXT_TYPE

diff --git a/packages/react-reconciler/src/ReactFiberThenable.new.js b/packages/react-reconciler/src/ReactFiberThenable.new.js
@@ -31,57 +31,40 @@ export const SuspenseException: mixed = new Error(
     "call the promise's `.catch` method and pass the result to `use`",
 );
 
-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.
-  const state = thenableState;
-  thenableState = null;
-  return state;
-}
-
 export function isThenableResolved(thenable: Thenable<mixed>): boolean {
   const status = thenable.status;
   return status === 'fulfilled' || status === 'rejected';
 }
 
 function noop(): void {}
 
-export function trackUsedThenable<T>(thenable: Thenable<T>, index: number): T {
+export function trackUsedThenable<T>(
+  thenableState: ThenableState,
+  thenable: Thenable<T>,
+  index: number,
+): T {
   if (__DEV__ && ReactCurrentActQueue.current !== null) {
     ReactCurrentActQueue.didUsePromise = true;
   }
 
-  if (thenableState === null) {
-    thenableState = [thenable];
+  const previous = thenableState[index];
+  if (previous === undefined) {
+    thenableState.push(thenable);
   } else {
-    const previous = thenableState[index];
-    if (previous === undefined) {
-      thenableState.push(thenable);
-    } else {
-      if (previous !== thenable) {
-        // Reuse the previous thenable, and drop the new one. We can assume
-        // they represent the same value, because components are idempotent.
-
-        // Avoid an unhandled rejection errors for the Promises that we'll
-        // intentionally ignore.
-        thenable.then(noop, noop);
-        thenable = previous;
-      }
+    if (previous !== thenable) {
+      // Reuse the previous thenable, and drop the new one. We can assume
+      // they represent the same value, because components are idempotent.
+
+      // Avoid an unhandled rejection errors for the Promises that we'll
+      // intentionally ignore.
+      thenable.then(noop, noop);
+      thenable = previous;
     }
   }