Fix: uDV skipped initial value if earlier transition suspended (#34376)

Fixes a bug in useDeferredValue's optional `initialValue` argument. In
the regression case, if a new useDeferredValue hook is mounted while an
earlier transition is suspended, the `initialValue` argument of the new
hook was ignored. After the fix, the `initialValue` argument is
correctly rendered during the initial mount, regardless of whether other
transitions were suspended.

The culprit was related to the mechanism we use to track whether a
render is the result of a `useDeferredValue` hook: we assign the
deferred lane a TransitionLane, then entangle that lane with the
DeferredLane bit. During the subsequent render, we check for the
presence of the DeferredLane bit to determine whether to switch to the
final, canonical value.

But because transition lanes can themselves become entangled with other
transitions, the effect is that every entangled transition was being
treated as if it were the result of a `useDeferredValue` hook, causing
us to skip the initial value and go straight to the final one.

The fix I've chosen is to reserve some subset of TransitionLanes to be
used only for deferred work, instead of using entanglement. This is
similar to how retries are already implemented. Originally I tried not
to implement it this way because it means there are now slightly fewer
lanes allocated for regular transitions, but I underestimated how
similar deferred work is to retries; they end up having a lot of the
same requirements. Eventually it may be possible to merge the two
concepts.

Author: acdlite

packages/react-reconciler/src/ReactFiberHooks.js

@@ -73,6 +73,7 @@ import {
   includesSomeLane,
   isGestureRender,
   GestureLane,
+  UpdateLanes,
 } from './ReactFiberLane';
 import {
   ContinuousEventPriority,
@@ -2983,6 +2984,20 @@ function rerenderDeferredValue<T>(value: T, initialValue?: T): T {
   }
 }
 
+function isRenderingDeferredWork(): boolean {
+  if (!includesSomeLane(renderLanes, DeferredLane)) {
+    // None of the render lanes are deferred lanes.
+    return false;
+  }
+  // At least one of the render lanes are deferred lanes. However, if the
+  // current render is also batched together with an update, then we can't
+  // say that the render is wholly the result of deferred work. We can check
+  // this by checking if the root render lanes contain any "update" lanes, i.e.
+  // lanes that are only assigned to updates, like setState.
+  const rootRenderLanes = getWorkInProgressRootRenderLanes();
+  return !includesSomeLane(rootRenderLanes, UpdateLanes);
+}
+
 function mountDeferredValueImpl<T>(hook: Hook, value: T, initialValue?: T): T {
   if (
     // When `initialValue` is provided, we defer the initial render even if the
@@ -2991,7 +3006,7 @@ function mountDeferredValueImpl<T>(hook: Hook, value: T, initialValue?: T): T {
     // However, to avoid waterfalls, we do not defer if this render
     // was itself spawned by an earlier useDeferredValue. Check if DeferredLane
     // is part of the render lanes.
-    !includesSomeLane(renderLanes, DeferredLane)
+    !isRenderingDeferredWork()
   ) {
     // Render with the initial value
     hook.memoizedState = initialValue;
@@ -3038,8 +3053,7 @@ function updateDeferredValueImpl<T>(
     }
 
     const shouldDeferValue =
-      !includesOnlyNonUrgentLanes(renderLanes) &&
-      !includesSomeLane(renderLanes, DeferredLane);
+      !includesOnlyNonUrgentLanes(renderLanes) && !isRenderingDeferredWork();
     if (shouldDeferValue) {
       // This is an urgent update. Since the value has changed, keep using the
       // previous value and spawn a deferred render to update it later.

packages/react-reconciler/src/ReactFiberLane.js

@@ -73,6 +73,20 @@ const TransitionLane12: Lane = /*                       */ 0b0000000000010000000
 const TransitionLane13: Lane = /*                       */ 0b0000000000100000000000000000000;
 const TransitionLane14: Lane = /*                       */ 0b0000000001000000000000000000000;
 
+const TransitionUpdateLanes =
+  TransitionLane1 |
+  TransitionLane2 |
+  TransitionLane3 |
+  TransitionLane4 |
+  TransitionLane5 |
+  TransitionLane6 |
+  TransitionLane7 |
+  TransitionLane8 |
+  TransitionLane9 |
+  TransitionLane10;
+const TransitionDeferredLanes =
+  TransitionLane11 | TransitionLane12 | TransitionLane13 | TransitionLane14;
+
 const RetryLanes: Lanes = /*                            */ 0b0000011110000000000000000000000;
 const RetryLane1: Lane = /*                             */ 0b0000000010000000000000000000000;
 const RetryLane2: Lane = /*                             */ 0b0000000100000000000000000000000;
@@ -94,7 +108,7 @@ export const DeferredLane: Lane = /*                    */ 0b1000000000000000000
 // Any lane that might schedule an update. This is used to detect infinite
 // update loops, so it doesn't include hydration lanes or retries.
 export const UpdateLanes: Lanes =
-  SyncLane | InputContinuousLane | DefaultLane | TransitionLanes;
+  SyncLane | InputContinuousLane | DefaultLane | TransitionUpdateLanes;
 
 export const HydrationLanes =
   SyncHydrationLane |
@@ -155,7 +169,8 @@ export function getLabelForLane(lane: Lane): string | void {
 
 export const NoTimestamp = -1;
 
-let nextTransitionLane: Lane = TransitionLane1;
+let nextTransitionUpdateLane: Lane = TransitionLane1;
+let nextTransitionDeferredLane: Lane = TransitionLane11;
 let nextRetryLane: Lane = RetryLane1;
 
 function getHighestPriorityLanes(lanes: Lanes | Lane): Lanes {
@@ -190,11 +205,12 @@ function getHighestPriorityLanes(lanes: Lanes | Lane): Lanes {
     case TransitionLane8:
     case TransitionLane9:
     case TransitionLane10:
+      return lanes & TransitionUpdateLanes;
     case TransitionLane11:
     case TransitionLane12:
     case TransitionLane13:
     case TransitionLane14:
-      return lanes & TransitionLanes;
+      return lanes & TransitionDeferredLanes;
     case RetryLane1:
     case RetryLane2:
     case RetryLane3:
@@ -679,14 +695,23 @@ export function isGestureRender(lanes: Lanes): boolean {
   return lanes === GestureLane;
 }
 
-export function claimNextTransitionLane(): Lane {
+export function claimNextTransitionUpdateLane(): Lane {
   // Cycle through the lanes, assigning each new transition to the next lane.
   // In most cases, this means every transition gets its own lane, until we
   // run out of lanes and cycle back to the beginning.
-  const lane = nextTransitionLane;
-  nextTransitionLane <<= 1;
-  if ((nextTransitionLane & TransitionLanes) === NoLanes) {
-    nextTransitionLane = TransitionLane1;
+  const lane = nextTransitionUpdateLane;
+  nextTransitionUpdateLane <<= 1;
+  if ((nextTransitionUpdateLane & TransitionUpdateLanes) === NoLanes) {
+    nextTransitionUpdateLane = TransitionLane1;
+  }
+  return lane;
+}
+
+export function claimNextTransitionDeferredLane(): Lane {
+  const lane = nextTransitionDeferredLane;
+  nextTransitionDeferredLane <<= 1;
+  if ((nextTransitionDeferredLane & TransitionDeferredLanes) === NoLanes) {
+    nextTransitionDeferredLane = TransitionLane11;
   }
   return lane;
 }
@@ -952,6 +977,14 @@ function markSpawnedDeferredLane(
   // Entangle the spawned lane with the DeferredLane bit so that we know it
   // was the result of another render. This lets us avoid a useDeferredValue
   // waterfall — only the first level will defer.
+  // TODO: Now that there is a reserved set of transition lanes that are used
+  // exclusively for deferred work, we should get rid of this special
+  // DeferredLane bit; the same information can be inferred by checking whether
+  // the lane is one of the TransitionDeferredLanes. The only reason this still
+  // exists is because we need to also do the same for OffscreenLane. That
+  // requires additional changes because there are more places around the
+  // codebase that treat OffscreenLane as a magic value; would need to check
+  // for a new OffscreenDeferredLane, too. Will leave this for a follow-up.
   const spawnedLaneIndex = laneToIndex(spawnedLane);
   root.entangledLanes |= spawnedLane;
   root.entanglements[spawnedLaneIndex] |=

packages/react-reconciler/src/ReactFiberRootScheduler.js

@@ -31,7 +31,7 @@ import {
   getNextLanes,
   includesSyncLane,
   markStarvedLanesAsExpired,
-  claimNextTransitionLane,
+  claimNextTransitionUpdateLane,
   getNextLanesToFlushSync,
   checkIfRootIsPrerendering,
   isGestureRender,
@@ -716,7 +716,7 @@ export function requestTransitionLane(
         : // We may or may not be inside an async action scope. If we are, this
           // is the first update in that scope. Either way, we need to get a
           // fresh transition lane.
-          claimNextTransitionLane();
+          claimNextTransitionUpdateLane();
   }
   return currentEventTransitionLane;
 }

packages/react-reconciler/src/ReactFiberWorkLoop.js

@@ -192,7 +192,7 @@ import {
   OffscreenLane,
   SyncUpdateLanes,
   UpdateLanes,
-  claimNextTransitionLane,
+  claimNextTransitionDeferredLane,
   checkIfRootIsPrerendering,
   includesOnlyViewTransitionEligibleLanes,
   isGestureRender,
@@ -827,7 +827,7 @@ export function requestDeferredLane(): Lane {
       workInProgressDeferredLane = OffscreenLane;
     } else {
       // Everything else is spawned as a transition.
-      workInProgressDeferredLane = claimNextTransitionLane();
+      workInProgressDeferredLane = claimNextTransitionDeferredLane();
     }
   }
 

packages/react-reconciler/src/__tests__/ReactDeferredValue-test.js

@@ -608,6 +608,48 @@ describe('ReactDeferredValue', () => {
     },
   );
 
+  it(
+    "regression: useDeferredValue's initial value argument works even if an unrelated " +
+      'transition is suspended',
+    async () => {
+      // Simulates a previous bug where a new useDeferredValue hook is mounted
+      // while some unrelated transition is suspended. In the regression case,
+      // the initial values was skipped/ignored.
+
+      function Content({text}) {
+        return (
+          <AsyncText text={useDeferredValue(text, `Preview ${text}...`)} />
+        );
+      }
+
+      function App({text}) {
+        // Use a key to force a new Content instance to be mounted each time
+        // the text changes.
+        return <Content key={text} text={text} />;
+      }
+
+      const root = ReactNoop.createRoot();
+
+      // Render a previous UI using useDeferredValue. Suspend on the
+      // final value.
+      resolveText('Preview A...');
+      await act(() => startTransition(() => root.render(<App text="A" />)));
+      assertLog(['Preview A...', 'Suspend! [A]']);
+
+      // While it's still suspended, update the UI to show a different screen
+      // with a different preview value. We should be able to show the new
+      // preview even though the previous transition never finished.
+      resolveText('Preview B...');
+      await act(() => startTransition(() => root.render(<App text="B" />)));
+      assertLog(['Preview B...', 'Suspend! [B]']);
+
+      // Now finish loading the final value.
+      await act(() => resolveText('B'));
+      assertLog(['B']);
+      expect(root).toMatchRenderedOutput('B');
+    },
+  );
+
   it('avoids a useDeferredValue waterfall when separated by a Suspense boundary', async () => {
     // Same as the previous test but with a Suspense boundary separating the
     // two useDeferredValue hooks.