useTransition improvements

packages/react-reconciler/src/ReactFiberHooks.js

@@ -17,6 +17,7 @@ import type {Fiber} from './ReactFiber';
 import type {ExpirationTime} from './ReactFiberExpirationTime';
 import type {HookEffectTag} from './ReactHookEffectTags';
 import type {SuspenseConfig} from './ReactFiberSuspenseConfig';
+import type {TransitionInstance} from './ReactFiberTransition';
 import type {ReactPriorityLevel} from './SchedulerWithReactIntegration';
 
 import ReactSharedInternals from 'shared/ReactSharedInternals';
@@ -51,13 +52,14 @@ import warning from 'shared/warning';
 import getComponentName from 'shared/getComponentName';
 import is from 'shared/objectIs';
 import {markWorkInProgressReceivedUpdate} from './ReactFiberBeginWork';
+import {SuspendOnTask} from './ReactFiberThrow';
 import {requestCurrentSuspenseConfig} from './ReactFiberSuspenseConfig';
 import {
-  UserBlockingPriority,
-  NormalPriority,
-  runWithPriority,
-  getCurrentPriorityLevel,
-} from './SchedulerWithReactIntegration';
+  startTransition,
+  requestCurrentTransition,
+  cancelPendingTransition,
+} from './ReactFiberTransition';
+import {getCurrentPriorityLevel} from './SchedulerWithReactIntegration';
 
 const {ReactCurrentDispatcher, ReactCurrentBatchConfig} = ReactSharedInternals;
 
@@ -117,6 +119,7 @@ type Update<S, A> = {
 type UpdateQueue<S, A> = {
   last: Update<S, A> | null,
   dispatch: (A => mixed) | null,
+  pendingTransition: TransitionInstance | null,
   lastRenderedReducer: ((S, A) => S) | null,
   lastRenderedState: S | null,
 };
@@ -663,6 +666,7 @@ function mountReducer<S, I, A>(
   const queue = (hook.queue = {
     last: null,
     dispatch: null,
+    pendingTransition: null,
     lastRenderedReducer: reducer,
     lastRenderedState: (initialState: any),
   });
@@ -758,7 +762,10 @@ function updateReducer<S, I, A>(
     let prevUpdate = baseUpdate;
     let update = first;
     let didSkip = false;
+    let lastProcessedTransitionTime = NoWork;
+    let lastSkippedTransitionTime = NoWork;
     do {
+      const suspenseConfig = update.suspenseConfig;
       const updateExpirationTime = update.expirationTime;
       if (updateExpirationTime < renderExpirationTime) {
         // Priority is insufficient. Skip this update. If this is the first
@@ -774,19 +781,25 @@ function updateReducer<S, I, A>(
           remainingExpirationTime = updateExpirationTime;
           markUnprocessedUpdateTime(remainingExpirationTime);
         }
+
+        if (suspenseConfig !== null) {
+          // This update is part of a transition
+          if (
+            lastSkippedTransitionTime === NoWork ||
+            lastSkippedTransitionTime > updateExpirationTime
+          ) {
+            lastSkippedTransitionTime = updateExpirationTime;
+          }
+        }
       } else {
         // This update does have sufficient priority.
-
         // Mark the event time of this update as relevant to this render pass.
         // TODO: This should ideally use the true event time of this update rather than
         // its priority which is a derived and not reverseable value.
         // TODO: We should skip this update if it was already committed but currently
         // we have no way of detecting the difference between a committed and suspended
         // update here.
-        markRenderEventTimeAndConfig(
-          updateExpirationTime,
-          update.suspenseConfig,
-        );
+        markRenderEventTimeAndConfig(updateExpirationTime, suspenseConfig);
 
         // Process this update.
         if (update.eagerReducer === reducer) {
@@ -797,11 +810,32 @@ function updateReducer<S, I, A>(
           const action = update.action;
           newState = reducer(newState, action);
         }
+
+        if (suspenseConfig !== null) {
+          // This update is part of a transition
+          if (
+            lastProcessedTransitionTime === NoWork ||
+            lastProcessedTransitionTime > updateExpirationTime
+          ) {
+            lastProcessedTransitionTime = updateExpirationTime;
+          }
+        }
       }
+
       prevUpdate = update;
       update = update.next;
     } while (update !== null && update !== first);
 
+    if (
+      lastProcessedTransitionTime !== NoWork &&
+      lastSkippedTransitionTime !== NoWork
+    ) {
+      // There are multiple updates scheduled on this queue, but only some of
+      // them were processed. To avoid showing an intermediate state, abort
+      // the current render and restart at a level that includes them all.
+      throw new SuspendOnTask(lastSkippedTransitionTime);
+    }
+
     if (!didSkip) {
       newBaseUpdate = prevUpdate;
       newBaseState = newState;
@@ -835,6 +869,7 @@ function mountState<S>(
   const queue = (hook.queue = {
     last: null,
     dispatch: null,
+    pendingTransition: null,
     lastRenderedReducer: basicStateReducer,
     lastRenderedState: (initialState: any),
   });
@@ -1172,49 +1207,122 @@ function updateDeferredValue<T>(
   return prevValue;
 }
 
-function startTransition(setPending, config, callback) {
-  const priorityLevel = getCurrentPriorityLevel();
-  runWithPriority(
-    priorityLevel < UserBlockingPriority ? UserBlockingPriority : priorityLevel,
-    () => {
-      setPending(true);
-    },
-  );
-  runWithPriority(
-    priorityLevel > NormalPriority ? NormalPriority : priorityLevel,
-    () => {
-      const previousConfig = ReactCurrentBatchConfig.suspense;
-      ReactCurrentBatchConfig.suspense = config === undefined ? null : config;
-      try {
-        setPending(false);
-        callback();
-      } finally {
-        ReactCurrentBatchConfig.suspense = previousConfig;
-      }
-    },
-  );
-}
-
 function mountTransition(
   config: SuspenseConfig | void | null,
 ): [(() => void) => void, boolean] {
-  const [isPending, setPending] = mountState(false);
-  const start = mountCallback(startTransition.bind(null, setPending, config), [
-    setPending,
+  const hook = mountWorkInProgressHook();
+  const fiber = ((currentlyRenderingFiber: any): Fiber);
+  const instance: TransitionInstance = {
+    pendingExpirationTime: NoWork,
+    fiber,
+  };
+  // TODO: Intentionally storing this on the queue field to avoid adding a new/
+  // one; `queue` should be a union.
+  hook.queue = (instance: any);
+
+  const isPending = false;
+
+  // TODO: Consider passing `config` to `startTransition` instead of the hook.
+  // Then we don't have to recompute the callback whenever it changes. However,
+  // if we don't end up changing the API, we should at least optimize this
+  // to use the same hook instead of a separate hook just for the callback.
+  const start = mountCallback(startTransition.bind(null, instance, config), [
     config,
   ]);
+
+  const resolvedExpirationTime = NoWork;
+  hook.memoizedState = {
+    isPending,
+
+    // Represents the last processed expiration time.
+    resolvedExpirationTime,
+  };
+
   return [start, isPending];
 }
 
 function updateTransition(
   config: SuspenseConfig | void | null,
 ): [(() => void) => void, boolean] {
-  const [isPending, setPending] = updateState(false);
-  const start = updateCallback(startTransition.bind(null, setPending, config), [
-    setPending,
+  const hook = updateWorkInProgressHook();
+
+  const instance: TransitionInstance = (hook.queue: any);
+
+  const pendingExpirationTime = instance.pendingExpirationTime;
+  const oldState = hook.memoizedState;
+  const oldIsPending = oldState.isPending;
+  const oldResolvedExpirationTime = oldState.resolvedExpirationTime;
+
+  // Check if the most recent transition is pending. The following logic is
+  // a little confusing, but it conceptually maps to same logic used to process
+  // state update queues (see: updateReducer). We're cheating a bit because
+  // we know that there is only ever a single pending transition, and the last
+  // one always wins. So we don't need to maintain an actual queue of updates;
+  // we only need to track 1) which is the most recent pending level 2) did
+  // we already resolve
+  //
+  // Note: This could be even simpler if we used a commit effect to mark when a
+  // pending transition is resolved. The cleverness that follows is meant to
+  // avoid the overhead of an extra effect; however, if this ends up being *too*
+  // clever, an effect probably isn't that bad, since it would only fire once
+  // per transition.
+  let newIsPending;
+  let newResolvedExpirationTime;
+
+  if (pendingExpirationTime === NoWork) {
+    // There are no pending transitions. Reset all fields.
+    newIsPending = false;
+    newResolvedExpirationTime = NoWork;
+  } else {
+    // There is a pending transition. It may or may not have resolved. Compare
+    // the time at which we last resolved to the pending time. If the pending
+    // time is in the future, then we're still pending.
+    if (
+      oldResolvedExpirationTime === NoWork ||
+      oldResolvedExpirationTime > pendingExpirationTime
+    ) {
+      // We have not already resolved at the pending time. Check if this render
+      // includes the pending level.
+      if (renderExpirationTime <= pendingExpirationTime) {
+        // This render does include the pending level. Mark it as resolved.
+        newIsPending = false;
+        newResolvedExpirationTime = renderExpirationTime;
+      } else {
+        // This render does not include the pending level. Still pending.
+        newIsPending = true;
+        newResolvedExpirationTime = oldResolvedExpirationTime;
+      }
+    } else {
+      // Already resolved at this expiration time.
+      newIsPending = false;
+      newResolvedExpirationTime = oldResolvedExpirationTime;
+    }
+  }
+
+  if (newIsPending !== oldIsPending) {
+    markWorkInProgressReceivedUpdate();
+  } else if (oldIsPending === false) {
+    // This is a trick to mutate the instance without a commit effect. If
+    // neither the current nor work-in-progress hook are pending, and there's no
+    // pending transition at a lower priority (which we know because there can
+    // only be one pending level per useTransition hook), then we can be certain
+    // there are no pending transitions even if this render does not finish.
+    // It's similar to the trick we use for eager setState bailouts. Like that
+    // optimization, this should have no semantic effect.
+    instance.pendingExpirationTime = NoWork;
+    newResolvedExpirationTime = NoWork;
+  }
+
+  hook.memoizedState = {
+    isPending: newIsPending,
+    resolvedExpirationTime: newResolvedExpirationTime,
+  };
+
+  const start = updateCallback(startTransition.bind(null, instance, config), [
     config,
   ]);
-  return [start, isPending];
+
+  return [start, newIsPending];
 }
 
 function dispatchAction<S, A>(
@@ -1274,6 +1382,7 @@ function dispatchAction<S, A>(
   } else {
     const currentTime = requestCurrentTimeForUpdate();
     const suspenseConfig = requestCurrentSuspenseConfig();
+    const transition = requestCurrentTransition();
     const expirationTime = computeExpirationForFiber(
       currentTime,
       fiber,
@@ -1354,6 +1463,20 @@ function dispatchAction<S, A>(
         warnIfNotCurrentlyActingUpdatesInDev(fiber);
       }
     }
+
+    if (transition !== null) {
+      const prevPendingTransition = queue.pendingTransition;
+      if (transition !== prevPendingTransition) {
+        queue.pendingTransition = transition;
+        if (prevPendingTransition !== null) {
+          // There's already a pending transition on this queue. The new
+          // transition supersedes the old one. Turn of the `isPending` state
+          // of the previous transition.
+          cancelPendingTransition(prevPendingTransition);
+        }
+      }
+    }
+
     scheduleWork(fiber, expirationTime);
   }
 }

packages/react-reconciler/src/ReactFiberSuspenseConfig.js

@@ -9,6 +9,7 @@
 
 import ReactSharedInternals from 'shared/ReactSharedInternals';
 
+// TODO: Remove React.unstable_withSuspenseConfig and move this to the renderer
 const {ReactCurrentBatchConfig} = ReactSharedInternals;
 
 export type SuspenseConfig = {|

packages/react-reconciler/src/ReactFiberThrow.js

@@ -62,6 +62,12 @@ import {Sync} from './ReactFiberExpirationTime';
 
 const PossiblyWeakMap = typeof WeakMap === 'function' ? WeakMap : Map;
 
+// Throw an object with this type to abort the current render and restart at
+// a different level.
+export function SuspendOnTask(expirationTime: ExpirationTime) {
+  this.expirationTime = expirationTime;
+}
+
 function createRootErrorUpdate(
   fiber: Fiber,
   errorInfo: CapturedValue<mixed>,