Fix: Synchronous popstate transitions (#30759)

This is a refactor of the fix in #27505.

When a transition update is scheduled by a popstate event, (i.e. a back/
forward navigation) we attempt to render it synchronously even though
it's a transition, since it's likely the previous page's data is cached.

In #27505, I changed the implementation so that it only "upgrades" the
priority of the transition for a single attempt. If the attempt
suspends, say because the data is not cached after all, from then on it
should be treated as a normal transition.

But it turns out #27505 did not work as intended, because it relied on
marking the root with pending synchronous work (root.pendingLanes),
which was never cleared until the popstate update completed.

The test scenarios I wrote accidentally worked for a different reason
related to suspending the work loop, which I'm currently in the middle
of refactoring.

packages/react-dom/src/__tests__/ReactDOMFiberAsync-test.js

@@ -744,7 +744,7 @@ describe('ReactDOMFiberAsync', () => {
       // Because it suspended, it remains on the current path
       expect(div.textContent).toBe('/path/a');
     });
-    assertLog(['Suspend! [/path/b]']);
+    assertLog([]);
 
     await act(async () => {
       resolvePromise();

packages/react-reconciler/src/ReactFiberLane.js

@@ -92,6 +92,14 @@ export const DeferredLane: Lane = /*                    */ 0b1000000000000000000
 export const UpdateLanes: Lanes =
   SyncLane | InputContinuousLane | DefaultLane | TransitionLanes;
 
+export const HydrationLanes =
+  SyncHydrationLane |
+  InputContinuousHydrationLane |
+  DefaultHydrationLane |
+  TransitionHydrationLane |
+  SelectiveHydrationLane |
+  IdleHydrationLane;
+
 // This function is used for the experimental timeline (react-devtools-timeline)
 // It should be kept in sync with the Lanes values above.
 export function getLabelForLane(lane: Lane): string | void {
@@ -282,6 +290,51 @@ export function getNextLanes(root: FiberRoot, wipLanes: Lanes): Lanes {
   return nextLanes;
 }
 
+export function getNextLanesToFlushSync(
+  root: FiberRoot,
+  extraLanesToForceSync: Lane | Lanes,
+): Lanes {
+  // Similar to getNextLanes, except instead of choosing the next lanes to work
+  // on based on their priority, it selects all the lanes that have equal or
+  // higher priority than those are given. That way they can be synchronously
+  // rendered in a single batch.
+  //
+  // The main use case is updates scheduled by popstate events, which are
+  // flushed synchronously even though they are transitions.
+  const lanesToFlush = SyncUpdateLanes | extraLanesToForceSync;
+
+  // Early bailout if there's no pending work left.
+  const pendingLanes = root.pendingLanes;
+  if (pendingLanes === NoLanes) {
+    return NoLanes;
+  }
+
+  const suspendedLanes = root.suspendedLanes;
+  const pingedLanes = root.pingedLanes;
+
+  // Remove lanes that are suspended (but not pinged)
+  const unblockedLanes = pendingLanes & ~(suspendedLanes & ~pingedLanes);
+  const unblockedLanesWithMatchingPriority =
+    unblockedLanes & getLanesOfEqualOrHigherPriority(lanesToFlush);
+
+  // If there are matching hydration lanes, we should do those by themselves.
+  // Hydration lanes must never include updates.
+  if (unblockedLanesWithMatchingPriority & HydrationLanes) {
+    return (
+      (unblockedLanesWithMatchingPriority & HydrationLanes) | SyncHydrationLane
+    );
+  }
+
+  if (unblockedLanesWithMatchingPriority) {
+    // Always include the SyncLane as part of the result, even if there's no
+    // pending sync work, to indicate the priority of the entire batch of work
+    // is considered Sync.
+    return unblockedLanesWithMatchingPriority | SyncLane;
+  }
+
+  return NoLanes;
+}
+
 export function getEntangledLanes(root: FiberRoot, renderLanes: Lanes): Lanes {
   let entangledLanes = renderLanes;
 
@@ -534,6 +587,14 @@ export function getHighestPriorityLane(lanes: Lanes): Lane {
   return lanes & -lanes;
 }
 
+function getLanesOfEqualOrHigherPriority(lanes: Lane | Lanes): Lanes {
+  // Create a mask with all bits to the right or same as the highest bit.
+  // So if lanes is 0b100, the result would be 0b111.
+  // If lanes is 0b101, the result would be 0b111.
+  const lowestPriorityLaneIndex = 31 - clz32(lanes);
+  return (1 << (lowestPriorityLaneIndex + 1)) - 1;
+}
+
 export function pickArbitraryLane(lanes: Lanes): Lane {
   // This wrapper function gets inlined. Only exists so to communicate that it
   // doesn't matter which bit is selected; you can pick any bit without
@@ -757,17 +818,6 @@ export function markRootEntangled(root: FiberRoot, entangledLanes: Lanes) {
   }
 }
 
-export function upgradePendingLaneToSync(root: FiberRoot, lane: Lane) {
-  // Since we're upgrading the priority of the given lane, there is now pending
-  // sync work.
-  root.pendingLanes |= SyncLane;
-
-  // Entangle the sync lane with the lane we're upgrading. This means SyncLane
-  // will not be allowed to finish without also finishing the given lane.
-  root.entangledLanes |= SyncLane;
-  root.entanglements[SyncLaneIndex] |= lane;
-}
-
 export function upgradePendingLanesToSync(
   root: FiberRoot,
   lanesToUpgrade: Lanes,

packages/react-reconciler/src/ReactFiberRootScheduler.js

@@ -8,7 +8,7 @@
  */
 
 import type {FiberRoot} from './ReactInternalTypes';
-import type {Lane} from './ReactFiberLane';
+import type {Lane, Lanes} from './ReactFiberLane';
 import type {PriorityLevel} from 'scheduler/src/SchedulerPriorities';
 import type {BatchConfigTransition} from './ReactFiberTracingMarkerComponent';
 
@@ -24,8 +24,8 @@ import {
   getNextLanes,
   includesSyncLane,
   markStarvedLanesAsExpired,
-  upgradePendingLaneToSync,
   claimNextTransitionLane,
+  getNextLanesToFlushSync,
 } from './ReactFiberLane';
 import {
   CommitContext,
@@ -145,18 +145,21 @@ export function ensureRootIsScheduled(root: FiberRoot): void {
 export function flushSyncWorkOnAllRoots() {
   // This is allowed to be called synchronously, but the caller should check
   // the execution context first.
-  flushSyncWorkAcrossRoots_impl(false);
+  flushSyncWorkAcrossRoots_impl(NoLanes, false);
 }
 
 export function flushSyncWorkOnLegacyRootsOnly() {
   // This is allowed to be called synchronously, but the caller should check
   // the execution context first.
   if (!disableLegacyMode) {
-    flushSyncWorkAcrossRoots_impl(true);
+    flushSyncWorkAcrossRoots_impl(NoLanes, true);
   }
 }
 
-function flushSyncWorkAcrossRoots_impl(onlyLegacy: boolean) {
+function flushSyncWorkAcrossRoots_impl(
+  syncTransitionLanes: Lanes | Lane,
+  onlyLegacy: boolean,
+) {
   if (isFlushingWork) {
     // Prevent reentrancy.
     // TODO: Is this overly defensive? The callers must check the execution
@@ -179,17 +182,28 @@ function flushSyncWorkAcrossRoots_impl(onlyLegacy: boolean) {
       if (onlyLegacy && (disableLegacyMode || root.tag !== LegacyRoot)) {
         // Skip non-legacy roots.
       } else {
-        const workInProgressRoot = getWorkInProgressRoot();
-        const workInProgressRootRenderLanes =
-          getWorkInProgressRootRenderLanes();
-        const nextLanes = getNextLanes(
-          root,
-          root === workInProgressRoot ? workInProgressRootRenderLanes : NoLanes,
-        );
-        if (includesSyncLane(nextLanes)) {
-          // This root has pending sync work. Flush it now.
-          didPerformSomeWork = true;
-          performSyncWorkOnRoot(root, nextLanes);
+        if (syncTransitionLanes !== NoLanes) {
+          const nextLanes = getNextLanesToFlushSync(root, syncTransitionLanes);
+          if (nextLanes !== NoLanes) {
+            // This root has pending sync work. Flush it now.
+            didPerformSomeWork = true;
+            performSyncWorkOnRoot(root, nextLanes);
+          }
+        } else {
+          const workInProgressRoot = getWorkInProgressRoot();
+          const workInProgressRootRenderLanes =
+            getWorkInProgressRootRenderLanes();
+          const nextLanes = getNextLanes(
+            root,
+            root === workInProgressRoot
+              ? workInProgressRootRenderLanes
+              : NoLanes,
+          );
+          if (includesSyncLane(nextLanes)) {
+            // This root has pending sync work. Flush it now.
+            didPerformSomeWork = true;
+            performSyncWorkOnRoot(root, nextLanes);
+          }
         }
       }
       root = root.next;
@@ -209,23 +223,23 @@ function processRootScheduleInMicrotask() {
   // We'll recompute this as we iterate through all the roots and schedule them.
   mightHavePendingSyncWork = false;
 
+  let syncTransitionLanes = NoLanes;
+  if (currentEventTransitionLane !== NoLane) {
+    if (shouldAttemptEagerTransition()) {
+      // A transition was scheduled during an event, but we're going to try to
+      // render it synchronously anyway. We do this during a popstate event to
+      // preserve the scroll position of the previous page.
+      syncTransitionLanes = currentEventTransitionLane;
+    }
+    currentEventTransitionLane = NoLane;
+  }
+
   const currentTime = now();
 
   let prev = null;
   let root = firstScheduledRoot;
   while (root !== null) {
     const next = root.next;
-
-    if (
-      currentEventTransitionLane !== NoLane &&
-      shouldAttemptEagerTransition()
-    ) {
-      // A transition was scheduled during an event, but we're going to try to
-      // render it synchronously anyway. We do this during a popstate event to
-      // preserve the scroll position of the previous page.
-      upgradePendingLaneToSync(root, currentEventTransitionLane);
-    }
-
     const nextLanes = scheduleTaskForRootDuringMicrotask(root, currentTime);
     if (nextLanes === NoLane) {
       // This root has no more pending work. Remove it from the schedule. To
@@ -248,18 +262,27 @@ function processRootScheduleInMicrotask() {
     } else {
       // This root still has work. Keep it in the list.
       prev = root;
-      if (includesSyncLane(nextLanes)) {
+
+      // This is a fast-path optimization to early exit from
+      // flushSyncWorkOnAllRoots if we can be certain that there is no remaining
+      // synchronous work to perform. Set this to true if there might be sync
+      // work left.
+      if (
+        // Skip the optimization if syncTransitionLanes is set
+        syncTransitionLanes !== NoLanes ||
+        // Common case: we're not treating any extra lanes as synchronous, so we
+        // can just check if the next lanes are sync.
+        includesSyncLane(nextLanes)
+      ) {
         mightHavePendingSyncWork = true;
       }
     }
     root = next;
   }
 
-  currentEventTransitionLane = NoLane;
-
   // At the end of the microtask, flush any pending synchronous work. This has
   // to come at the end, because it does actual rendering work that might throw.
-  flushSyncWorkOnAllRoots();
+  flushSyncWorkAcrossRoots_impl(syncTransitionLanes, false);
 }
 
 function scheduleTaskForRootDuringMicrotask(