[WIP] Schedule prerender after something suspends

Part of a series of PRs related to sibling prerendering. I have the
implementation working locally, but still working on splitting it up
into a reasonable sequence of steps so we can land it incrementally.
There's likely to be some regressions due to the scope of the change,
so I've also done my best to keep the change behind a feature flag.

Opening this in draft mode while I continue to work on updating the
test suite, which requires many changes.

---

Adds the concept of a "prerender". These special renders are spawned
whenever something suspends (and we're not already prerendering).

The purpose is to move speculative rendering work into a separate
phase that does not block the UI from updating. For example, during a
transition, if something suspends, we should not speculatively
prerender siblings that will be replaced by a fallback in the UI until
*after* the fallback has been shown to the user.

Author: acdlite

packages/react-cache/src/__tests__/ReactCacheOld-test.internal.js

@@ -148,7 +148,15 @@ describe('ReactCache', () => {
       error = e;
     }
     expect(error.message).toMatch('Failed to load: Hi');
-    assertLog(['Promise rejected [Hi]', 'Error! [Hi]', 'Error! [Hi]']);
+    assertLog([
+      'Promise rejected [Hi]',
+      'Error! [Hi]',
+      'Error! [Hi]',
+
+      ...(gate('enableSiblingPrerendering')
+        ? ['Error! [Hi]', 'Error! [Hi]']
+        : []),
+    ]);
 
     // Should throw again on a subsequent read
     root.render(<App />);
@@ -191,6 +199,7 @@ describe('ReactCache', () => {
     }
   });
 
+  // @gate enableSiblingPrerendering
   it('evicts least recently used values', async () => {
     ReactCache.unstable_setGlobalCacheLimit(3);
 
@@ -206,15 +215,13 @@ describe('ReactCache', () => {
     await waitForAll(['Suspend! [1]', 'Loading...']);
     jest.advanceTimersByTime(100);
     assertLog(['Promise resolved [1]']);
-    await waitForAll([1, 'Suspend! [2]']);
+    await waitForAll([1, 'Suspend! [2]', 1, 'Suspend! [2]', 'Suspend! [3]']);
 
     jest.advanceTimersByTime(100);
-    assertLog(['Promise resolved [2]']);
-    await waitForAll([1, 2, 'Suspend! [3]']);
+    assertLog(['Promise resolved [2]', 'Promise resolved [3]']);
+    await waitForAll([1, 2, 3]);
 
     await act(() => jest.advanceTimersByTime(100));
-    assertLog(['Promise resolved [3]', 1, 2, 3]);
-
     expect(root).toMatchRenderedOutput('123');
 
     // Render 1, 4, 5
@@ -234,6 +241,9 @@ describe('ReactCache', () => {
       1,
       4,
       'Suspend! [5]',
+      1,
+      4,
+      'Suspend! [5]',
       'Promise resolved [5]',
       1,
       4,
@@ -267,6 +277,9 @@ describe('ReactCache', () => {
       1,
       2,
       'Suspend! [3]',
+      1,
+      2,
+      'Suspend! [3]',
       'Promise resolved [3]',
       1,
       2,

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([]);
+    assertLog(gate('enableSiblingPrerendering') ? ['Suspend! [/path/b]'] : []);
 
     await act(async () => {
       resolvePromise();

packages/react-reconciler/src/ReactFiberLane.js

@@ -229,28 +229,49 @@ export function getNextLanes(root: FiberRoot, wipLanes: Lanes): Lanes {
 
   const suspendedLanes = root.suspendedLanes;
   const pingedLanes = root.pingedLanes;
+  const warmLanes = root.warmLanes;
 
   // Do not work on any idle work until all the non-idle work has finished,
   // even if the work is suspended.
   const nonIdlePendingLanes = pendingLanes & NonIdleLanes;
   if (nonIdlePendingLanes !== NoLanes) {
+    // First check for fresh updates.
     const nonIdleUnblockedLanes = nonIdlePendingLanes & ~suspendedLanes;
     if (nonIdleUnblockedLanes !== NoLanes) {
       nextLanes = getHighestPriorityLanes(nonIdleUnblockedLanes);
     } else {
+      // No fresh updates. Check if suspended work has been pinged.
       const nonIdlePingedLanes = nonIdlePendingLanes & pingedLanes;
       if (nonIdlePingedLanes !== NoLanes) {
         nextLanes = getHighestPriorityLanes(nonIdlePingedLanes);
+      } else {
+        // Nothing has been pinged. Check for lanes that need to be prewarmed.
+        const lanesToPrewarm = nonIdlePendingLanes & ~warmLanes;
+        if (lanesToPrewarm !== NoLanes) {
+          nextLanes = getHighestPriorityLanes(lanesToPrewarm);
+        }
       }
     }
   } else {
     // The only remaining work is Idle.
+    // TODO: Idle isn't really used anywhere, and the thinking around
+    // speculative rendering has evolved since this was implemented. Consider
+    // removing until we've thought about this again.
+
+    // First check for fresh updates.
     const unblockedLanes = pendingLanes & ~suspendedLanes;
     if (unblockedLanes !== NoLanes) {
       nextLanes = getHighestPriorityLanes(unblockedLanes);
     } else {
+      // No fresh updates. Check if suspended work has been pinged.
       if (pingedLanes !== NoLanes) {
         nextLanes = getHighestPriorityLanes(pingedLanes);
+      } else {
+        // Nothing has been pinged. Check for lanes that need to be prewarmed.
+        const lanesToPrewarm = pendingLanes & ~warmLanes;
+        if (lanesToPrewarm !== NoLanes) {
+          nextLanes = getHighestPriorityLanes(lanesToPrewarm);
+        }
       }
     }
   }
@@ -335,6 +356,21 @@ export function getNextLanesToFlushSync(
   return NoLanes;
 }
 
+export function checkIfRootIsPrerendering(
+  root: FiberRoot,
+  renderLanes: Lanes,
+): boolean {
+  const pendingLanes = root.pendingLanes;
+  const suspendedLanes = root.suspendedLanes;
+  const pingedLanes = root.pingedLanes;
+  // Remove lanes that are suspended (but not pinged)
+  const unblockedLanes = pendingLanes & ~(suspendedLanes & ~pingedLanes);
+
+  // If there are no unsuspended or pinged lanes, that implies that we're
+  // performing a prerender.
+  return unblockedLanes === 0;
+}
+
 export function getEntangledLanes(root: FiberRoot, renderLanes: Lanes): Lanes {
   let entangledLanes = renderLanes;
 
@@ -670,17 +706,27 @@ export function markRootUpdated(root: FiberRoot, updateLane: Lane) {
   if (updateLane !== IdleLane) {
     root.suspendedLanes = NoLanes;
     root.pingedLanes = NoLanes;
+    root.warmLanes = NoLanes;
   }
 }
 
 export function markRootSuspended(
   root: FiberRoot,
   suspendedLanes: Lanes,
   spawnedLane: Lane,
+  didSkipSuspendedSiblings: boolean,
 ) {
   root.suspendedLanes |= suspendedLanes;
   root.pingedLanes &= ~suspendedLanes;
 
+  if (!didSkipSuspendedSiblings) {
+    // Mark these lanes as warm so we know there's nothing else to work on.
+    root.warmLanes |= suspendedLanes;
+  } else {
+    // Render unwound without attempting all the siblings. Do no mark the lanes
+    // as warm. This will cause a prewarm render to be scheduled.
+  }
+
   // The suspended lanes are no longer CPU-bound. Clear their expiration times.
   const expirationTimes = root.expirationTimes;
   let lanes = suspendedLanes;
@@ -700,6 +746,9 @@ export function markRootSuspended(
 
 export function markRootPinged(root: FiberRoot, pingedLanes: Lanes) {
   root.pingedLanes |= root.suspendedLanes & pingedLanes;
+  // The data that just resolved could have unblocked additional children, which
+  // will also need to be prewarmed if something suspends again.
+  root.warmLanes &= ~pingedLanes;
 }
 
 export function markRootFinished(
@@ -714,6 +763,7 @@ export function markRootFinished(
   // Let's try everything again
   root.suspendedLanes = NoLanes;
   root.pingedLanes = NoLanes;
+  root.warmLanes = NoLanes;
 
   root.expiredLanes &= remainingLanes;
 

packages/react-reconciler/src/ReactFiberRoot.js

@@ -75,6 +75,7 @@ function FiberRootNode(
   this.pendingLanes = NoLanes;
   this.suspendedLanes = NoLanes;
   this.pingedLanes = NoLanes;
+  this.warmLanes = NoLanes;
   this.expiredLanes = NoLanes;
   this.finishedLanes = NoLanes;
   this.errorRecoveryDisabledLanes = NoLanes;