Track entangled lanes separately from update lane

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

@@ -16,6 +16,8 @@ let ReactDOMClient;
 let Scheduler;
 let act;
 let waitForAll;
+let waitFor;
+let waitForMicrotasks;
 let assertLog;
 
 const setUntrackedInputValue = Object.getOwnPropertyDescriptor(
@@ -36,6 +38,8 @@ describe('ReactDOMFiberAsync', () => {
 
     const InternalTestUtils = require('internal-test-utils');
     waitForAll = InternalTestUtils.waitForAll;
+    waitFor = InternalTestUtils.waitFor;
+    waitForMicrotasks = InternalTestUtils.waitForMicrotasks;
     assertLog = InternalTestUtils.assertLog;
 
     document.body.appendChild(container);
@@ -653,52 +657,94 @@ describe('ReactDOMFiberAsync', () => {
     });
   });
 
-  it('transition lane in popState should yield if it suspends', async () => {
-    const never = {then() {}};
-    let _setText;
+  it('transition lane in popState should be allowed to suspend', async () => {
+    let resolvePromise;
+    const promise = new Promise(res => {
+      resolvePromise = res;
+    });
+
+    function Text({text}) {
+      Scheduler.log(text);
+      return text;
+    }
 
     function App() {
-      const [shouldSuspend, setShouldSuspend] = React.useState(false);
-      const [text, setText] = React.useState('0');
-      _setText = setText;
-      if (shouldSuspend) {
-        Scheduler.log('Suspend!');
-        throw never;
-      }
-      function onPopstate() {
-        React.startTransition(() => {
-          setShouldSuspend(val => !val);
-        });
+      const [pathname, setPathname] = React.useState('/path/a');
+
+      if (pathname !== '/path/a') {
+        try {
+          React.use(promise);
+        } catch (e) {
+          Scheduler.log(`Suspend! [${pathname}]`);
+          throw e;
+        }
       }
+
       React.useEffect(() => {
+        function onPopstate() {
+          React.startTransition(() => {
+            setPathname('/path/b');
+          });
+        }
         window.addEventListener('popstate', onPopstate);
         return () => window.removeEventListener('popstate', onPopstate);
       }, []);
-      Scheduler.log(`Child:${shouldSuspend}/${text}`);
-      return text;
+
+      return (
+        <>
+          <Text text="Before" />
+          <div>
+            <Text text={pathname} />
+          </div>
+          <Text text="After" />
+        </>
+      );
     }
 
     const root = ReactDOMClient.createRoot(container);
     await act(async () => {
       root.render(<App />);
     });
-    assertLog(['Child:false/0']);
+    assertLog(['Before', '/path/a', 'After']);
 
-    await act(() => {
+    const div = container.getElementsByTagName('div')[0];
+    expect(div.textContent).toBe('/path/a');
+
+    // Simulate a popstate event
+    await act(async () => {
       const popStateEvent = new Event('popstate');
+
+      // Simulate a popstate event
       window.event = popStateEvent;
       window.dispatchEvent(popStateEvent);
-      queueMicrotask(() => {
-        window.event = undefined;
-      });
+      await waitForMicrotasks();
+      window.event = undefined;
+
+      // The transition lane should have been attempted synchronously (in
+      // a microtask)
+      assertLog(['Suspend! [/path/b]']);
+      // Because it suspended, it remains on the current path
+      expect(div.textContent).toBe('/path/a');
     });
-    assertLog(['Suspend!']);
+    assertLog(['Suspend! [/path/b]']);
 
     await act(async () => {
-      _setText('1');
-    });
-    assertLog(['Child:false/1', 'Suspend!']);
+      resolvePromise();
+
+      // Since the transition previously suspended, there's no need for this
+      // transition to be rendered synchronously on susbequent attempts; if we
+      // fail to commit synchronously the first time, the scroll restoration
+      // state won't be restored anyway.
+      //
+      // Yield in between each child to prove that it's concurrent.
+      await waitForMicrotasks();
+      assertLog([]);
 
-    root.unmount();
+      await waitFor(['Before']);
+      await waitFor(['/path/b']);
+      await waitFor(['After']);
+    });
+    assertLog([]);
+    expect(div.textContent).toBe('/path/b');
   });
 });

packages/react-reconciler/src/ReactFiberHiddenContext.js

@@ -13,7 +13,10 @@ import type {Lanes} from './ReactFiberLane';
 
 import {createCursor, push, pop} from './ReactFiberStack';
 
-import {getRenderLanes, setRenderLanes} from './ReactFiberWorkLoop';
+import {
+  getEntangledRenderLanes,
+  setEntangledRenderLanes,
+} from './ReactFiberWorkLoop';
 import {NoLanes, mergeLanes} from './ReactFiberLane';
 
 // TODO: Remove `renderLanes` context in favor of hidden context
@@ -29,26 +32,28 @@ type HiddenContext = {
 // InvisibleParentContext that is currently managed by SuspenseContext.
 export const currentTreeHiddenStackCursor: StackCursor<HiddenContext | null> =
   createCursor(null);
-export const prevRenderLanesStackCursor: StackCursor<Lanes> =
+export const prevEntangledRenderLanesCursor: StackCursor<Lanes> =
   createCursor(NoLanes);
 
 export function pushHiddenContext(fiber: Fiber, context: HiddenContext): void {
-  const prevRenderLanes = getRenderLanes();
-  push(prevRenderLanesStackCursor, prevRenderLanes, fiber);
+  const prevEntangledRenderLanes = getEntangledRenderLanes();
+  push(prevEntangledRenderLanesCursor, prevEntangledRenderLanes, fiber);
   push(currentTreeHiddenStackCursor, context, fiber);
 
   // When rendering a subtree that's currently hidden, we must include all
   // lanes that would have rendered if the hidden subtree hadn't been deferred.
   // That is, in order to reveal content from hidden -> visible, we must commit
   // all the updates that we skipped when we originally hid the tree.
-  setRenderLanes(mergeLanes(prevRenderLanes, context.baseLanes));
+  setEntangledRenderLanes(
+    mergeLanes(prevEntangledRenderLanes, context.baseLanes),
+  );
 }
 
 export function reuseHiddenContextOnStack(fiber: Fiber): void {
   // This subtree is not currently hidden, so we don't need to add any lanes
   // to the render lanes. But we still need to push something to avoid a
   // context mismatch. Reuse the existing context on the stack.
-  push(prevRenderLanesStackCursor, getRenderLanes(), fiber);
+  push(prevEntangledRenderLanesCursor, getEntangledRenderLanes(), fiber);
   push(
     currentTreeHiddenStackCursor,
     currentTreeHiddenStackCursor.current,
@@ -58,10 +63,10 @@ export function reuseHiddenContextOnStack(fiber: Fiber): void {
 
 export function popHiddenContext(fiber: Fiber): void {
   // Restore the previous render lanes from the stack
-  setRenderLanes(prevRenderLanesStackCursor.current);
+  setEntangledRenderLanes(prevEntangledRenderLanesCursor.current);
 
   pop(currentTreeHiddenStackCursor, fiber);
-  pop(prevRenderLanesStackCursor, fiber);
+  pop(prevEntangledRenderLanesCursor, fiber);
 }
 
 export function isCurrentTreeHidden(): boolean {

packages/react-reconciler/src/ReactFiberHooks.js

@@ -1525,7 +1525,8 @@ function mountSyncExternalStore<T>(
       );
     }
 
-    if (!includesBlockingLane(root, renderLanes)) {
+    const rootRenderLanes = getWorkInProgressRootRenderLanes();
+    if (!includesBlockingLane(root, rootRenderLanes)) {
       pushStoreConsistencyCheck(fiber, getSnapshot, nextSnapshot);
     }
   }

packages/react-reconciler/src/ReactFiberLane.js

@@ -39,6 +39,7 @@ export const NoLane: Lane = /*                          */ 0b0000000000000000000
 
 export const SyncHydrationLane: Lane = /*               */ 0b0000000000000000000000000000001;
 export const SyncLane: Lane = /*                        */ 0b0000000000000000000000000000010;
+export const SyncLaneIndex: number = 1;
 
 export const InputContinuousHydrationLane: Lane = /*    */ 0b0000000000000000000000000000100;
 export const InputContinuousLane: Lane = /*             */ 0b0000000000000000000000000001000;
@@ -274,17 +275,23 @@ export function getNextLanes(root: FiberRoot, wipLanes: Lanes): Lanes {
     }
   }
 
+  return nextLanes;
+}
+
+export function getEntangledLanes(root: FiberRoot, renderLanes: Lanes): Lanes {
+  let entangledLanes = renderLanes;
+
   if (
     allowConcurrentByDefault &&
     (root.current.mode & ConcurrentUpdatesByDefaultMode) !== NoMode
   ) {
     // Do nothing, use the lanes as they were assigned.
-  } else if ((nextLanes & InputContinuousLane) !== NoLanes) {
+  } else if ((entangledLanes & InputContinuousLane) !== NoLanes) {
     // When updates are sync by default, we entangle continuous priority updates
     // and default updates, so they render in the same batch. The only reason
     // they use separate lanes is because continuous updates should interrupt
     // transitions, but default updates should not.
-    nextLanes |= pendingLanes & DefaultLane;
+    entangledLanes |= entangledLanes & DefaultLane;
   }
 
   // Check for entangled lanes and add them to the batch.
@@ -309,21 +316,21 @@ export function getNextLanes(root: FiberRoot, wipLanes: Lanes): Lanes {
   // For those exceptions where entanglement is semantically important,
   // we should ensure that there is no partial work at the
   // time we apply the entanglement.
-  const entangledLanes = root.entangledLanes;
-  if (entangledLanes !== NoLanes) {
+  const allEntangledLanes = root.entangledLanes;
+  if (allEntangledLanes !== NoLanes) {
     const entanglements = root.entanglements;
-    let lanes = nextLanes & entangledLanes;
+    let lanes = entangledLanes & allEntangledLanes;
     while (lanes > 0) {
       const index = pickArbitraryLaneIndex(lanes);
       const lane = 1 << index;
 
-      nextLanes |= entanglements[index];
+      entangledLanes |= entanglements[index];
 
       lanes &= ~lane;
     }
   }
 
-  return nextLanes;
+  return entangledLanes;
 }
 
 function computeExpirationTime(lane: Lane, currentTime: number) {
@@ -404,6 +411,7 @@ export function markStarvedLanesAsExpired(
   // Iterate through the pending lanes and check if we've reached their
   // expiration time. If so, we'll assume the update is being starved and mark
   // it as expired to force it to finish.
+  // TODO: We should be able to replace this with upgradePendingLanesToSync
   //
   // We exclude retry lanes because those must always be time sliced, in order
   // to unwrap uncached promises.
@@ -708,6 +716,34 @@ 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,
+) {
+  // Same as upgradePendingLaneToSync but accepts multiple lanes, so it's a
+  // bit slower.
+  root.pendingLanes |= SyncLane;
+  root.entangledLanes |= SyncLane;
+  let lanes = lanesToUpgrade;
+  while (lanes) {
+    const index = pickArbitraryLaneIndex(lanes);
+    const lane = 1 << index;
+    root.entanglements[SyncLaneIndex] |= lane;
+    lanes &= ~lane;
+  }
+}
+
 export function markHiddenUpdate(
   root: FiberRoot,
   update: ConcurrentUpdate,

packages/react-reconciler/src/ReactFiberRootScheduler.js

@@ -20,8 +20,7 @@ import {
   getNextLanes,
   includesSyncLane,
   markStarvedLanesAsExpired,
-  markRootEntangled,
-  mergeLanes,
+  upgradePendingLaneToSync,
   claimNextTransitionLane,
 } from './ReactFiberLane';
 import {
@@ -250,7 +249,10 @@ function processRootScheduleInMicrotask() {
       currentEventTransitionLane !== NoLane &&
       shouldAttemptEagerTransition()
     ) {
-      markRootEntangled(root, mergeLanes(currentEventTransitionLane, SyncLane));
+      // 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);