[FORKED] 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

@@ -37,7 +37,6 @@ import type {UpdateQueue} from './ReactFiberClassUpdateQueue.new';
 import type {RootState} from './ReactFiberRoot.new';
 import type {TracingMarkerInstance} from './ReactFiberTracingMarkerComponent.new';
 import {
-  enableSuspenseAvoidThisFallback,
   enableCPUSuspense,
   enableUseMutableSource,
 } from 'shared/ReactFeatureFlags';
@@ -167,13 +166,14 @@ import {shouldError, shouldSuspend} from './ReactFiberReconciler';
 import {pushHostContext, pushHostContainer} from './ReactFiberHostContext.new';
 import {
   suspenseStackCursor,
-  pushSuspenseContext,
+  pushSuspenseListContext,
-  InvisibleParentSuspenseContext,
   ForceSuspenseFallback,
-  hasSuspenseContext,
+  hasSuspenseListContext,
-  setDefaultShallowSuspenseContext,
+  setDefaultShallowSuspenseListContext,
-  addSubtreeSuspenseContext,
+  setShallowSuspenseListContext,
-  setShallowSuspenseContext,
+  pushPrimaryTreeSuspenseHandler,
+  pushFallbackTreeSuspenseHandler,
+  popSuspenseHandler,
 } from './ReactFiberSuspenseContext.new';
 import {
   pushHiddenContext,
@@ -1969,7 +1969,6 @@ function updateSuspenseOffscreenState(
 
 // TODO: Probably should inline this back
 function shouldRemainOnFallback(
-  suspenseContext: SuspenseContext,
   current: null | Fiber,
   workInProgress: Fiber,
   renderLanes: Lanes,
@@ -1989,7 +1988,8 @@ function shouldRemainOnFallback(
   }
 
   // Not currently showing content. Consult the Suspense context.
-  return hasSuspenseContext(
+  const suspenseContext: SuspenseContext = suspenseStackCursor.current;
+  return hasSuspenseListContext(
     suspenseContext,
     (ForceSuspenseFallback: SuspenseContext),
   );
@@ -2010,50 +2010,18 @@ function updateSuspenseComponent(current, workInProgress, renderLanes) {
     }
   }
 
-  let suspenseContext: SuspenseContext = suspenseStackCursor.current;
-
   let showFallback = false;
   const didSuspend = (workInProgress.flags & DidCapture) !== NoFlags;
-
   if (
     didSuspend ||
-    shouldRemainOnFallback(
+    shouldRemainOnFallback(current, workInProgress, renderLanes)
-      suspenseContext,
-      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.
@@ -2081,24 +2049,40 @@ function updateSuspenseComponent(current, workInProgress, renderLanes) {
 
     // Special path for hydration
     // If we're currently hydrating, try to hydrate this boundary.
-    tryToClaimNextHydratableInstance(workInProgress);
+    if (getIsHydrating()) {
-    // This could've been a dehydrated suspense component.
+      // We must push the suspense handler context *before* attempting to
-    const suspenseState: null | SuspenseState = workInProgress.memoizedState;
+      // hydrate, to avoid a mismatch in case it errors.
-    if (suspenseState !== null) {
+      if (showFallback) {
-      const dehydrated = suspenseState.dehydrated;
+        pushPrimaryTreeSuspenseHandler(workInProgress);
-      if (dehydrated !== null) {
+      } else {
-        return mountDehydratedSuspenseComponent(
+        pushFallbackTreeSuspenseHandler(workInProgress);
-          workInProgress,
-          dehydrated,
-          renderLanes,
-        );
       }
+      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,
@@ -2131,6 +2115,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,
@@ -2154,6 +2139,7 @@ function updateSuspenseComponent(current, workInProgress, renderLanes) {
       workInProgress.lanes = SomeRetryLane;
       return fallbackFragment;
     } else {
+      pushPrimaryTreeSuspenseHandler(workInProgress);
       return mountSuspensePrimaryChildren(
         workInProgress,
         nextPrimaryChildren,
@@ -2181,6 +2167,8 @@ function updateSuspenseComponent(current, workInProgress, renderLanes) {
     }
 
     if (showFallback) {
+      pushFallbackTreeSuspenseHandler(workInProgress);
+
       const nextFallbackChildren = nextProps.fallback;
       const nextPrimaryChildren = nextProps.children;
       const fallbackChildFragment = updateSuspenseFallbackChildren(
@@ -2215,6 +2203,8 @@ function updateSuspenseComponent(current, workInProgress, renderLanes) {
       workInProgress.memoizedState = SUSPENDED_MARKER;
       return fallbackChildFragment;
     } else {
+      pushPrimaryTreeSuspenseHandler(workInProgress);
+
       const nextPrimaryChildren = nextProps.children;
       const primaryChildFragment = updateSuspensePrimaryChildren(
         current,
@@ -2585,6 +2575,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.
@@ -2751,6 +2742,8 @@ function updateDehydratedSuspenseComponent(
 
     if (workInProgress.flags & ForceClientRender) {
       // Something errored during hydration. Try again without hydrating.
+      pushPrimaryTreeSuspenseHandler(workInProgress);
+
       workInProgress.flags &= ~ForceClientRender;
       const capturedValue = createCapturedValue(
         new Error(
@@ -2767,6 +2760,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.
@@ -2775,6 +2772,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(
@@ -3070,12 +3069,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,
     );
@@ -3093,9 +3092,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
@@ -3559,10 +3558,9 @@ function attemptEarlyBailoutIfNoScheduledUpdate(
       const state: SuspenseState | null = workInProgress.memoizedState;
       if (state !== null) {
         if (state.dehydrated !== null) {
-          pushSuspenseContext(
+          // We're not going to render the children, so this is just to maintain
-            workInProgress,
+          // push/pop symmetry
-            setDefaultShallowSuspenseContext(suspenseStackCursor.current),
+          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.
@@ -3585,10 +3583,7 @@ function attemptEarlyBailoutIfNoScheduledUpdate(
         } else {
           // The primary child fragment does not have pending work marked
           // on it
-          pushSuspenseContext(
+          pushPrimaryTreeSuspenseHandler(workInProgress);
-            workInProgress,
-            setDefaultShallowSuspenseContext(suspenseStackCursor.current),
-          );
           // The primary children do not have pending work with sufficient
           // priority. Bailout.
           const child = bailoutOnAlreadyFinishedWork(
@@ -3608,10 +3603,7 @@ function attemptEarlyBailoutIfNoScheduledUpdate(
           }
         }
       } else {
-        pushSuspenseContext(
+        pushPrimaryTreeSuspenseHandler(workInProgress);
-          workInProgress,
-          setDefaultShallowSuspenseContext(suspenseStackCursor.current),
-        );
       }
       break;
     }
@@ -3669,7 +3661,7 @@ function attemptEarlyBailoutIfNoScheduledUpdate(
         renderState.tail = null;
         renderState.lastEffect = null;
       }
-      pushSuspenseContext(workInProgress, suspenseStackCursor.current);
+      pushSuspenseListContext(workInProgress, suspenseStackCursor.current);
 
       if (hasChildWork) {
         break;

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

@@ -27,7 +27,6 @@ import type {
   SuspenseState,
   SuspenseListRenderState,
 } from './ReactFiberSuspenseComponent.new';
-import type {SuspenseContext} from './ReactFiberSuspenseContext.new';
 import type {OffscreenState} from './ReactFiberOffscreenComponent';
 import type {Cache} from './ReactFiberCacheComponent.new';
 import {
@@ -109,15 +108,17 @@ import {
 } from './ReactFiberHostContext.new';
 import {
   suspenseStackCursor,
-  InvisibleParentSuspenseContext,
+  popSuspenseListContext,
-  hasSuspenseContext,
+  popSuspenseHandler,
-  popSuspenseContext,
+  pushSuspenseListContext,
-  pushSuspenseContext,
+  setShallowSuspenseListContext,
-  setShallowSuspenseContext,
   ForceSuspenseFallback,
-  setDefaultShallowSuspenseContext,
+  setDefaultShallowSuspenseListContext,
 } from './ReactFiberSuspenseContext.new';
-import {popHiddenContext} from './ReactFiberHiddenContext.new';
+import {
+  popHiddenContext,
+  isCurrentTreeHidden,
+} from './ReactFiberHiddenContext.new';
 import {findFirstSuspended} from './ReactFiberSuspenseComponent.new';
 import {
   isContextProvider as isLegacyContextProvider,
@@ -1076,7 +1077,7 @@ function completeWork(
       return null;
     }
     case SuspenseComponent: {
-      popSuspenseContext(workInProgress);
+      popSuspenseHandler(workInProgress);
       const nextState: null | SuspenseState = workInProgress.memoizedState;
 
       // Special path for dehydrated boundaries. We may eventually move this
@@ -1185,25 +1186,23 @@ function completeWork(
             // If this render already had a ping or lower pri updates,
             // and this is the first time we know we're going to suspend we
             // should be able to immediately restart from within throwException.
-            const hasInvisibleChildContext =
+
-              current === null &&
+            // Check if this is a "bad" fallback state or a good one. A bad
-              (workInProgress.memoizedProps.unstable_avoidThisFallback !==
+            // fallback state is one that we only show as a last resort; if this
-                true ||
+            // is a transition, we'll block it from displaying, and wait for
-                !enableSuspenseAvoidThisFallback);
+            // more data to arrive.
-            if (
+            const isBadFallback =
-              hasInvisibleChildContext ||
+              // It's bad to switch to a fallback if content is already visible
-              hasSuspenseContext(
+              (current !== null && !prevDidTimeout && !isCurrentTreeHidden()) ||
-                suspenseStackCursor.current,
+              // Experimental: Some fallbacks are always bad
-                (InvisibleParentSuspenseContext: SuspenseContext),
+              (enableSuspenseAvoidThisFallback &&
-              )
+                workInProgress.memoizedProps.unstable_avoidThisFallback ===
-            ) {
+                  true);
-              // If this was in an invisible tree or a new render, then showing
+
-              // this boundary is ok.
+            if (isBadFallback) {
-              renderDidSuspend();
-            } else {
-              // Otherwise, we're going to have to hide content so we should
-              // suspend for longer if possible.
               renderDidSuspendDelayIfPossible();
+            } else {
+              renderDidSuspend();
             }
           }
         }
@@ -1265,7 +1264,7 @@ function completeWork(
       return null;
     }
     case SuspenseListComponent: {
-      popSuspenseContext(workInProgress);
+      popSuspenseListContext(workInProgress);
 
       const renderState: null | SuspenseListRenderState =
         workInProgress.memoizedState;
@@ -1331,11 +1330,11 @@ function completeWork(
                 workInProgress.subtreeFlags = NoFlags;
                 resetChildFibers(workInProgress, renderLanes);
 
-                // Set up the Suspense Context to force suspense and immediately
+                // Set up the Suspense List Context to force suspense and
-                // rerender the children.
+                // immediately rerender the children.
-                pushSuspenseContext(
+                pushSuspenseListContext(
                   workInProgress,
-                  setShallowSuspenseContext(
+                  setShallowSuspenseListContext(
                     suspenseStackCursor.current,
                     ForceSuspenseFallback,
                   ),
@@ -1458,14 +1457,16 @@ function completeWork(
         // setting it the first time we go from not suspended to suspended.
         let suspenseContext = suspenseStackCursor.current;
         if (didSuspendAlready) {
-          suspenseContext = setShallowSuspenseContext(
+          suspenseContext = setShallowSuspenseListContext(
             suspenseContext,
             ForceSuspenseFallback,
           );
         } else {
-          suspenseContext = setDefaultShallowSuspenseContext(suspenseContext);
+          suspenseContext = setDefaultShallowSuspenseListContext(
+            suspenseContext,
+          );
         }
-        pushSuspenseContext(workInProgress, suspenseContext);
+        pushSuspenseListContext(workInProgress, suspenseContext);
         // Do a pass over the next row.
         // Don't bubble properties in this case.
         return next;

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

@@ -64,3 +64,7 @@ export function popHiddenContext(fiber: Fiber): void {
   pop(currentTreeHiddenStackCursor, fiber);
   pop(prevRenderLanesStackCursor, fiber);
 }
+
+export function isCurrentTreeHidden() {
+  return currentTreeHiddenStackCursor.current !== null;
+}

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

@@ -13,7 +13,6 @@ import type {SuspenseInstance} from './ReactFiberHostConfig';
 import type {Lane} from './ReactFiberLane.new';
 import type {TreeContext} from './ReactFiberTreeContext.new';
 
-import {enableSuspenseAvoidThisFallback} from 'shared/ReactFeatureFlags';
 import {SuspenseComponent, SuspenseListComponent} from './ReactWorkTags';
 import {NoFlags, DidCapture} from './ReactFiberFlags';
 import {
@@ -67,37 +66,6 @@ export type SuspenseListRenderState = {|
   tailMode: SuspenseListTailMode,
 |};
 
-export function shouldCaptureSuspense(
-  workInProgress: Fiber,
-  hasInvisibleParent: boolean,
-): boolean {
-  // If it was the primary children that just suspended, capture and render the
-  // fallback. Otherwise, don't capture and bubble to the next boundary.
-  const nextState: SuspenseState | null = workInProgress.memoizedState;
-  if (nextState !== null) {
-    if (nextState.dehydrated !== null) {
-      // A dehydrated boundary always captures.
-      return true;
-    }
-    return false;
-  }
-  const props = workInProgress.memoizedProps;
-  // Regular boundaries always capture.
-  if (
-    !enableSuspenseAvoidThisFallback ||
-    props.unstable_avoidThisFallback !== true
-  ) {
-    return true;
-  }
-  // If it's a boundary we should avoid, then we prefer to bubble up to the
-  // parent boundary if it is currently invisible.
-  if (hasInvisibleParent) {
-    return false;
-  }
-  // If the parent is not able to handle it, we must handle it.
-  return true;
-}
-
 export function findFirstSuspended(row: Fiber): null | Fiber {
   let node = row;
   while (node !== null) {

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

@@ -9,33 +9,96 @@
 
 import type {Fiber} from './ReactInternalTypes';
 import type {StackCursor} from './ReactFiberStack.new';
+import type {SuspenseState} from './ReactFiberSuspenseComponent.new';
 
+import {enableSuspenseAvoidThisFallback} from 'shared/ReactFeatureFlags';
 import {createCursor, push, pop} from './ReactFiberStack.new';
+import {isCurrentTreeHidden} from './ReactFiberHiddenContext.new';
+import {SuspenseComponent} from './ReactWorkTags';
 
+// The Suspense handler is the boundary that should capture if something
+// suspends, i.e. it's the nearest `catch` block on the stack.
+const suspenseHandlerStackCursor: StackCursor<Fiber | null> = createCursor(
+  null,
+);
+
+function shouldAvoidedBoundaryCapture(
+  workInProgress: Fiber,
+  handlerOnStack: Fiber,
+  props: any,
+): boolean {
+  if (enableSuspenseAvoidThisFallback) {
+    // If the parent is already showing content, and we're not inside a hidden
+    // tree, then we should show the avoided fallback.
+    if (handlerOnStack.alternate !== null && !isCurrentTreeHidden()) {
+      return true;
+    }
+
+    // If the handler on the stack is also an avoided boundary, then we should
+    // favor this inner one.
+    if (
+      handlerOnStack.tag === SuspenseComponent &&
+      handlerOnStack.memoizedProps.unstable_avoidThisFallback === true
+    ) {
+      return true;
+    }
+
+    // If this avoided boundary is dehydrated, then it should capture.
+    const suspenseState: SuspenseState | null = workInProgress.memoizedState;
+    if (suspenseState !== null && suspenseState.dehydrated !== null) {
+      return true;
+    }
+  }
+
+  // If none of those cases apply, then we should avoid this fallback and show
+  // the outer one instead.
+  return false;
+}
+
+export function pushPrimaryTreeSuspenseHandler(handler: Fiber): void {
+  const props = handler.pendingProps;
+  const handlerOnStack = suspenseHandlerStackCursor.current;
+  if (
+    enableSuspenseAvoidThisFallback &&
+    props.unstable_avoidThisFallback === true &&
+    handlerOnStack !== null &&
+    !shouldAvoidedBoundaryCapture(handler, handlerOnStack, props)
+  ) {
+    // This boundary should not capture if something suspends. Reuse the
+    // existing handler on the stack.
+    push(suspenseHandlerStackCursor, handlerOnStack, handler);
+  } else {
+    // Push this handler onto the stack.
+    push(suspenseHandlerStackCursor, handler, handler);
+  }
+}
+
+export function pushFallbackTreeSuspenseHandler(fiber: Fiber): void {
+  // We're about to render the fallback. If something in the fallback suspends,
+  // it's akin to throwing inside of a `catch` block. This boundary should not
+  // capture. Reuse the existing handler on the stack.
+  push(suspenseHandlerStackCursor, getSuspenseHandler(), fiber);
+}
+
+export function getSuspenseHandler(): Fiber | null {
+  return suspenseHandlerStackCursor.current;
+}
+
+export function popSuspenseHandler(fiber: Fiber): void {
+  pop(suspenseHandlerStackCursor, fiber);
+}
+
+// SuspenseList context
+// TODO: Move to a separate module? We may change the SuspenseList
+// implementation to hide/show in the commit phase, anyway.
 export opaque type SuspenseContext = number;
 export opaque type SubtreeSuspenseContext: SuspenseContext = number;
 export opaque type ShallowSuspenseContext: SuspenseContext = number;
 
 const DefaultSuspenseContext: SuspenseContext = 0b00;
 
-// The Suspense Context is split into two parts. The lower bits is
-// inherited deeply down the subtree. The upper bits only affect
-// this immediate suspense boundary and gets reset each new
-// boundary or suspense list.
 const SubtreeSuspenseContextMask: SuspenseContext = 0b01;
 
-// Subtree Flags:
-
-// InvisibleParentSuspenseContext indicates that one of our parent Suspense
-// boundaries is not currently showing visible main content.
-// Either because it is already showing a fallback or is not mounted at all.
-// We can use this to determine if it is desirable to trigger a fallback at
-// the parent. If not, then we might need to trigger undesirable boundaries
-// and/or suspend the commit to avoid hiding the parent content.
-export const InvisibleParentSuspenseContext: SubtreeSuspenseContext = 0b01;
-
-// Shallow Flags:
-
 // ForceSuspenseFallback can be used by SuspenseList to force newly added
 // items into their fallback state during one of the render passes.
 export const ForceSuspenseFallback: ShallowSuspenseContext = 0b10;
@@ -44,40 +107,33 @@ export const suspenseStackCursor: StackCursor<SuspenseContext> = createCursor(
   DefaultSuspenseContext,
 );
 
-export function hasSuspenseContext(
+export function hasSuspenseListContext(
   parentContext: SuspenseContext,
   flag: SuspenseContext,
 ): boolean {
   return (parentContext & flag) !== 0;
 }
 
-export function setDefaultShallowSuspenseContext(
+export function setDefaultShallowSuspenseListContext(
   parentContext: SuspenseContext,
 ): SuspenseContext {
   return parentContext & SubtreeSuspenseContextMask;
 }
 
-export function setShallowSuspenseContext(
+export function setShallowSuspenseListContext(
   parentContext: SuspenseContext,
   shallowContext: ShallowSuspenseContext,
 ): SuspenseContext {
   return (parentContext & SubtreeSuspenseContextMask) | shallowContext;
 }
 
-export function addSubtreeSuspenseContext(
+export function pushSuspenseListContext(
-  parentContext: SuspenseContext,
-  subtreeContext: SubtreeSuspenseContext,
-): SuspenseContext {
-  return parentContext | subtreeContext;
-}
-
-export function pushSuspenseContext(
   fiber: Fiber,
   newContext: SuspenseContext,
 ): void {
   push(suspenseStackCursor, newContext, fiber);
 }
 
-export function popSuspenseContext(fiber: Fiber): void {
+export function popSuspenseListContext(fiber: Fiber): void {
   pop(suspenseStackCursor, fiber);
 }

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

@@ -13,13 +13,11 @@ import type {Lane, Lanes} from './ReactFiberLane.new';
 import type {CapturedValue} from './ReactCapturedValue';
 import type {Update} from './ReactFiberClassUpdateQueue.new';
 import type {Wakeable} from 'shared/ReactTypes';
-import type {SuspenseContext} from './ReactFiberSuspenseContext.new';
 
 import getComponentNameFromFiber from 'react-reconciler/src/getComponentNameFromFiber';
 import {
   ClassComponent,
   HostRoot,
-  SuspenseComponent,
   IncompleteClassComponent,
   FunctionComponent,
   ForwardRef,
@@ -34,7 +32,6 @@ import {
   ForceUpdateForLegacySuspense,
   ForceClientRender,
 } from './ReactFiberFlags';
-import {shouldCaptureSuspense} from './ReactFiberSuspenseComponent.new';
 import {NoMode, ConcurrentMode, DebugTracingMode} from './ReactTypeOfMode';
 import {
   enableDebugTracing,
@@ -50,11 +47,7 @@ import {
   enqueueUpdate,
 } from './ReactFiberClassUpdateQueue.new';
 import {markFailedErrorBoundaryForHotReloading} from './ReactFiberHotReloading.new';
-import {
+import {getSuspenseHandler} from './ReactFiberSuspenseContext.new';
-  suspenseStackCursor,
-  InvisibleParentSuspenseContext,
-  hasSuspenseContext,
-} from './ReactFiberSuspenseContext.new';
 import {
   renderDidError,
   renderDidSuspendDelayIfPossible,
@@ -269,26 +262,6 @@ function resetSuspendedComponent(sourceFiber: Fiber, rootRenderLanes: Lanes) {
   }
 }
 
-function getNearestSuspenseBoundaryToCapture(returnFiber: Fiber) {
-  let node = returnFiber;
-  const hasInvisibleParentBoundary = hasSuspenseContext(
-    suspenseStackCursor.current,
-    (InvisibleParentSuspenseContext: SuspenseContext),
-  );
-  do {
-    if (
-      node.tag === SuspenseComponent &&
-      shouldCaptureSuspense(node, hasInvisibleParentBoundary)
-    ) {
-      return node;
-    }
-    // This boundary already captured during this render. Continue to the next
-    // boundary.
-    node = node.return;
-  } while (node !== null);
-  return null;
-}
-
 function markSuspenseBoundaryShouldCapture(
   suspenseBoundary: Fiber,
   returnFiber: Fiber,
@@ -444,7 +417,7 @@ function throwException(
     }
 
     // Schedule the nearest Suspense to re-render the timed out view.
-    const suspenseBoundary = getNearestSuspenseBoundaryToCapture(returnFiber);
+    const suspenseBoundary = getSuspenseHandler();
     if (suspenseBoundary !== null) {
       suspenseBoundary.flags &= ~ForceClientRender;
       markSuspenseBoundaryShouldCapture(
@@ -496,7 +469,7 @@ function throwException(
     // This is a regular error, not a Suspense wakeable.
     if (getIsHydrating() && sourceFiber.mode & ConcurrentMode) {
       markDidThrowWhileHydratingDEV();
-      const suspenseBoundary = getNearestSuspenseBoundaryToCapture(returnFiber);
+      const suspenseBoundary = getSuspenseHandler();
       // If the error was thrown during hydration, we may be able to recover by
       // discarding the dehydrated content and switching to a client render.
       // Instead of surfacing the error, find the nearest Suspense boundary

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

@@ -36,7 +36,10 @@ import {
 } from 'shared/ReactFeatureFlags';
 
 import {popHostContainer, popHostContext} from './ReactFiberHostContext.new';
-import {popSuspenseContext} from './ReactFiberSuspenseContext.new';
+import {
+  popSuspenseListContext,
+  popSuspenseHandler,
+} from './ReactFiberSuspenseContext.new';
 import {popHiddenContext} from './ReactFiberHiddenContext.new';
 import {resetHydrationState} from './ReactFiberHydrationContext.new';
 import {
@@ -109,7 +112,7 @@ function unwindWork(
       return null;
     }
     case SuspenseComponent: {
-      popSuspenseContext(workInProgress);
+      popSuspenseHandler(workInProgress);
       const suspenseState: null | SuspenseState = workInProgress.memoizedState;
       if (suspenseState !== null && suspenseState.dehydrated !== null) {
         if (workInProgress.alternate === null) {
@@ -137,7 +140,7 @@ function unwindWork(
       return null;
     }
     case SuspenseListComponent: {
-      popSuspenseContext(workInProgress);
+      popSuspenseListContext(workInProgress);
       // SuspenseList doesn't actually catch anything. It should've been
       // caught by a nested boundary. If not, it should bubble through.
       return null;
@@ -208,10 +211,10 @@ function unwindInterruptedWork(
       popHostContainer(interruptedWork);
       break;
     case SuspenseComponent:
-      popSuspenseContext(interruptedWork);
+      popSuspenseHandler(interruptedWork);
       break;
     case SuspenseListComponent:
-      popSuspenseContext(interruptedWork);
+      popSuspenseListContext(interruptedWork);
       break;
     case ContextProvider:
       const context: ReactContext<any> = interruptedWork.type._context;