Extract hydration logic in complete phase, too

Same as previous step but for the complete phase. This is a separate
commit to make bisecting easier in case something breaks. The logic
is very subtle but mostly all I've done is extract it to
another function.

Author: acdlite

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

@@ -759,6 +759,93 @@ function bubbleProperties(completedWork: Fiber) {
   return didBailout;
 }
 
+function completeDehydratedSuspenseBoundary(
+  current: Fiber | null,
+  workInProgress: Fiber,
+  nextState: SuspenseState | null,
+): boolean {
+  if (
+    hasUnhydratedTailNodes() &&
+    (workInProgress.mode & ConcurrentMode) !== NoMode &&
+    (workInProgress.flags & DidCapture) === NoFlags
+  ) {
+    warnIfUnhydratedTailNodes(workInProgress);
+    resetHydrationState();
+    workInProgress.flags |= ForceClientRender | Incomplete | ShouldCapture;
+
+    return false;
+  }
+
+  const wasHydrated = popHydrationState(workInProgress);
+
+  if (nextState !== null && nextState.dehydrated !== null) {
+    // We might be inside a hydration state the first time we're picking up this
+    // Suspense boundary, and also after we've reentered it for further hydration.
+    if (current === null) {
+      if (!wasHydrated) {
+        throw new Error(
+          'A dehydrated suspense component was completed without a hydrated node. ' +
+            'This is probably a bug in React.',
+        );
+      }
+      prepareToHydrateHostSuspenseInstance(workInProgress);
+      bubbleProperties(workInProgress);
+      if (enableProfilerTimer) {
+        if ((workInProgress.mode & ProfileMode) !== NoMode) {
+          const isTimedOutSuspense = nextState !== null;
+          if (isTimedOutSuspense) {
+            // Don't count time spent in a timed out Suspense subtree as part of the base duration.
+            const primaryChildFragment = workInProgress.child;
+            if (primaryChildFragment !== null) {
+              // $FlowFixMe Flow doesn't support type casting in combination with the -= operator
+              workInProgress.treeBaseDuration -= ((primaryChildFragment.treeBaseDuration: any): number);
+            }
+          }
+        }
+      }
+      return false;
+    } else {
+      // We might have reentered this boundary to hydrate it. If so, we need to reset the hydration
+      // state since we're now exiting out of it. popHydrationState doesn't do that for us.
+      resetHydrationState();
+      if ((workInProgress.flags & DidCapture) === NoFlags) {
+        // This boundary did not suspend so it's now hydrated and unsuspended.
+        workInProgress.memoizedState = null;
+      }
+      // If nothing suspended, we need to schedule an effect to mark this boundary
+      // as having hydrated so events know that they're free to be invoked.
+      // It's also a signal to replay events and the suspense callback.
+      // If something suspended, schedule an effect to attach retry listeners.
+      // So we might as well always mark this.
+      workInProgress.flags |= Update;
+      bubbleProperties(workInProgress);
+      if (enableProfilerTimer) {
+        if ((workInProgress.mode & ProfileMode) !== NoMode) {
+          const isTimedOutSuspense = nextState !== null;
+          if (isTimedOutSuspense) {
+            // Don't count time spent in a timed out Suspense subtree as part of the base duration.
+            const primaryChildFragment = workInProgress.child;
+            if (primaryChildFragment !== null) {
+              // $FlowFixMe Flow doesn't support type casting in combination with the -= operator
+              workInProgress.treeBaseDuration -= ((primaryChildFragment.treeBaseDuration: any): number);
+            }
+          }
+        }
+      }
+      return false;
+    }
+  } else {
+    // Successfully completed this tree. If this was a forced client render,
+    // there may have been recoverable errors during first hydration
+    // attempt. If so, add them to a queue so we can log them in the
+    // commit phase.
+    upgradeHydrationErrorsToRecoverable();
+
+    // Fall through to normal Suspense path
+    return true;
+  }
+}
+
 function completeWork(
   current: Fiber | null,
   workInProgress: Fiber,
@@ -996,80 +1083,35 @@ function completeWork(
       popSuspenseContext(workInProgress);
       const nextState: null | SuspenseState = workInProgress.memoizedState;
 
+      // Special path for dehydrated boundaries. We may eventually move this
+      // to its own fiber type so that we can add other kinds of hydration
+      // boundaries that aren't associated with a Suspense tree. In anticipation
+      // of such a refactor, all the hydration logic is contained in
+      // this branch.
       if (
-        hasUnhydratedTailNodes() &&
-        (workInProgress.mode & ConcurrentMode) !== NoMode &&
-        (workInProgress.flags & DidCapture) === NoFlags
+        current === null ||
+        (current.memoizedState !== null &&
+          current.memoizedState.dehydrated !== null)
       ) {
-        warnIfUnhydratedTailNodes(workInProgress);
-        resetHydrationState();
-        workInProgress.flags |= ForceClientRender | Incomplete | ShouldCapture;
-        return workInProgress;
-      }
-      if (nextState !== null && nextState.dehydrated !== null) {
-        // We might be inside a hydration state the first time we're picking up this
-        // Suspense boundary, and also after we've reentered it for further hydration.
-        const wasHydrated = popHydrationState(workInProgress);
-        if (current === null) {
-          if (!wasHydrated) {
-            throw new Error(
-              'A dehydrated suspense component was completed without a hydrated node. ' +
-                'This is probably a bug in React.',
-            );
-          }
-          prepareToHydrateHostSuspenseInstance(workInProgress);
-          bubbleProperties(workInProgress);
-          if (enableProfilerTimer) {
-            if ((workInProgress.mode & ProfileMode) !== NoMode) {
-              const isTimedOutSuspense = nextState !== null;
-              if (isTimedOutSuspense) {
-                // Don't count time spent in a timed out Suspense subtree as part of the base duration.
-                const primaryChildFragment = workInProgress.child;
-                if (primaryChildFragment !== null) {
-                  // $FlowFixMe Flow doesn't support type casting in combination with the -= operator
-                  workInProgress.treeBaseDuration -= ((primaryChildFragment.treeBaseDuration: any): number);
-                }
-              }
-            }
-          }
-          return null;
-        } else {
-          // We might have reentered this boundary to hydrate it. If so, we need to reset the hydration
-          // state since we're now exiting out of it. popHydrationState doesn't do that for us.
-          resetHydrationState();
-          if ((workInProgress.flags & DidCapture) === NoFlags) {
-            // This boundary did not suspend so it's now hydrated and unsuspended.
-            workInProgress.memoizedState = null;
-          }
-          // If nothing suspended, we need to schedule an effect to mark this boundary
-          // as having hydrated so events know that they're free to be invoked.
-          // It's also a signal to replay events and the suspense callback.
-          // If something suspended, schedule an effect to attach retry listeners.
-          // So we might as well always mark this.
-          workInProgress.flags |= Update;
-          bubbleProperties(workInProgress);
-          if (enableProfilerTimer) {
-            if ((workInProgress.mode & ProfileMode) !== NoMode) {
-              const isTimedOutSuspense = nextState !== null;
-              if (isTimedOutSuspense) {
-                // Don't count time spent in a timed out Suspense subtree as part of the base duration.
-                const primaryChildFragment = workInProgress.child;
-                if (primaryChildFragment !== null) {
-                  // $FlowFixMe Flow doesn't support type casting in combination with the -= operator
-                  workInProgress.treeBaseDuration -= ((primaryChildFragment.treeBaseDuration: any): number);
-                }
-              }
-            }
+        const fallthroughToNormalSuspensePath = completeDehydratedSuspenseBoundary(
+          current,
+          workInProgress,
+          nextState,
+        );
+        if (!fallthroughToNormalSuspensePath) {
+          if (workInProgress.flags & ShouldCapture) {
+            // Special case. There were remaining unhydrated nodes. We treat
+            // this as a mismatch. Revert to client rendering.
+            return workInProgress;
+          } else {
+            // Did not finish hydrating, either because this is the initial
+            // render or because something suspended.
+            return null;
           }
-          return null;
         }
-      }
 
-      // Successfully completed this tree. If this was a forced client render,
-      // there may have been recoverable errors during first hydration
-      // attempt. If so, add them to a queue so we can log them in the
-      // commit phase.
-      upgradeHydrationErrorsToRecoverable();
+        // Continue with the normal Suspense path.
+      }
 
       if ((workInProgress.flags & DidCapture) !== NoFlags) {
         // Something suspended. Re-render with the fallback children.
@@ -1086,13 +1128,9 @@ function completeWork(
       }
 
       const nextDidTimeout = nextState !== null;
-      let prevDidTimeout = false;
-      if (current === null) {
-        popHydrationState(workInProgress);
-      } else {
-        const prevState: null | SuspenseState = current.memoizedState;
-        prevDidTimeout = prevState !== null;
-      }
+      const prevDidTimeout =
+        current !== null &&
+        (current.memoizedState: null | SuspenseState) !== null;
 
       if (enableCache && nextDidTimeout) {
         const offscreenFiber: Fiber = (workInProgress.child: any);