Track nearest Suspense handler on stack (#24585)

* [FORKED] Add HiddenContext to track if subtree is hidden

This adds a new stack cursor for tracking whether we're rendering inside
a subtree that's currently hidden.

This corresponds to the same place where we're already tracking the
"base lanes" needed to reveal a hidden subtree — that is, when going
from hidden -> visible, the base lanes are the ones that we skipped
over when we deferred the subtree. We must includes all the base lanes
and their updates in order to avoid an inconsistency with the
surrounding content that already committed.

I consolidated the base lanes logic and the hidden logic into the same
set of push/pop calls.

This is intended to replace the InvisibleParentContext that is currently
part of SuspenseContext, but I haven't done that part yet.

* Add previous commit to forked revisions

* [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.

* Add previous commit to forked revisions

Author: acdlite

Diff for: packages/react-reconciler/src/ReactFiberBeginWork.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

Diff for: packages/react-reconciler/src/ReactFiberCompleteWork.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;
+}

Diff for: packages/react-reconciler/src/ReactFiberHiddenContext.new.js

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

Diff for: packages/react-reconciler/src/ReactFiberHiddenContext.old.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) {