Re-enable risky work loop changes

packages/react-reconciler/src/ReactFiberBeginWork.js

@@ -177,6 +177,7 @@ import {
   retryDehydratedSuspenseBoundary,
   scheduleWork,
   renderDidSuspendDelayIfPossible,
+  markUnprocessedUpdateTime,
 } from './ReactFiberWorkLoop';
 
 const ReactCurrentOwner = ReactSharedInternals.ReactCurrentOwner;
@@ -2709,6 +2710,11 @@ function bailoutOnAlreadyFinishedWork(
     stopProfilerTimerIfRunning(workInProgress);
   }
 
+  const updateExpirationTime = workInProgress.expirationTime;
+  if (updateExpirationTime !== NoWork) {
+    markUnprocessedUpdateTime(updateExpirationTime);
+  }
+
   // Check if the children have any pending work.
   const childExpirationTime = workInProgress.childExpirationTime;
   if (childExpirationTime < renderExpirationTime) {

packages/react-reconciler/src/ReactFiberExpirationTime.js

@@ -21,7 +21,10 @@ import {
 export type ExpirationTime = number;
 
 export const NoWork = 0;
+// TODO: Think of a better name for Never.
 export const Never = 1;
+// TODO: Use the Idle expiration time for idle state updates
+export const Idle = 2;
 export const Sync = MAX_SIGNED_31_BIT_INT;
 export const Batched = Sync - 1;
 

packages/react-reconciler/src/ReactFiberHooks.js

@@ -43,6 +43,7 @@ import {
   warnIfNotCurrentlyActingUpdatesInDev,
   warnIfNotScopedWithMatchingAct,
   markRenderEventTimeAndConfig,
+  markUnprocessedUpdateTime,
 } from './ReactFiberWorkLoop';
 
 import invariant from 'shared/invariant';
@@ -531,6 +532,7 @@ export function resetHooks(): void {
   // This is used to reset the state of this module when a component throws.
   // It's also called inside mountIndeterminateComponent if we determine the
   // component is a module-style component.
+
   renderExpirationTime = NoWork;
   currentlyRenderingFiber = null;
 
@@ -755,6 +757,7 @@ function updateReducer<S, I, A>(
         // Update the remaining priority in the queue.
         if (updateExpirationTime > remainingExpirationTime) {
           remainingExpirationTime = updateExpirationTime;
+          markUnprocessedUpdateTime(remainingExpirationTime);
         }
       } else {
         // This update does have sufficient priority.

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,12 +71,20 @@ 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
-  lastPendingTime: ExpirationTime,
-  // The time at which a suspended component pinged the root to render again
-  pingTime: ExpirationTime,
+  // The earliest suspended expiration time that exists in the tree
+  firstSuspendedTime: ExpirationTime,
+  // The latest suspended expiration time that exists in the tree
+  lastSuspendedTime: ExpirationTime,
+  // The next known expiration time after the suspended range
+  nextKnownPendingLevel: ExpirationTime,
+  // The latest time at which a suspended component pinged the root to
+  // render again
+  lastPingedTime: ExpirationTime,
+  lastExpiredTime: ExpirationTime,
 |};
 
 // The following attributes are only used by interaction tracing builds.
@@ -117,10 +127,13 @@ 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.pingTime = NoWork;
+  this.firstSuspendedTime = NoWork;
+  this.lastSuspendedTime = NoWork;
+  this.nextKnownPendingLevel = NoWork;
+  this.lastPingedTime = NoWork;
+  this.lastExpiredTime = NoWork;
 
   if (enableSchedulerTracing) {
     this.interactionThreadID = unstable_getThreadID();
@@ -151,3 +164,108 @@ export function createFiberRoot(
 
   return root;
 }
+
+export function isRootSuspendedAtTime(
+  root: FiberRoot,
+  expirationTime: ExpirationTime,
+): boolean {
+  const firstSuspendedTime = root.firstSuspendedTime;
+  const lastSuspendedTime = root.lastSuspendedTime;
+  return (
+    firstSuspendedTime !== NoWork &&
+    (firstSuspendedTime >= expirationTime &&
+      lastSuspendedTime <= expirationTime)
+  );
+}
+
+export function markRootSuspendedAtTime(
+  root: FiberRoot,
+  expirationTime: ExpirationTime,
+): void {
+  const firstSuspendedTime = root.firstSuspendedTime;
+  const lastSuspendedTime = root.lastSuspendedTime;
+  if (firstSuspendedTime < expirationTime) {
+    root.firstSuspendedTime = expirationTime;
+  }
+  if (lastSuspendedTime > expirationTime || firstSuspendedTime === NoWork) {
+    root.lastSuspendedTime = expirationTime;
+  }
+
+  if (expirationTime <= root.lastPingedTime) {
+    root.lastPingedTime = NoWork;
+  }
+
+  if (expirationTime <= root.lastExpiredTime) {
+    root.lastExpiredTime = NoWork;
+  }
+}
+
+export function markRootUpdatedAtTime(
+  root: FiberRoot,
+  expirationTime: ExpirationTime,
+): void {
+  // Update the range of pending times
+  const firstPendingTime = root.firstPendingTime;
+  if (expirationTime > firstPendingTime) {
+    root.firstPendingTime = 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;
+
+  // 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.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 = finishedExpirationTime - 1;
+  }
+
+  if (finishedExpirationTime <= root.lastPingedTime) {
+    // Clear the pinged time
+    root.lastPingedTime = NoWork;
+  }
+
+  if (finishedExpirationTime <= root.lastExpiredTime) {
+    // Clear the expired time
+    root.lastExpiredTime = NoWork;
+  }
+}
+
+export function markRootExpiredAtTime(
+  root: FiberRoot,
+  expirationTime: ExpirationTime,
+): void {
+  const lastExpiredTime = root.lastExpiredTime;
+  if (lastExpiredTime === NoWork || lastExpiredTime > expirationTime) {
+    root.lastExpiredTime = expirationTime;
+  }
+}