Get next expiration time from FiberRoot

Given a FiberRoot, we should be able to determine the next expiration
time that needs to be worked on, taking into account the levels that
are pending, suspended, pinged, and so on.

This removes the `expirationTime` argument from
`scheduleCallbackForRoot`, and renames it to `ensureRootIsScheduled` to
reflect the new signature. The expiration time is instead read from the
root using a new function, `getNextExpirationTimeToWorkOn`.

The next step will be to remove the `expirationTime` argument from
`renderRoot`, too.

Author: acdlite

packages/react-reconciler/src/ReactFiberRoot.js

@@ -14,6 +14,7 @@ import type {TimeoutHandle, NoTimeout} from './ReactFiberHostConfig';
 import type {Thenable} from './ReactFiberWorkLoop';
 import type {Interaction} from 'scheduler/src/Tracing';
 import type {SuspenseHydrationCallbacks} from './ReactFiberSuspenseComponent';
+import type {ReactPriorityLevel} from './SchedulerWithReactIntegration';
 
 import {noTimeout} from './ReactFiberHostConfig';
 import {createHostRootFiber} from './ReactFiber';
@@ -23,6 +24,7 @@ import {
   enableSuspenseCallback,
 } from 'shared/ReactFeatureFlags';
 import {unstable_getThreadID} from 'scheduler/tracing';
+import {NoPriority} from './SchedulerWithReactIntegration';
 
 // TODO: This should be lifted into the renderer.
 export type Batch = {
@@ -69,6 +71,8 @@ type BaseFiberRootProperties = {|
   callbackNode: *,
   // Expiration of the callback associated with this root
   callbackExpirationTime: ExpirationTime,
+  // Priority of the callback associated with this root
+  callbackPriority: ReactPriorityLevel,
   // The earliest pending expiration time that exists in the tree
   firstPendingTime: ExpirationTime,
   // The latest pending expiration time that exists in the tree
@@ -78,7 +82,7 @@ type BaseFiberRootProperties = {|
   // The latest suspended expiration time that exists in the tree
   lastSuspendedTime: ExpirationTime,
   // The next known expiration time after the suspended range
-  nextAfterSuspendedTime: ExpirationTime,
+  nextKnownPendingLevel: ExpirationTime,
   // The latest time at which a suspended component pinged the root to
   // render again
   lastPingedTime: ExpirationTime,
@@ -124,12 +128,12 @@ function FiberRootNode(containerInfo, tag, hydrate) {
   this.hydrate = hydrate;
   this.firstBatch = null;
   this.callbackNode = null;
-  this.callbackExpirationTime = NoWork;
+  this.callbackPriority = NoPriority;
   this.firstPendingTime = NoWork;
   this.lastPendingTime = NoWork;
   this.firstSuspendedTime = NoWork;
   this.lastSuspendedTime = NoWork;
-  this.nextAfterSuspendedTime = NoWork;
+  this.nextKnownPendingLevel = NoWork;
   this.lastPingedTime = NoWork;
 
   if (enableSchedulerTracing) {
@@ -193,21 +197,66 @@ export function markRootSuspendedAtTime(
   }
 }
 
-export function markRootUnsuspendedAtTime(
+export function markRootUpdatedAtTime(
   root: FiberRoot,
   expirationTime: ExpirationTime,
 ): void {
-  if (expirationTime <= root.lastSuspendedTime) {
+  // Update the range of pending times
+  const firstPendingTime = root.firstPendingTime;
+  if (expirationTime > firstPendingTime) {
+    root.firstPendingTime = expirationTime;
+  }
+  const lastPendingTime = root.lastPendingTime;
+  if (lastPendingTime === NoWork || expirationTime < lastPendingTime) {
+    root.lastPendingTime = expirationTime;
+  }
+
+  // Update the range of suspended times. Treat everything lower priority or
+  // equal to this update as unsuspended.
+  const firstSuspendedTime = root.firstSuspendedTime;
+  if (firstSuspendedTime !== NoWork) {
+    if (expirationTime >= firstSuspendedTime) {
+      // The entire suspended range is now unsuspended.
+      root.firstSuspendedTime = root.lastSuspendedTime = root.nextKnownPendingLevel = NoWork;
+    } else if (expirationTime >= root.lastSuspendedTime) {
+      root.lastSuspendedTime = expirationTime + 1;
+    }
+
+    // This is a pending level. Check if it's higher priority than the next
+    // known pending level.
+    if (expirationTime > root.nextKnownPendingLevel) {
+      root.nextKnownPendingLevel = expirationTime;
+    }
+  }
+}
+
+export function markRootFinishedAtTime(
+  root: FiberRoot,
+  finishedExpirationTime: ExpirationTime,
+  remainingExpirationTime: ExpirationTime,
+): void {
+  // Update the range of pending times
+  root.firstPendingTime = remainingExpirationTime;
+  if (remainingExpirationTime < root.lastPendingTime) {
+    // This usually means we've finished all the work, but it can also happen
+    // when something gets downprioritized during render, like a hidden tree.
+    root.lastPendingTime = remainingExpirationTime;
+  }
+
+  // Update the range of suspended times. Treat everything higher priority or
+  // equal to this update as unsuspended.
+  if (finishedExpirationTime <= root.lastSuspendedTime) {
     // The entire suspended range is now unsuspended.
-    root.firstSuspendedTime = root.lastSuspendedTime = root.nextAfterSuspendedTime = NoWork;
-  } else if (expirationTime <= root.firstSuspendedTime) {
+    root.firstSuspendedTime = root.lastSuspendedTime = root.nextKnownPendingLevel = NoWork;
+  } else if (finishedExpirationTime <= root.firstSuspendedTime) {
     // Part of the suspended range is now unsuspended. Narrow the range to
     // include everything between the unsuspended time (non-inclusive) and the
     // last suspended time.
-    root.firstSuspendedTime = expirationTime - 1;
+    root.firstSuspendedTime = finishedExpirationTime - 1;
   }
 
-  if (expirationTime <= root.lastPingedTime) {
+  if (finishedExpirationTime <= root.lastPingedTime) {
+    // Clear the pinged time
     root.lastPingedTime = NoWork;
   }
 }