Track nearest Suspense handler on stack

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,14 +108,17 @@ import {
 } from './ReactFiberHostContext.new';
 import {
   suspenseStackCursor,
-  InvisibleParentSuspenseContext,
-  hasSuspenseContext,
-  popSuspenseContext,
-  pushSuspenseContext,
-  setShallowSuspenseContext,
+  popSuspenseListContext,
+  popSuspenseHandler,
+  pushSuspenseListContext,
+  setShallowSuspenseListContext,
   ForceSuspenseFallback,
-  setDefaultShallowSuspenseContext,
+  setDefaultShallowSuspenseListContext,
 } from './ReactFiberSuspenseContext.new';
+import {
+  popHiddenContext,
+  isCurrentTreeHidden,
+} from './ReactFiberHiddenContext.new';
 import {findFirstSuspended} from './ReactFiberSuspenseComponent.new';
 import {
   isContextProvider as isLegacyContextProvider,
@@ -146,9 +148,7 @@ import {
   renderDidSuspend,
   renderDidSuspendDelayIfPossible,
   renderHasNotSuspendedYet,
-  popRenderLanes,
   getRenderTargetTime,
-  subtreeRenderLanes,
   getWorkInProgressTransitions,
 } from './ReactFiberWorkLoop.new';
 import {
@@ -1077,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
@@ -1186,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 &&
-              (workInProgress.memoizedProps.unstable_avoidThisFallback !==
-                true ||
-                !enableSuspenseAvoidThisFallback);
-            if (
-              hasInvisibleChildContext ||
-              hasSuspenseContext(
-                suspenseStackCursor.current,
-                (InvisibleParentSuspenseContext: SuspenseContext),
-              )
-            ) {
-              // If this was in an invisible tree or a new render, then showing
-              // this boundary is ok.
-              renderDidSuspend();
-            } else {
-              // Otherwise, we're going to have to hide content so we should
-              // suspend for longer if possible.
+
+            // Check if this is a "bad" fallback state or a good one. A bad
+            // fallback state is one that we only show as a last resort; if this
+            // is a transition, we'll block it from displaying, and wait for
+            // more data to arrive.
+            const isBadFallback =
+              // It's bad to switch to a fallback if content is already visible
+              (current !== null && !prevDidTimeout && !isCurrentTreeHidden()) ||
+              // Experimental: Some fallbacks are always bad
+              (enableSuspenseAvoidThisFallback &&
+                workInProgress.memoizedProps.unstable_avoidThisFallback ===
+                  true);
+
+            if (isBadFallback) {
               renderDidSuspendDelayIfPossible();
+            } else {
+              renderDidSuspend();
             }
           }
         }
@@ -1266,7 +1264,7 @@ function completeWork(
       return null;
     }
     case SuspenseListComponent: {
-      popSuspenseContext(workInProgress);
+      popSuspenseListContext(workInProgress);
 
       const renderState: null | SuspenseListRenderState =
         workInProgress.memoizedState;
@@ -1332,11 +1330,11 @@ function completeWork(
                 workInProgress.subtreeFlags = NoFlags;
                 resetChildFibers(workInProgress, renderLanes);
 
-                // Set up the Suspense Context to force suspense and immediately
-                // rerender the children.
-                pushSuspenseContext(
+                // Set up the Suspense List Context to force suspense and
+                // immediately rerender the children.
+                pushSuspenseListContext(
                   workInProgress,
-                  setShallowSuspenseContext(
+                  setShallowSuspenseListContext(
                     suspenseStackCursor.current,
                     ForceSuspenseFallback,
                   ),
@@ -1459,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;
@@ -1499,7 +1499,7 @@ function completeWork(
     }
     case OffscreenComponent:
     case LegacyHiddenComponent: {
-      popRenderLanes(workInProgress);
+      popHiddenContext(workInProgress);
       const nextState: OffscreenState | null = workInProgress.memoizedState;
       const nextIsHidden = nextState !== null;
 
@@ -1520,7 +1520,7 @@ function completeWork(
       } else {
         // Don't bubble properties for hidden children unless we're rendering
         // at offscreen priority.
-        if (includesSomeLane(subtreeRenderLanes, (OffscreenLane: Lane))) {
+        if (includesSomeLane(renderLanes, (OffscreenLane: Lane))) {
           bubbleProperties(workInProgress);
           // Check if there was an insertion or update in the hidden subtree.
           // If so, we need to hide those nodes in the commit phase, so

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

@@ -0,0 +1,70 @@
+/**
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ *
+ * @flow
+ */
+
+import type {Fiber} from './ReactInternalTypes';
+import type {StackCursor} from './ReactFiberStack.new';
+import type {Lanes} from './ReactFiberLane.new';
+
+import {createCursor, push, pop} from './ReactFiberStack.new';
+
+import {getRenderLanes, setRenderLanes} from './ReactFiberWorkLoop.new';
+import {NoLanes, mergeLanes} from './ReactFiberLane.new';
+
+// TODO: Remove `renderLanes` context in favor of hidden context
+type HiddenContext = {
+  // Represents the lanes that must be included when processing updates in
+  // order to reveal the hidden content.
+  // TODO: Remove `subtreeLanes` context from work loop in favor of this one.
+  baseLanes: number,
+};
+
+// TODO: This isn't being used yet, but it's intended to replace the
+// InvisibleParentContext that is currently managed by SuspenseContext.
+export const currentTreeHiddenStackCursor: StackCursor<HiddenContext | null> = createCursor(
+  null,
+);
+export const prevRenderLanesStackCursor: StackCursor<Lanes> = createCursor(
+  NoLanes,
+);
+
+export function pushHiddenContext(fiber: Fiber, context: HiddenContext): void {
+  const prevRenderLanes = getRenderLanes();
+  push(prevRenderLanesStackCursor, prevRenderLanes, fiber);
+  push(currentTreeHiddenStackCursor, context, fiber);
+
+  // When rendering a subtree that's currently hidden, we must include all
+  // lanes that would have rendered if the hidden subtree hadn't been deferred.
+  // That is, in order to reveal content from hidden -> visible, we must commit
+  // all the updates that we skipped when we originally hid the tree.
+  setRenderLanes(mergeLanes(prevRenderLanes, context.baseLanes));
+}
+
+export function reuseHiddenContextOnStack(fiber: Fiber): void {
+  // This subtree is not currently hidden, so we don't need to add any lanes
+  // to the render lanes. But we still need to push something to avoid a
+  // context mismatch. Reuse the existing context on the stack.
+  push(prevRenderLanesStackCursor, getRenderLanes(), fiber);
+  push(
+    currentTreeHiddenStackCursor,
+    currentTreeHiddenStackCursor.current,
+    fiber,
+  );
+}
+
+export function popHiddenContext(fiber: Fiber): void {
+  // Restore the previous render lanes from the stack
+  setRenderLanes(prevRenderLanesStackCursor.current);
+
+  pop(currentTreeHiddenStackCursor, fiber);
+  pop(prevRenderLanesStackCursor, fiber);
+}
+
+export function isCurrentTreeHidden() {
+  return currentTreeHiddenStackCursor.current !== null;
+}

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

@@ -0,0 +1 @@
+// Intentionally blank. File only exists in new reconciler fork.

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) {