Hoist suspending logic into outermost work loop

This hoists the logic for checking whether the work loop is suspended
into the outermost work loop. Depending on certain conditions, we may
choose to exit the work loop, wait for a promise to resolve, wait for
microtasks to fire, and so on.

It needs to be in the outermost work loop so we can break out of it.

No intentional behavior change in this commit.

Author: acdlite

packages/react-reconciler/src/ReactFiberWorkLoop.new.js

@@ -1936,6 +1936,24 @@ function renderRootSync(root: FiberRoot, lanes: Lanes) {
 
   do {
     try {
+      if (workInProgressSuspendedReason !== NotSuspended) {
+        if (workInProgress !== null) {
+          // The current work-in-progress was already attempted. We need to unwind
+          // it before we continue the normal work loop.
+          const unitOfWork = workInProgress;
+          const thrownValue = workInProgressThrownValue;
+          workInProgressSuspendedReason = NotSuspended;
+          workInProgressThrownValue = null;
+          const wasPinged =
+            workInProgressSuspendedThenableState !== null &&
+            isThenableStateResolved(workInProgressSuspendedThenableState);
+          if (wasPinged) {
+            replaySuspendedUnitOfWork(unitOfWork, thrownValue);
+          } else {
+            unwindSuspendedUnitOfWork(unitOfWork, thrownValue);
+          }
+        }
+      }
       workLoopSync();
       break;
     } catch (thrownValue) {
@@ -1980,18 +1998,6 @@ function renderRootSync(root: FiberRoot, lanes: Lanes) {
 /** @noinline */
 function workLoopSync() {
   // Perform work without checking if we need to yield between fiber.
-
-  if (workInProgressSuspendedReason !== NotSuspended) {
-    // The current work-in-progress was already attempted. We need to unwind
-    // it before we continue the normal work loop.
-    const thrownValue = workInProgressThrownValue;
-    workInProgressSuspendedReason = NotSuspended;
-    workInProgressThrownValue = null;
-    if (workInProgress !== null) {
-      resumeSuspendedUnitOfWork(workInProgress, thrownValue);
-    }
-  }
-
   while (workInProgress !== null) {
     performUnitOfWork(workInProgress);
   }
@@ -2039,6 +2045,24 @@ function renderRootConcurrent(root: FiberRoot, lanes: Lanes) {
 
   do {
     try {
+      if (workInProgressSuspendedReason !== NotSuspended) {
+        if (workInProgress !== null) {
+          // The current work-in-progress was already attempted. We need to unwind
+          // it before we continue the normal work loop.
+          const unitOfWork = workInProgress;
+          const thrownValue = workInProgressThrownValue;
+          workInProgressSuspendedReason = NotSuspended;
+          workInProgressThrownValue = null;
+          const wasPinged =
+            workInProgressSuspendedThenableState !== null &&
+            isThenableStateResolved(workInProgressSuspendedThenableState);
+          if (wasPinged) {
+            replaySuspendedUnitOfWork(unitOfWork, thrownValue);
+          } else {
+            unwindSuspendedUnitOfWork(unitOfWork, thrownValue);
+          }
+        }
+      }
       workLoopConcurrent();
       break;
     } catch (thrownValue) {
@@ -2091,18 +2115,6 @@ function renderRootConcurrent(root: FiberRoot, lanes: Lanes) {
 /** @noinline */
 function workLoopConcurrent() {
   // Perform work until Scheduler asks us to yield
-
-  if (workInProgressSuspendedReason !== NotSuspended) {
-    // The current work-in-progress was already attempted. We need to unwind
-    // it before we continue the normal work loop.
-    const thrownValue = workInProgressThrownValue;
-    workInProgressSuspendedReason = NotSuspended;
-    workInProgressThrownValue = null;
-    if (workInProgress !== null) {
-      resumeSuspendedUnitOfWork(workInProgress, thrownValue);
-    }
-  }
-
   while (workInProgress !== null && !shouldYield()) {
     // $FlowFixMe[incompatible-call] found when upgrading Flow
     performUnitOfWork(workInProgress);
@@ -2137,69 +2149,15 @@ function performUnitOfWork(unitOfWork: Fiber): void {
   ReactCurrentOwner.current = null;
 }
 
-function resumeSuspendedUnitOfWork(
+function replaySuspendedUnitOfWork(
   unitOfWork: Fiber,
   thrownValue: mixed,
 ): void {
-  // This is a fork of performUnitOfWork specifcally for resuming a fiber that
-  // just suspended. In some cases, we may choose to retry the fiber immediately
-  // instead of unwinding the stack. It's a separate function to keep the
-  // additional logic out of the work loop's hot path.
-
-  const wasPinged =
-    workInProgressSuspendedThenableState !== null &&
-    isThenableStateResolved(workInProgressSuspendedThenableState);
-
-  if (!wasPinged) {
-    // The thenable wasn't pinged. Return to the normal work loop. This will
-    // unwind the stack, and potentially result in showing a fallback.
-    workInProgressSuspendedThenableState = null;
-
-    const returnFiber = unitOfWork.return;
-    if (returnFiber === null || workInProgressRoot === null) {
-      // Expected to be working on a non-root fiber. This is a fatal error
-      // because there's no ancestor that can handle it; the root is
-      // supposed to capture all errors that weren't caught by an error
-      // boundary.
-      workInProgressRootExitStatus = RootFatalErrored;
-      workInProgressRootFatalError = thrownValue;
-      // Set `workInProgress` to null. This represents advancing to the next
-      // sibling, or the parent if there are no siblings. But since the root
-      // has no siblings nor a parent, we set it to null. Usually this is
-      // handled by `completeUnitOfWork` or `unwindWork`, but since we're
-      // intentionally not calling those, we need set it here.
-      // TODO: Consider calling `unwindWork` to pop the contexts.
-      workInProgress = null;
-      return;
-    }
-
-    try {
-      // Find and mark the nearest Suspense or error boundary that can handle
-      // this "exception".
-      throwException(
-        workInProgressRoot,
-        returnFiber,
-        unitOfWork,
-        thrownValue,
-        workInProgressRootRenderLanes,
-      );
-    } catch (error) {
-      // We had trouble processing the error. An example of this happening is
-      // when accessing the `componentDidCatch` property of an error boundary
-      // throws an error. A weird edge case. There's a regression test for this.
-      // To prevent an infinite loop, bubble the error up to the next parent.
-      workInProgress = returnFiber;
-      throw error;
-    }
-
-    // Return to the normal work loop.
-    completeUnitOfWork(unitOfWork);
-    return;
-  }
-
-  // The work-in-progress was immediately pinged. Instead of unwinding the
-  // stack and potentially showing a fallback, unwind only the last stack frame,
-  // reset the fiber, and try rendering it again.
+  // This is a fork of performUnitOfWork specifcally for replaying a fiber that
+  // just suspended.
+  //
+  // Instead of unwinding the stack and potentially showing a fallback, unwind
+  // only the last stack frame, reset the fiber, and try rendering it again.
   const current = unitOfWork.alternate;
   unwindInterruptedWork(current, unitOfWork, workInProgressRootRenderLanes);
   unitOfWork = workInProgress = resetWorkInProgress(unitOfWork, renderLanes);
@@ -2232,6 +2190,55 @@ function resumeSuspendedUnitOfWork(
   ReactCurrentOwner.current = null;
 }
 
+function unwindSuspendedUnitOfWork(unitOfWork: Fiber, thrownValue: mixed) {
+  // This is a fork of performUnitOfWork specifcally for unwinding a fiber
+  // that threw an exception.
+  //
+  // Return to the normal work loop. This will unwind the stack, and potentially
+  // result in showing a fallback.
+  workInProgressSuspendedThenableState = null;
+
+  const returnFiber = unitOfWork.return;
+  if (returnFiber === null || workInProgressRoot === null) {
+    // Expected to be working on a non-root fiber. This is a fatal error
+    // because there's no ancestor that can handle it; the root is
+    // supposed to capture all errors that weren't caught by an error
+    // boundary.
+    workInProgressRootExitStatus = RootFatalErrored;
+    workInProgressRootFatalError = thrownValue;
+    // Set `workInProgress` to null. This represents advancing to the next
+    // sibling, or the parent if there are no siblings. But since the root
+    // has no siblings nor a parent, we set it to null. Usually this is
+    // handled by `completeUnitOfWork` or `unwindWork`, but since we're
+    // intentionally not calling those, we need set it here.
+    // TODO: Consider calling `unwindWork` to pop the contexts.
+    workInProgress = null;
+    return;
+  }
+
+  try {
+    // Find and mark the nearest Suspense or error boundary that can handle
+    // this "exception".
+    throwException(
+      workInProgressRoot,
+      returnFiber,
+      unitOfWork,
+      thrownValue,
+      workInProgressRootRenderLanes,
+    );
+  } catch (error) {
+    // We had trouble processing the error. An example of this happening is
+    // when accessing the `componentDidCatch` property of an error boundary
+    // throws an error. A weird edge case. There's a regression test for this.
+    // To prevent an infinite loop, bubble the error up to the next parent.
+    workInProgress = returnFiber;
+    throw error;
+  }
+
+  // Return to the normal work loop.
+  completeUnitOfWork(unitOfWork);
+}
+
 export function getSuspendedThenableState(): ThenableState | null {
   return workInProgressSuspendedThenableState;
 }