Move the context out of SuspenseComponent

Author: sebmarkbage

packages/react-reconciler/src/ReactFiberSuspenseComponent.js

@@ -10,18 +10,14 @@
 import type {Fiber} from './ReactFiber';
 import type {ExpirationTime} from './ReactFiberExpirationTime';
 
-import type {SuspenseContext} from './ReactFiberSuspenseContext';
-import {
-  suspenseStackCursor,
-  InvisibleParentSuspenseContext,
-  hasSuspenseContext,
-} from './ReactFiberSuspenseContext';
-
 export type SuspenseState = {|
   fallbackExpirationTime: ExpirationTime,
 |};
 
-export function shouldCaptureSuspense(workInProgress: Fiber): boolean {
+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;
@@ -39,12 +35,7 @@ export function shouldCaptureSuspense(workInProgress: Fiber): boolean {
   }
   // If it's a boundary we should avoid, then we prefer to bubble up to the
   // parent boundary if it is currently invisible.
-  if (
-    hasSuspenseContext(
-      suspenseStackCursor.current,
-      (InvisibleParentSuspenseContext: SuspenseContext),
-    )
-  ) {
+  if (hasInvisibleParent) {
     return false;
   }
   // If the parent is not able to handle it, we must handle it.

packages/react-reconciler/src/ReactFiberUnwindWork.js

@@ -13,6 +13,7 @@ import type {ExpirationTime} from './ReactFiberExpirationTime';
 import type {CapturedValue} from './ReactCapturedValue';
 import type {Update} from './ReactUpdateQueue';
 import type {Thenable} from './ReactFiberScheduler';
+import type {SuspenseContext} from './ReactFiberSuspenseContext';
 
 import {unstable_wrap as Schedule_tracing_wrap} from 'scheduler/tracing';
 import getComponentName from 'shared/getComponentName';
@@ -55,7 +56,12 @@ import {
 import {logError} from './ReactFiberCommitWork';
 import {getStackByFiberInDevAndProd} from './ReactCurrentFiber';
 import {popHostContainer, popHostContext} from './ReactFiberHostContext';
-import {popSuspenseContext} from './ReactFiberSuspenseContext';
+import {
+  suspenseStackCursor,
+  InvisibleParentSuspenseContext,
+  hasSuspenseContext,
+  popSuspenseContext,
+} from './ReactFiberSuspenseContext';
 import {
   isContextProvider as isLegacyContextProvider,
   popContext as popLegacyContext,
@@ -207,15 +213,17 @@ function throwException(
 
     checkForWrongSuspensePriorityInDEV(sourceFiber);
 
-    // TODO: If we're not in an invisible subtree, then we need to mark this render
-    // as needing to suspend for longer to avoid showing this fallback state.
+    let hasInvisibleParentBoundary = hasSuspenseContext(
+      suspenseStackCursor.current,
+      (InvisibleParentSuspenseContext: SuspenseContext),
+    );
 
     // Schedule the nearest Suspense to re-render the timed out view.
     let workInProgress = returnFiber;
     do {
       if (
         workInProgress.tag === SuspenseComponent &&
-        shouldCaptureSuspense(workInProgress)
+        shouldCaptureSuspense(workInProgress, hasInvisibleParentBoundary)
       ) {
         // Found the nearest boundary.
 
@@ -278,6 +286,13 @@ function throwException(
 
         workInProgress.effectTag |= ShouldCapture;
         workInProgress.expirationTime = renderExpirationTime;
+
+        if (!hasInvisibleParentBoundary) {
+          // TODO: If we're not in an invisible subtree, then we need to mark this render
+          // pass as needing to suspend for longer to avoid showing this fallback state.
+          // We could do it here or when we render the fallback.
+        }
+
         return;
       } else if (
         enableSuspenseServerRenderer &&