Synchronously render the transition lane on popstate

Author: tyao1

packages/react-art/src/ReactARTHostConfig.js

@@ -346,6 +346,10 @@ export function getCurrentEventPriority() {
   return DefaultEventPriority;
 }
 
+export function getIsCurrentEventPopState() {
+  return false;
+}
+
 // The ART renderer is secondary to the React DOM renderer.
 export const isPrimaryRenderer = false;
 

packages/react-dom-bindings/src/client/ReactDOMHostConfig.js

@@ -388,6 +388,10 @@ export function getCurrentEventPriority(): EventPriority {
   return getEventPriority(currentEvent.type);
 }
 
+export function getIsCurrentEventPopState(): boolean {
+  return window.event && window.event.type === 'popstate';
+}
+
 export const isPrimaryRenderer = true;
 export const warnsIfNotActing = true;
 // This initialization code may run even on server environments

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

@@ -658,4 +658,84 @@ describe('ReactDOMFiberAsync', () => {
 
     expect(container.textContent).toBe('new');
   });
+
+  it('should synchronously render the transition lane in a popState', async () => {
+    function App() {
+      const [syncState, setSyncState] = React.useState(false);
+      const [hasNavigated, setHasNavigated] = React.useState(false);
+      function onPopstate() {
+        Scheduler.unstable_yieldValue(`popState`);
+        setSyncState(true);
+        React.startTransition(() => {
+          setHasNavigated(true);
+        });
+      }
+      React.useEffect(() => {
+        window.addEventListener('popstate', onPopstate);
+        return () => window.removeEventListener('popstate', onPopstate);
+      }, []);
+      Scheduler.unstable_yieldValue(`render:${hasNavigated}/${syncState}`);
+      return null;
+    }
+    const root = ReactDOMClient.createRoot(container);
+    await act(async () => {
+      root.render(<App />);
+    });
+    expect(Scheduler).toHaveYielded(['render:false/false']);
+
+    await act(async () => {
+      const popStateEvent = new Event('popstate');
+      window.dispatchEvent(popStateEvent);
+    });
+    expect(Scheduler).toHaveYielded(['popState', 'render:true/true']);
+
+    root.unmount();
+  });
+
+  it('transition lane in popState should still yield if it suspends', async () => {
+    const never = {then() {}};
+    let _setText;
+
+    function App() {
+      const [shouldSuspend, setShouldSuspend] = React.useState(false);
+      const [text, setText] = React.useState('0');
+      _setText = setText;
+      if (shouldSuspend) {
+        Scheduler.unstable_yieldValue('Suspend!');
+        throw never;
+      }
+      function onPopstate() {
+        Scheduler.unstable_yieldValue(`popState`);
+        React.startTransition(() => {
+          setShouldSuspend(val => !val);
+        });
+      }
+      React.useEffect(() => {
+        window.addEventListener('popstate', onPopstate);
+        return () => window.removeEventListener('popstate', onPopstate);
+      }, []);
+      Scheduler.unstable_yieldValue(`Child:${shouldSuspend}/${text}`);
+      return null;
+    }
+
+    const root = ReactDOMClient.createRoot(container);
+    await act(async () => {
+      root.render(<App />);
+    });
+    expect(Scheduler).toHaveYielded(['Child:false/0']);
+
+    await act(() => {
+      window.dispatchEvent(new Event('popstate'));
+    });
+    expect(Scheduler).toHaveYielded(['popState', 'Suspend!']);
+
+    await act(async () => {
+      _setText(true);
+    });
+
+    // Transition lane yields
+    expect(Scheduler).toHaveYielded(['Child:false/true', 'Suspend!']);
+
+    root.unmount();
+  });
 });

packages/react-native-renderer/src/ReactFabricHostConfig.js

@@ -506,6 +506,10 @@ export function getCurrentEventPriority(): * {
   return DefaultEventPriority;
 }
 
+export function getIsCurrentEventPopState(): boolean {
+  return false;
+}
+
 // The Fabric renderer is secondary to the existing React Native renderer.
 export const isPrimaryRenderer = false;
 

packages/react-native-renderer/src/ReactNativeHostConfig.js

@@ -260,6 +260,10 @@ export function getCurrentEventPriority(): * {
   return DefaultEventPriority;
 }
 
+export function getIsCurrentEventPopState(): boolean {
+  return false;
+}
+
 // -------------------
 //     Mutation
 // -------------------

packages/react-noop-renderer/src/createReactNoop.js

@@ -442,6 +442,10 @@ function createReactNoop(reconciler: Function, useMutation: boolean) {
       return currentEventPriority;
     },
 
+    getIsCurrentEventPopState(): boolean {
+      return false;
+    },
+
     now: Scheduler.unstable_now,
 
     isPrimaryRenderer: true,

packages/react-reconciler/src/ReactFiberLane.js

@@ -223,6 +223,10 @@ export function getNextLanes(root: FiberRoot, wipLanes: Lanes): Lanes {
         nextLanes = getHighestPriorityLanes(nonIdlePingedLanes);
       }
     }
+    // Batch transitions in `popstate` with the sync lane
+    if (root.forceSync && isTransitionLane(nonIdleUnblockedLanes)) {
+      nextLanes |= nonIdleUnblockedLanes & TransitionLanes;
+    }
   } else {
     // The only remaining work is Idle.
     const unblockedLanes = pendingLanes & ~suspendedLanes;
@@ -315,7 +319,6 @@ export function getNextLanes(root: FiberRoot, wipLanes: Lanes): Lanes {
       lanes &= ~lane;
     }
   }
-
   return nextLanes;
 }
 
@@ -471,8 +474,11 @@ export function getLanesToRetrySynchronouslyOnError(
   return NoLanes;
 }
 
-export function includesSyncLane(lanes: Lanes): boolean {
-  return (lanes & (SyncLane | SyncHydrationLane)) !== NoLanes;
+export function includesSyncLaneOrForceSync(
+  lanes: Lanes,
+  root: FiberRoot,
+): boolean {
+  return (lanes & (SyncLane | SyncHydrationLane)) !== NoLanes || root.forceSync;
 }
 
 export function includesNonIdleWork(lanes: Lanes): boolean {
@@ -633,6 +639,7 @@ export function markRootUpdated(
 export function markRootSuspended(root: FiberRoot, suspendedLanes: Lanes) {
   root.suspendedLanes |= suspendedLanes;
   root.pingedLanes &= ~suspendedLanes;
+  root.forceSync = false;
 
   // The suspended lanes are no longer CPU-bound. Clear their expiration times.
   const expirationTimes = root.expirationTimes;
@@ -671,6 +678,8 @@ export function markRootFinished(root: FiberRoot, remainingLanes: Lanes) {
 
   root.errorRecoveryDisabledLanes &= remainingLanes;
 
+  root.forceSync = false;
+
   const entanglements = root.entanglements;
   const eventTimes = root.eventTimes;
   const expirationTimes = root.expirationTimes;

packages/react-reconciler/src/ReactFiberWorkLoop.js

@@ -84,6 +84,7 @@ import {
   scheduleMicrotask,
   prepareRendererToRender,
   resetRendererAfterRender,
+  getIsCurrentEventPopState,
 } from './ReactFiberHostConfig';
 
 import {
@@ -137,7 +138,7 @@ import {
   NoTimestamp,
   claimNextTransitionLane,
   claimNextRetryLane,
-  includesSyncLane,
+  includesSyncLaneOrForceSync,
   isSubsetOfLanes,
   mergeLanes,
   removeLanes,
@@ -147,6 +148,7 @@ import {
   includesOnlyTransitions,
   includesBlockingLane,
   includesExpiredLane,
+  isTransitionLane,
   getNextLanes,
   markStarvedLanesAsExpired,
   getLanesToRetrySynchronouslyOnError,
@@ -727,6 +729,7 @@ export function scheduleUpdateOnFiber(
   }
 
   // Mark that the root has a pending update.
+  markRootWithPopState(root, lane);
   markRootUpdated(root, lane, eventTime);
 
   if (
@@ -914,7 +917,7 @@ function ensureRootIsScheduled(root: FiberRoot, currentTime: number) {
       // TODO: Temporary until we confirm this warning is not fired.
       if (
         existingCallbackNode == null &&
-        !includesSyncLane(existingCallbackPriority)
+        !includesSyncLaneOrForceSync(existingCallbackPriority, root)
       ) {
         console.error(
           'Expected scheduled callback to exist. This error is likely caused by a bug in React. Please file an issue.',
@@ -932,7 +935,7 @@ function ensureRootIsScheduled(root: FiberRoot, currentTime: number) {
 
   // Schedule a new callback.
   let newCallbackNode;
-  if (includesSyncLane(newCallbackPriority)) {
+  if (includesSyncLaneOrForceSync(newCallbackPriority, root)) {
     // Special case: Sync React callbacks are scheduled on a special
     // internal queue
     if (root.tag === LegacyRoot) {
@@ -1455,6 +1458,12 @@ function markRootSuspended(root: FiberRoot, suspendedLanes: Lanes) {
   markRootSuspended_dontCallThisOneDirectly(root, suspendedLanes);
 }
 
+function markRootWithPopState(root: FiberRoot, updateLane: Lane) {
+  if (isTransitionLane(updateLane) && getIsCurrentEventPopState()) {
+    root.forceSync = true;
+  }
+}
+
 // This is the entry point for synchronous tasks that don't go
 // through Scheduler
 function performSyncWorkOnRoot(root: FiberRoot) {
@@ -1469,7 +1478,7 @@ function performSyncWorkOnRoot(root: FiberRoot) {
   flushPassiveEffects();
 
   let lanes = getNextLanes(root, NoLanes);
-  if (!includesSyncLane(lanes)) {
+  if (!includesSyncLaneOrForceSync(lanes, root)) {
     // There's no remaining sync work left.
     ensureRootIsScheduled(root, now());
     return null;
@@ -2920,13 +2929,16 @@ function commitRootImpl(
   // TODO: We can optimize this by not scheduling the callback earlier. Since we
   // currently schedule the callback in multiple places, will wait until those
   // are consolidated.
-  if (includesSyncLane(pendingPassiveEffectsLanes) && root.tag !== LegacyRoot) {
+  if (
+    includesSyncLaneOrForceSync(pendingPassiveEffectsLanes, root) &&
+    root.tag !== LegacyRoot
+  ) {
     flushPassiveEffects();
   }
 
   // Read this again, since a passive effect might have updated it
   remainingLanes = root.pendingLanes;
-  if (includesSyncLane(remainingLanes)) {
+  if (includesSyncLaneOrForceSync(remainingLanes, root)) {
     if (enableProfilerTimer && enableProfilerNestedUpdatePhase) {
       markNestedUpdateScheduled();
     }

packages/react-reconciler/src/ReactInternalTypes.js

@@ -259,6 +259,8 @@ type BaseFiberRootProperties = {
   pooledCache: Cache | null,
   pooledCacheLanes: Lanes,
 
+  forceSync: boolean,
+
   // TODO: In Fizz, id generation is specific to each server config. Maybe we
   // should do this in Fiber, too? Deferring this decision for now because
   // there's no other place to store the prefix except for an internal field on

packages/react-reconciler/src/__tests__/ReactFiberHostContext-test.internal.js

@@ -70,6 +70,9 @@ describe('ReactFiberHostContext', () => {
       getCurrentEventPriority: function() {
         return DefaultEventPriority;
       },
+      getIsCurrentEventPopState() {
+        return false;
+      },
       requestPostPaintCallback: function() {},
       prepareRendererToRender: function() {},
       resetRendererAfterRender: function() {},