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

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

@@ -175,6 +175,10 @@ import {
   addSubtreeSuspenseContext,
   setShallowSuspenseContext,
 } from './ReactFiberSuspenseContext.new';
+import {
+  pushHiddenContext,
+  reuseHiddenContextOnStack,
+} from './ReactFiberHiddenContext.new';
 import {findFirstSuspended} from './ReactFiberSuspenseComponent.new';
 import {
   pushProvider,
@@ -232,7 +236,6 @@ import {
   renderDidSuspendDelayIfPossible,
   markSkippedUpdateLanes,
   getWorkInProgressRoot,
-  pushRenderLanes,
 } from './ReactFiberWorkLoop.new';
 import {enqueueConcurrentRenderForLane} from './ReactFiberConcurrentUpdates.new';
 import {setWorkInProgressVersion} from './ReactMutableSource.new';
@@ -688,21 +691,14 @@ function updateOffscreenComponent(
           pushTransition(workInProgress, null, null);
         }
       }
-      pushRenderLanes(workInProgress, renderLanes);
+      reuseHiddenContextOnStack(workInProgress);
     } else if (!includesSomeLane(renderLanes, (OffscreenLane: Lane))) {
-      let spawnedCachePool: SpawnedCachePool | null = null;
       // We're hidden, and we're not rendering at Offscreen. We will bail out
       // and resume this tree later.
-      let nextBaseLanes;
+      let nextBaseLanes = renderLanes;
       if (prevState !== null) {
-        const prevBaseLanes = prevState.baseLanes;
-        nextBaseLanes = mergeLanes(prevBaseLanes, renderLanes);
-        if (enableCache) {
-          // Save the cache pool so we can resume later.
-          spawnedCachePool = getOffscreenDeferredCache();
-        }
-      } else {
-        nextBaseLanes = renderLanes;
+        // Include the base lanes from the last render
+        nextBaseLanes = mergeLanes(nextBaseLanes, prevState.baseLanes);
       }
 
       // Schedule this fiber to re-render at offscreen priority. Then bailout.
@@ -711,7 +707,8 @@ function updateOffscreenComponent(
       );
       const nextState: OffscreenState = {
         baseLanes: nextBaseLanes,
-        cachePool: spawnedCachePool,
+        // Save the cache pool so we can resume later.
+        cachePool: enableCache ? getOffscreenDeferredCache() : null,
       };
       workInProgress.memoizedState = nextState;
       workInProgress.updateQueue = null;
@@ -725,7 +722,7 @@ function updateOffscreenComponent(
 
       // We're about to bail out, but we need to push this to the stack anyway
       // to avoid a push/pop misalignment.
-      pushRenderLanes(workInProgress, nextBaseLanes);
+      reuseHiddenContextOnStack(workInProgress);
 
       if (enableLazyContextPropagation && current !== null) {
         // Since this tree will resume rendering in a separate render, we need
@@ -749,9 +746,6 @@ function updateOffscreenComponent(
         cachePool: null,
       };
       workInProgress.memoizedState = nextState;
-      // Push the lanes that were skipped when we bailed out.
-      const subtreeRenderLanes =
-        prevState !== null ? prevState.baseLanes : renderLanes;
       if (enableCache && current !== null) {
         // If the render that spawned this one accessed the cache pool, resume
         // using the same cache. Unless the parent changed, since that means
@@ -762,16 +756,17 @@ function updateOffscreenComponent(
         pushTransition(workInProgress, prevCachePool, null);
       }
 
-      pushRenderLanes(workInProgress, subtreeRenderLanes);
+      // Push the lanes that were skipped when we bailed out.
+      if (prevState !== null) {
+        pushHiddenContext(workInProgress, prevState);
+      } else {
+        reuseHiddenContextOnStack(workInProgress);
+      }
     }
   } else {
     // Rendering a visible tree.
-    let subtreeRenderLanes;
     if (prevState !== null) {
       // We're going from hidden -> visible.
-
-      subtreeRenderLanes = mergeLanes(prevState.baseLanes, renderLanes);
-
       let prevCachePool = null;
       if (enableCache) {
         // If the render that spawned this one accessed the cache pool, resume
@@ -789,13 +784,15 @@ function updateOffscreenComponent(
 
       pushTransition(workInProgress, prevCachePool, transitions);
 
+      // Push the lanes that were skipped when we bailed out.
+      pushHiddenContext(workInProgress, prevState);
+
       // Since we're not hidden anymore, reset the state
       workInProgress.memoizedState = null;
     } else {
       // We weren't previously hidden, and we still aren't, so there's nothing
       // special to do. Need to push to the stack regardless, though, to avoid
       // a push/pop misalignment.
-      subtreeRenderLanes = renderLanes;
 
       if (enableCache) {
         // If the render that spawned this one accessed the cache pool, resume
@@ -805,8 +802,11 @@ function updateOffscreenComponent(
           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);
     }
-    pushRenderLanes(workInProgress, subtreeRenderLanes);
   }
 
   reconcileChildren(current, workInProgress, nextChildren, renderLanes);

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

@@ -117,6 +117,7 @@ import {
   ForceSuspenseFallback,
   setDefaultShallowSuspenseContext,
 } from './ReactFiberSuspenseContext.new';
+import {popHiddenContext} from './ReactFiberHiddenContext.new';
 import {findFirstSuspended} from './ReactFiberSuspenseComponent.new';
 import {
   isContextProvider as isLegacyContextProvider,
@@ -146,9 +147,7 @@ import {
   renderDidSuspend,
   renderDidSuspendDelayIfPossible,
   renderHasNotSuspendedYet,
-  popRenderLanes,
   getRenderTargetTime,
-  subtreeRenderLanes,
   getWorkInProgressTransitions,
 } from './ReactFiberWorkLoop.new';
 import {
@@ -1499,7 +1498,7 @@ function completeWork(
     }
     case OffscreenComponent:
     case LegacyHiddenComponent: {
-      popRenderLanes(workInProgress);
+      popHiddenContext(workInProgress);
       const nextState: OffscreenState | null = workInProgress.memoizedState;
       const nextIsHidden = nextState !== null;
 
@@ -1520,7 +1519,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,66 @@
+/**
+ * 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);
+}

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

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

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

@@ -37,6 +37,7 @@ import {
 
 import {popHostContainer, popHostContext} from './ReactFiberHostContext.new';
 import {popSuspenseContext} from './ReactFiberSuspenseContext.new';
+import {popHiddenContext} from './ReactFiberHiddenContext.new';
 import {resetHydrationState} from './ReactFiberHydrationContext.new';
 import {
   isContextProvider as isLegacyContextProvider,
@@ -44,7 +45,6 @@ import {
   popTopLevelContextObject as popTopLevelLegacyContextObject,
 } from './ReactFiberContext.new';
 import {popProvider} from './ReactFiberNewContext.new';
-import {popRenderLanes} from './ReactFiberWorkLoop.new';
 import {popCacheProvider} from './ReactFiberCacheComponent.new';
 import {transferActualDuration} from './ReactProfilerTimer.new';
 import {popTreeContext} from './ReactFiberTreeContext.new';
@@ -151,7 +151,7 @@ function unwindWork(
       return null;
     case OffscreenComponent:
     case LegacyHiddenComponent:
-      popRenderLanes(workInProgress);
+      popHiddenContext(workInProgress);
       popTransition(workInProgress, current);
       return null;
     case CacheComponent:
@@ -219,7 +219,7 @@ function unwindInterruptedWork(
       break;
     case OffscreenComponent:
     case LegacyHiddenComponent:
-      popRenderLanes(interruptedWork);
+      popHiddenContext(interruptedWork);
       popTransition(interruptedWork, current);
       break;
     case CacheComponent:

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

@@ -11,7 +11,6 @@ import type {Wakeable} from 'shared/ReactTypes';
 import type {Fiber, FiberRoot} from './ReactInternalTypes';
 import type {Lanes, Lane} from './ReactFiberLane.new';
 import type {SuspenseState} from './ReactFiberSuspenseComponent.new';
-import type {StackCursor} from './ReactFiberStack.new';
 import type {Flags} from './ReactFiberFlags';
 import type {FunctionComponentUpdateQueue} from './ReactFiberHooks.new';
 import type {EventPriority} from './ReactEventPriorities.new';
@@ -192,11 +191,6 @@ import {
   createCapturedValueAtFiber,
   type CapturedValue,
 } from './ReactCapturedValue';
-import {
-  push as pushToStack,
-  pop as popFromStack,
-  createCursor,
-} from './ReactFiberStack.new';
 import {
   enqueueConcurrentRenderForLane,
   finishQueueingConcurrentUpdates,
@@ -285,26 +279,20 @@ let workInProgress: Fiber | null = null;
 // The lanes we're rendering
 let workInProgressRootRenderLanes: Lanes = NoLanes;
 
-// Stack that allows components to change the render lanes for its subtree
-// This is a superset of the lanes we started working on at the root. The only
-// case where it's different from `workInProgressRootRenderLanes` is when we
-// enter a subtree that is hidden and needs to be unhidden: Suspense and
-// Offscreen component.
+// A contextual version of workInProgressRootRenderLanes. It is a superset of
+// the lanes that we started working on at the root. When we enter a subtree
+// that is currently hidden, we add the lanes that would have committed if
+// the hidden tree hadn't been deferred. This is modified by the
+// HiddenContext module.
 //
 // Most things in the work loop should deal with workInProgressRootRenderLanes.
-// Most things in begin/complete phases should deal with subtreeRenderLanes.
-export let subtreeRenderLanes: Lanes = NoLanes;
-const subtreeRenderLanesCursor: StackCursor<Lanes> = createCursor(NoLanes);
+// Most things in begin/complete phases should deal with renderLanes.
+export let renderLanes: Lanes = NoLanes;
 
 // Whether to root completed, errored, suspended, etc.
 let workInProgressRootExitStatus: RootExitStatus = RootInProgress;
 // A fatal error, if one is thrown
 let workInProgressRootFatalError: mixed = null;
-// "Included" lanes refer to lanes that were worked on during this render. It's
-// slightly different than `renderLanes` because `renderLanes` can change as you
-// enter and exit an Offscreen tree. This value is the combination of all render
-// lanes for the entire render phase.
-let workInProgressRootIncludedLanes: Lanes = NoLanes;
 // The work left over by components that were visited during this render. Only
 // includes unprocessed updates, not work in bailed out children.
 let workInProgressRootSkippedLanes: Lanes = NoLanes;
@@ -1455,18 +1443,16 @@ export function flushControlled(fn: () => mixed): void {
   }
 }
 
-export function pushRenderLanes(fiber: Fiber, lanes: Lanes) {
-  pushToStack(subtreeRenderLanesCursor, subtreeRenderLanes, fiber);
-  subtreeRenderLanes = mergeLanes(subtreeRenderLanes, lanes);
-  workInProgressRootIncludedLanes = mergeLanes(
-    workInProgressRootIncludedLanes,
-    lanes,
-  );
+// This is called by the HiddenContext module when we enter or leave a
+// hidden subtree. The stack logic is managed there because that's the only
+// place that ever modifies it. Which module it lives in doesn't matter for
+// performance because this function will get inlined regardless
+export function setRenderLanes(subtreeRenderLanes: Lanes) {
+  renderLanes = subtreeRenderLanes;
 }
 
-export function popRenderLanes(fiber: Fiber) {
-  subtreeRenderLanes = subtreeRenderLanesCursor.current;
-  popFromStack(subtreeRenderLanesCursor, fiber);
+export function getRenderLanes(): Lanes {
+  return renderLanes;
 }
 
 function prepareFreshStack(root: FiberRoot, lanes: Lanes): Fiber {
@@ -1497,7 +1483,7 @@ function prepareFreshStack(root: FiberRoot, lanes: Lanes): Fiber {
   workInProgressRoot = root;
   const rootWorkInProgress = createWorkInProgress(root.current, null);
   workInProgress = rootWorkInProgress;
-  workInProgressRootRenderLanes = subtreeRenderLanes = workInProgressRootIncludedLanes = lanes;
+  workInProgressRootRenderLanes = renderLanes = lanes;
   workInProgressRootExitStatus = RootInProgress;
   workInProgressRootFatalError = null;
   workInProgressRootSkippedLanes = NoLanes;
@@ -1864,10 +1850,10 @@ function performUnitOfWork(unitOfWork: Fiber): void {
   let next;
   if (enableProfilerTimer && (unitOfWork.mode & ProfileMode) !== NoMode) {
     startProfilerTimer(unitOfWork);
-    next = beginWork(current, unitOfWork, subtreeRenderLanes);
+    next = beginWork(current, unitOfWork, renderLanes);
     stopProfilerTimerIfRunningAndRecordDelta(unitOfWork, true);
   } else {
-    next = beginWork(current, unitOfWork, subtreeRenderLanes);
+    next = beginWork(current, unitOfWork, renderLanes);
   }
 
   resetCurrentDebugFiberInDEV();
@@ -1901,10 +1887,10 @@ function completeUnitOfWork(unitOfWork: Fiber): void {
         !enableProfilerTimer ||
         (completedWork.mode & ProfileMode) === NoMode
       ) {
-        next = completeWork(current, completedWork, subtreeRenderLanes);
+        next = completeWork(current, completedWork, renderLanes);
       } else {
         startProfilerTimer(completedWork);
-        next = completeWork(current, completedWork, subtreeRenderLanes);
+        next = completeWork(current, completedWork, renderLanes);
         // Update render duration assuming we didn't error.
         stopProfilerTimerIfRunningAndRecordDelta(completedWork, false);
       }
@@ -1919,7 +1905,7 @@ function completeUnitOfWork(unitOfWork: Fiber): void {
       // This fiber did not complete because something threw. Pop values off
       // the stack without entering the complete phase. If this is a boundary,
       // capture values if possible.
-      const next = unwindWork(current, completedWork, subtreeRenderLanes);
+      const next = unwindWork(current, completedWork, renderLanes);
 
       // Because this fiber did not complete, don't reset its lanes.