Track nearest Suspense handler on stack

Instead of traversing the return path whenever something suspends to
find the nearest Suspense boundary, we can push the Suspense boundary
onto the stack before entering its subtree. This doesn't affect the
overall algorithm that much, but because we already do all the same
logic in the begin phase, we can save some redundant work by tracking
that information on the stack instead of recomputing it every time.

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

@@ -36,7 +36,6 @@ import type {
 import type {UpdateQueue} from './ReactUpdateQueue.new';
 import type {RootState} from './ReactFiberRoot.new';
 import {
-  enableSuspenseAvoidThisFallback,
   enableCPUSuspense,
   enableUseMutableSource,
 } from 'shared/ReactFeatureFlags';
@@ -166,13 +165,14 @@ import {shouldError, shouldSuspend} from './ReactFiberReconciler';
 import {pushHostContext, pushHostContainer} from './ReactFiberHostContext.new';
 import {
   suspenseStackCursor,
-  pushSuspenseContext,
-  InvisibleParentSuspenseContext,
+  pushSuspenseListContext,
   ForceSuspenseFallback,
-  hasSuspenseContext,
-  setDefaultShallowSuspenseContext,
-  addSubtreeSuspenseContext,
-  setShallowSuspenseContext,
+  hasSuspenseListContext,
+  setDefaultShallowSuspenseListContext,
+  setShallowSuspenseListContext,
+  pushPrimaryTreeSuspenseHandler,
+  pushFallbackTreeSuspenseHandler,
+  popSuspenseHandler,
 } from './ReactFiberSuspenseContext.new';
 import {
   pushHiddenContext,
@@ -1940,7 +1940,6 @@ function updateSuspenseOffscreenState(
 
 // TODO: Probably should inline this back
 function shouldRemainOnFallback(
-  suspenseContext: SuspenseContext,
   current: null | Fiber,
   workInProgress: Fiber,
   renderLanes: Lanes,
@@ -1960,7 +1959,8 @@ function shouldRemainOnFallback(
   }
 
   // Not currently showing content. Consult the Suspense context.
-  return hasSuspenseContext(
+  const suspenseContext: SuspenseContext = suspenseStackCursor.current;
+  return hasSuspenseListContext(
     suspenseContext,
     (ForceSuspenseFallback: SuspenseContext),
   );
@@ -1981,50 +1981,18 @@ function updateSuspenseComponent(current, workInProgress, renderLanes) {
     }
   }
 
-  let suspenseContext: SuspenseContext = suspenseStackCursor.current;
-
   let showFallback = false;
   const didSuspend = (workInProgress.flags & DidCapture) !== NoFlags;
-
   if (
     didSuspend ||
-    shouldRemainOnFallback(
-      suspenseContext,
-      current,
-      workInProgress,
-      renderLanes,
-    )
+    shouldRemainOnFallback(current, workInProgress, renderLanes)
   ) {
     // Something in this boundary's subtree already suspended. Switch to
     // rendering the fallback children.
     showFallback = true;
     workInProgress.flags &= ~DidCapture;
-  } else {
-    // Attempting the main content
-    if (
-      current === null ||
-      (current.memoizedState: null | SuspenseState) !== null
-    ) {
-      // This is a new mount or this boundary is already showing a fallback state.
-      // Mark this subtree context as having at least one invisible parent that could
-      // handle the fallback state.
-      // Avoided boundaries are not considered since they cannot handle preferred fallback states.
-      if (
-        !enableSuspenseAvoidThisFallback ||
-        nextProps.unstable_avoidThisFallback !== true
-      ) {
-        suspenseContext = addSubtreeSuspenseContext(
-          suspenseContext,
-          InvisibleParentSuspenseContext,
-        );
-      }
-    }
   }
 
-  suspenseContext = setDefaultShallowSuspenseContext(suspenseContext);
-
-  pushSuspenseContext(workInProgress, suspenseContext);
-
   // OK, the next part is confusing. We're about to reconcile the Suspense
   // boundary's children. This involves some custom reconciliation logic. Two
   // main reasons this is so complicated.
@@ -2052,24 +2020,40 @@ function updateSuspenseComponent(current, workInProgress, renderLanes) {
 
     // Special path for hydration
     // If we're currently hydrating, try to hydrate this boundary.
-    tryToClaimNextHydratableInstance(workInProgress);
-    // This could've been a dehydrated suspense component.
-    const suspenseState: null | SuspenseState = workInProgress.memoizedState;
-    if (suspenseState !== null) {
-      const dehydrated = suspenseState.dehydrated;
-      if (dehydrated !== null) {
-        return mountDehydratedSuspenseComponent(
-          workInProgress,
-          dehydrated,
-          renderLanes,
-        );
+    if (getIsHydrating()) {
+      // We must push the suspense handler context *before* attempting to
+      // hydrate, to avoid a mismatch in case it errors.
+      if (showFallback) {
+        pushPrimaryTreeSuspenseHandler(workInProgress);
+      } else {
+        pushFallbackTreeSuspenseHandler(workInProgress);
+      }
+      tryToClaimNextHydratableInstance(workInProgress);
+      // This could've been a dehydrated suspense component.
+      const suspenseState: null | SuspenseState = workInProgress.memoizedState;
+      if (suspenseState !== null) {
+        const dehydrated = suspenseState.dehydrated;
+        if (dehydrated !== null) {
+          return mountDehydratedSuspenseComponent(
+            workInProgress,
+            dehydrated,
+            renderLanes,
+          );
+        }
       }
+      // If hydration didn't succeed, fall through to the normal Suspense path.
+      // To avoid a stack mismatch we need to pop the Suspense handler that we
+      // pushed above. This will become less awkward when move the hydration
+      // logic to its own fiber.
+      popSuspenseHandler(workInProgress);
     }
 
     const nextPrimaryChildren = nextProps.children;
     const nextFallbackChildren = nextProps.fallback;
 
     if (showFallback) {
+      pushFallbackTreeSuspenseHandler(workInProgress);
+
       const fallbackFragment = mountSuspenseFallbackChildren(
         workInProgress,
         nextPrimaryChildren,
@@ -2099,6 +2083,7 @@ function updateSuspenseComponent(current, workInProgress, renderLanes) {
       // This is a CPU-bound tree. Skip this tree and show a placeholder to
       // unblock the surrounding content. Then immediately retry after the
       // initial commit.
+      pushFallbackTreeSuspenseHandler(workInProgress);
       const fallbackFragment = mountSuspenseFallbackChildren(
         workInProgress,
         nextPrimaryChildren,
@@ -2122,6 +2107,7 @@ function updateSuspenseComponent(current, workInProgress, renderLanes) {
       workInProgress.lanes = SomeRetryLane;
       return fallbackFragment;
     } else {
+      pushPrimaryTreeSuspenseHandler(workInProgress);
       return mountSuspensePrimaryChildren(
         workInProgress,
         nextPrimaryChildren,
@@ -2149,6 +2135,8 @@ function updateSuspenseComponent(current, workInProgress, renderLanes) {
     }
 
     if (showFallback) {
+      pushFallbackTreeSuspenseHandler(workInProgress);
+
       const nextFallbackChildren = nextProps.fallback;
       const nextPrimaryChildren = nextProps.children;
       const fallbackChildFragment = updateSuspenseFallbackChildren(
@@ -2181,6 +2169,8 @@ function updateSuspenseComponent(current, workInProgress, renderLanes) {
       workInProgress.memoizedState = SUSPENDED_MARKER;
       return fallbackChildFragment;
     } else {
+      pushPrimaryTreeSuspenseHandler(workInProgress);
+
       const nextPrimaryChildren = nextProps.children;
       const primaryChildFragment = updateSuspensePrimaryChildren(
         current,
@@ -2551,6 +2541,7 @@ function updateDehydratedSuspenseComponent(
 ): null | Fiber {
   if (!didSuspend) {
     // This is the first render pass. Attempt to hydrate.
+    pushPrimaryTreeSuspenseHandler(workInProgress);
 
     // We should never be hydrating at this point because it is the first pass,
     // but after we've already committed once.
@@ -2694,6 +2685,8 @@ function updateDehydratedSuspenseComponent(
 
     if (workInProgress.flags & ForceClientRender) {
       // Something errored during hydration. Try again without hydrating.
+      pushPrimaryTreeSuspenseHandler(workInProgress);
+
       workInProgress.flags &= ~ForceClientRender;
       return retrySuspenseComponentWithoutHydrating(
         current,
@@ -2707,6 +2700,10 @@ function updateDehydratedSuspenseComponent(
     } else if ((workInProgress.memoizedState: null | SuspenseState) !== null) {
       // Something suspended and we should still be in dehydrated mode.
       // Leave the existing child in place.
+
+      // Push to avoid a mismatch
+      pushFallbackTreeSuspenseHandler(workInProgress);
+
       workInProgress.child = current.child;
       // The dehydrated completion pass expects this flag to be there
       // but the normal suspense pass doesn't.
@@ -2715,6 +2712,8 @@ function updateDehydratedSuspenseComponent(
     } else {
       // Suspended but we should no longer be in dehydrated mode.
       // Therefore we now have to render the fallback.
+      pushFallbackTreeSuspenseHandler(workInProgress);
+
       const nextPrimaryChildren = nextProps.children;
       const nextFallbackChildren = nextProps.fallback;
       const fallbackChildFragment = mountSuspenseFallbackAfterRetryWithoutHydrating(
@@ -3010,12 +3009,12 @@ function updateSuspenseListComponent(
 
   let suspenseContext: SuspenseContext = suspenseStackCursor.current;
 
-  const shouldForceFallback = hasSuspenseContext(
+  const shouldForceFallback = hasSuspenseListContext(
     suspenseContext,
     (ForceSuspenseFallback: SuspenseContext),
   );
   if (shouldForceFallback) {
-    suspenseContext = setShallowSuspenseContext(
+    suspenseContext = setShallowSuspenseListContext(
       suspenseContext,
       ForceSuspenseFallback,
     );
@@ -3033,9 +3032,9 @@ function updateSuspenseListComponent(
         renderLanes,
       );
     }
-    suspenseContext = setDefaultShallowSuspenseContext(suspenseContext);
+    suspenseContext = setDefaultShallowSuspenseListContext(suspenseContext);
   }
-  pushSuspenseContext(workInProgress, suspenseContext);
+  pushSuspenseListContext(workInProgress, suspenseContext);
 
   if ((workInProgress.mode & ConcurrentMode) === NoMode) {
     // In legacy mode, SuspenseList doesn't work so we just
@@ -3499,10 +3498,9 @@ function attemptEarlyBailoutIfNoScheduledUpdate(
       const state: SuspenseState | null = workInProgress.memoizedState;
       if (state !== null) {
         if (state.dehydrated !== null) {
-          pushSuspenseContext(
-            workInProgress,
-            setDefaultShallowSuspenseContext(suspenseStackCursor.current),
-          );
+          // We're not going to render the children, so this is just to maintain
+          // push/pop symmetry
+          pushPrimaryTreeSuspenseHandler(workInProgress);
           // We know that this component will suspend again because if it has
           // been unsuspended it has committed as a resolved Suspense component.
           // If it needs to be retried, it should have work scheduled on it.
@@ -3525,10 +3523,7 @@ function attemptEarlyBailoutIfNoScheduledUpdate(
         } else {
           // The primary child fragment does not have pending work marked
           // on it
-          pushSuspenseContext(
-            workInProgress,
-            setDefaultShallowSuspenseContext(suspenseStackCursor.current),
-          );
+          pushPrimaryTreeSuspenseHandler(workInProgress);
           // The primary children do not have pending work with sufficient
           // priority. Bailout.
           const child = bailoutOnAlreadyFinishedWork(
@@ -3548,10 +3543,7 @@ function attemptEarlyBailoutIfNoScheduledUpdate(
           }
         }
       } else {
-        pushSuspenseContext(
-          workInProgress,
-          setDefaultShallowSuspenseContext(suspenseStackCursor.current),
-        );
+        pushPrimaryTreeSuspenseHandler(workInProgress);
       }
       break;
     }
@@ -3609,7 +3601,7 @@ function attemptEarlyBailoutIfNoScheduledUpdate(
         renderState.tail = null;
         renderState.lastEffect = null;
       }
-      pushSuspenseContext(workInProgress, suspenseStackCursor.current);
+      pushSuspenseListContext(workInProgress, suspenseStackCursor.current);
 
       if (hasChildWork) {
         break;