[FORKED] Hidden trees should capture Suspense

If something suspends inside a hidden tree, it should not affect
anything in the visible part of the UI. This means that Offscreen acts
like a Suspense boundary whenever it's in its hidden state.

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

@@ -719,6 +719,7 @@ export function createFiberFromOffscreen(
  const primaryChildInstance: OffscreenInstance = {
  isHidden: false,
  pendingMarkers: null,
+ retryCache: null,
  transitions: null,
  };
  fiber.stateNode = primaryChildInstance;
@@ -740,6 +741,7 @@ export function createFiberFromLegacyHidden(
  isHidden: false,
  pendingMarkers: null,
  transitions: null,
+ retryCache: null,
  };
  fiber.stateNode = instance;
  return fiber;

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

@@ -719,6 +719,7 @@ export function createFiberFromOffscreen(
  const primaryChildInstance: OffscreenInstance = {
  isHidden: false,
  pendingMarkers: null,
+ retryCache: null,
  transitions: null,
  };
  fiber.stateNode = primaryChildInstance;
@@ -740,6 +741,7 @@ export function createFiberFromLegacyHidden(
  isHidden: false,
  pendingMarkers: null,
  transitions: null,
+ retryCache: null,
  };
  fiber.stateNode = instance;
  return fiber;

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

@@ -174,6 +174,8 @@ import {
  setShallowSuspenseListContext,
  pushPrimaryTreeSuspenseHandler,
  pushFallbackTreeSuspenseHandler,
+ pushOffscreenSuspenseHandler,
+ reuseSuspenseHandlerOnStack,
  popSuspenseHandler,
 } from './ReactFiberSuspenseContext.new';
 import {
@@ -678,6 +680,52 @@ function updateOffscreenComponent(
  (enableLegacyHidden && nextProps.mode === 'unstable-defer-without-hiding')
  ) {
  // Rendering a hidden tree.
+
+ const didSuspend = (workInProgress.flags & DidCapture) !== NoFlags;
+ if (didSuspend) {
+ // Something suspended inside a hidden tree
+
+ // Include the base lanes from the last render
+ const nextBaseLanes =
+ prevState !== null
+ ? mergeLanes(prevState.baseLanes, renderLanes)
+ : renderLanes;
+
+ if (current !== null) {
+ // Reset to the current children
+ let currentChild = (workInProgress.child = current.child);
+
+ // The current render suspended, but there may be other lanes with
+ // pending work. We can't read `childLanes` from the current Offscreen
+ // fiber because we reset it when it was deferred; however, we can read
+ // the pending lanes from the child fibers.
+ let currentChildLanes = NoLanes;
+ while (currentChild !== null) {
+ currentChildLanes = mergeLanes(
+ mergeLanes(currentChildLanes, currentChild.lanes),
+ currentChild.childLanes,
+ );
+ currentChild = currentChild.sibling;
+ }
+ const lanesWeJustAttempted = nextBaseLanes;
+ const remainingChildLanes = removeLanes(
+ currentChildLanes,
+ lanesWeJustAttempted,
+ );
+ workInProgress.childLanes = remainingChildLanes;
+ } else {
+ workInProgress.childLanes = NoLanes;
+ workInProgress.child = null;
+ }
+
+ return deferHiddenOffscreenComponent(
+ current,
+ workInProgress,
+ nextBaseLanes,
+ renderLanes,
+ );
+ }
+
  if ((workInProgress.mode & ConcurrentMode) === NoMode) {
  // In legacy sync mode, don't defer the subtree. Render it now.
  // TODO: Consider how Offscreen should work with transitions in the future
@@ -694,50 +742,28 @@ function updateOffscreenComponent(
  }
  }
  reuseHiddenContextOnStack(workInProgress);
+ pushOffscreenSuspenseHandler(workInProgress);
  } else if (!includesSomeLane(renderLanes, (OffscreenLane: Lane))) {
  // We're hidden, and we're not rendering at Offscreen. We will bail out
  // and resume this tree later.
- let nextBaseLanes = renderLanes;
- if (prevState !== null) {
- // Include the base lanes from the last render
- nextBaseLanes = mergeLanes(nextBaseLanes, prevState.baseLanes);
- }
 
- // Schedule this fiber to re-render at offscreen priority. Then bailout.
+ // Schedule this fiber to re-render at Offscreen priority
  workInProgress.lanes = workInProgress.childLanes = laneToLanes(
  OffscreenLane,
  );
- const nextState: OffscreenState = {
- baseLanes: nextBaseLanes,
- // Save the cache pool so we can resume later.
- cachePool: enableCache ? getOffscreenDeferredCache() : null,
- };
- workInProgress.memoizedState = nextState;
- workInProgress.updateQueue = null;
- if (enableCache) {
- // push the cache pool even though we're going to bail out
- // because otherwise there'd be a context mismatch
- if (current !== null) {
- pushTransition(workInProgress, null, null);
- }
- }
 
- // We're about to bail out, but we need to push this to the stack anyway
- // to avoid a push/pop misalignment.
- reuseHiddenContextOnStack(workInProgress);
+ // Include the base lanes from the last render
+ const nextBaseLanes =
+ prevState !== null
+ ? mergeLanes(prevState.baseLanes, renderLanes)
+ : renderLanes;
 
- if (enableLazyContextPropagation && current !== null) {
- // Since this tree will resume rendering in a separate render, we need
- // to propagate parent contexts now so we don't lose track of which
- // ones changed.
- propagateParentContextChangesToDeferredTree(
- current,
- workInProgress,
- renderLanes,
- );
- }
-
- return null;
+ return deferHiddenOffscreenComponent(
+ current,
+ workInProgress,
+ nextBaseLanes,
+ renderLanes,
+ );
  } else {
  // This is the second render. The surrounding visible content has already
  // committed. Now we resume rendering the hidden tree.
@@ -764,6 +790,7 @@ function updateOffscreenComponent(
  } else {
  reuseHiddenContextOnStack(workInProgress);
  }
+ pushOffscreenSuspenseHandler(workInProgress);
  }
  } else {
  // Rendering a visible tree.
@@ -791,6 +818,7 @@ function updateOffscreenComponent(
 
  // Push the lanes that were skipped when we bailed out.
  pushHiddenContext(workInProgress, prevState);
+ reuseSuspenseHandlerOnStack(workInProgress);
 
  // Since we're not hidden anymore, reset the state
  workInProgress.memoizedState = null;
@@ -811,13 +839,54 @@ function updateOffscreenComponent(
  // We're about to bail out, but we need to push this to the stack anyway
  // to avoid a push/pop misalignment.
  reuseHiddenContextOnStack(workInProgress);
+ reuseSuspenseHandlerOnStack(workInProgress);
  }
  }
 
  reconcileChildren(current, workInProgress, nextChildren, renderLanes);
  return workInProgress.child;
 }
 
+function deferHiddenOffscreenComponent(
+ current: Fiber | null,
+ workInProgress: Fiber,
+ nextBaseLanes: Lanes,
+ renderLanes: Lanes,
+) {
+ const nextState: OffscreenState = {
+ baseLanes: nextBaseLanes,
+ // Save the cache pool so we can resume later.
+ cachePool: enableCache ? getOffscreenDeferredCache() : null,
+ };
+ workInProgress.memoizedState = nextState;
+ if (enableCache) {
+ // push the cache pool even though we're going to bail out
+ // because otherwise there'd be a context mismatch
+ if (current !== null) {
+ pushTransition(workInProgress, null, null);
+ }
+ }
+
+ // We're about to bail out, but we need to push this to the stack anyway
+ // to avoid a push/pop misalignment.
+ reuseHiddenContextOnStack(workInProgress);
+
+ pushOffscreenSuspenseHandler(workInProgress);
+
+ if (enableLazyContextPropagation && current !== null) {
+ // Since this tree will resume rendering in a separate render, we need
+ // to propagate parent contexts now so we don't lose track of which
+ // ones changed.
+ propagateParentContextChangesToDeferredTree(
+ current,
+ workInProgress,
+ renderLanes,
+ );
+ }
+
+ return null;
+}
+
 // Note: These happen to have identical begin phases, for now. We shouldn't hold
 // ourselves to this constraint, though. If the behavior diverges, we should
 // fork the function.

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

@@ -1953,17 +1953,44 @@ function commitSuspenseHydrationCallbacks(
  }
 }
 
+function getRetryCache(finishedWork) {
+ // TODO: Unify the interface for the retry cache so we don't have to switch
+ // on the tag like this.
+ switch (finishedWork.tag) {
+ case SuspenseComponent:
+ case SuspenseListComponent: {
+ let retryCache = finishedWork.stateNode;
+ if (retryCache === null) {
+ retryCache = finishedWork.stateNode = new PossiblyWeakSet();
+ }
+ return retryCache;
+ }
+ case OffscreenComponent: {
+ const instance: OffscreenInstance = finishedWork.stateNode;
+ let retryCache = instance.retryCache;
+ if (retryCache === null) {
+ retryCache = instance.retryCache = new PossiblyWeakSet();
+ }
+ return retryCache;
+ }
+ default: {
+ throw new Error(
+ `Unexpected Suspense handler tag (${finishedWork.tag}). This is a ` +
+ 'bug in React.',
+ );
+ }
+ }
+}
+
 function attachSuspenseRetryListeners(finishedWork: Fiber) {
  // If this boundary just timed out, then it will have a set of wakeables.
  // For each wakeable, attach a listener so that when it resolves, React
  // attempts to re-render the boundary in the primary (pre-timeout) state.
  const wakeables: Set<Wakeable> | null = (finishedWork.updateQueue: any);
  if (wakeables !== null) {
  finishedWork.updateQueue = null;
- let retryCache = finishedWork.stateNode;
- if (retryCache === null) {
- retryCache = finishedWork.stateNode = new PossiblyWeakSet();
- }
+
+ const retryCache = getRetryCache(finishedWork);
  wakeables.forEach(wakeable => {
  // Memoize using the boundary fiber to prevent redundant listeners.
  const retry = resolveRetryWakeable.bind(null, finishedWork, wakeable);
@@ -2362,6 +2389,11 @@ function commitMutationEffectsOnFiber(
  hideOrUnhideAllChildren(offscreenBoundary, isHidden);
  }
  }
+
+ // TODO: Move to passive phase
+ if (flags & Update) {
+ attachSuspenseRetryListeners(finishedWork);
+ }
  return;
  }
  case SuspenseListComponent: {

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

@@ -1508,6 +1508,7 @@ function completeWork(
  }
  case OffscreenComponent:
  case LegacyHiddenComponent: {
+ popSuspenseHandler(workInProgress);
  popHiddenContext(workInProgress);
  const nextState: OffscreenState | null = workInProgress.memoizedState;
  const nextIsHidden = nextState !== null;
@@ -1529,7 +1530,11 @@ function completeWork(
  } else {
  // Don't bubble properties for hidden children unless we're rendering
  // at offscreen priority.
- if (includesSomeLane(renderLanes, (OffscreenLane: Lane))) {
+ if (
+ includesSomeLane(renderLanes, (OffscreenLane: Lane)) &&
+ // Also don't bubble if the tree suspended
+ (workInProgress.flags & DidCapture) === NoLanes
+ ) {
  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
@@ -1544,6 +1549,12 @@ function completeWork(
  }
  }
 
+ if (workInProgress.updateQueue !== null) {
+ // Schedule an effect to attach Suspense retry listeners
+ // TODO: Move to passive phase
+ workInProgress.flags |= Update;
+ }
+
  if (enableCache) {
  let previousCache: Cache | null = null;
  if (

packages/react-reconciler/src/ReactFiberOffscreenComponent.js

@@ -7,7 +7,7 @@
  * @flow
  */
 
-import type {ReactNodeList, OffscreenMode} from 'shared/ReactTypes';
+import type {ReactNodeList, OffscreenMode, Wakeable} from 'shared/ReactTypes';
 import type {Lanes} from './ReactFiberLane.old';
 import type {SpawnedCachePool} from './ReactFiberCacheComponent.new';
 import type {
@@ -46,4 +46,5 @@ export type OffscreenInstance = {|
  isHidden: boolean,
  pendingMarkers: Set<PendingSuspenseBoundaries> | null,
  transitions: Set<Transition> | null,
+ retryCache: WeakSet<Wakeable> | Set<Wakeable> | null,
 |};

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

@@ -14,7 +14,7 @@ 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';
+import {SuspenseComponent, OffscreenComponent} 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.
@@ -77,6 +77,19 @@ 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.
+ reuseSuspenseHandlerOnStack(fiber);
+}
+
+export function pushOffscreenSuspenseHandler(fiber: Fiber): void {
+ if (fiber.tag === OffscreenComponent) {
+ push(suspenseHandlerStackCursor, fiber, fiber);
+ } else {
+ // This is a LegacyHidden component.
+ reuseSuspenseHandlerOnStack(fiber);
+ }
+}
+
+export function reuseSuspenseHandlerOnStack(fiber: Fiber) {
  push(suspenseHandlerStackCursor, getSuspenseHandler(), fiber);
 }