Refactor Cache Code (#23393)

This PR refactors the cache code by moving it out of ReactFiberCacheComponent to ReactFiberTransitionPool in anticipation of it being reused by multiple stacks (ie. transition tracing)

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

@@ -229,15 +229,7 @@ import {
   getWorkInProgressTransitions,
 } from './ReactFiberWorkLoop.new';
 import {setWorkInProgressVersion} from './ReactMutableSource.new';
-import {
-  requestCacheFromPool,
-  pushCacheProvider,
-  pushRootCachePool,
-  CacheContext,
-  getSuspendedCachePool,
-  pushSpawnedCachePool,
-  getOffscreenDeferredCachePool,
-} from './ReactFiberCacheComponent.new';
+import {pushCacheProvider, CacheContext} from './ReactFiberCacheComponent.new';
 import {createCapturedValue} from './ReactCapturedValue';
 import {createClassErrorUpdate} from './ReactFiberThrow.new';
 import {completeSuspendedOffscreenHostContainer} from './ReactFiberCompleteWork.new';
@@ -248,6 +240,13 @@ import {
   pushTreeId,
   pushMaterializedTreeId,
 } from './ReactFiberTreeContext.new';
+import {
+  requestCacheFromPool,
+  pushRootTransition,
+  getSuspendedCache,
+  pushTransition,
+  getOffscreenDeferredCache,
+} from './ReactFiberTransition.new';
 
 const ReactCurrentOwner = ReactSharedInternals.ReactCurrentOwner;
 
@@ -652,7 +651,7 @@ function updateOffscreenComponent(
         // push the cache pool even though we're going to bail out
         // because otherwise there'd be a context mismatch
         if (current !== null) {
-          pushSpawnedCachePool(workInProgress, null);
+          pushTransition(workInProgress, null);
         }
       }
       pushRenderLanes(workInProgress, renderLanes);
@@ -666,7 +665,7 @@ function updateOffscreenComponent(
         nextBaseLanes = mergeLanes(prevBaseLanes, renderLanes);
         if (enableCache) {
           // Save the cache pool so we can resume later.
-          spawnedCachePool = getOffscreenDeferredCachePool();
+          spawnedCachePool = getOffscreenDeferredCache();
         }
       } else {
         nextBaseLanes = renderLanes;
@@ -686,7 +685,7 @@ function updateOffscreenComponent(
         // push the cache pool even though we're going to bail out
         // because otherwise there'd be a context mismatch
         if (current !== null) {
-          pushSpawnedCachePool(workInProgress, null);
+          pushTransition(workInProgress, null);
         }
       }
 
@@ -724,7 +723,7 @@ function updateOffscreenComponent(
         // using the same cache. Unless the parent changed, since that means
         // there was a refresh.
         const prevCachePool = prevState !== null ? prevState.cachePool : null;
-        pushSpawnedCachePool(workInProgress, prevCachePool);
+        pushTransition(workInProgress, prevCachePool);
       }
 
       pushRenderLanes(workInProgress, subtreeRenderLanes);
@@ -742,7 +741,7 @@ function updateOffscreenComponent(
         // using the same cache. Unless the parent changed, since that means
         // there was a refresh.
         const prevCachePool = prevState.cachePool;
-        pushSpawnedCachePool(workInProgress, prevCachePool);
+        pushTransition(workInProgress, prevCachePool);
       }
 
       // Since we're not hidden anymore, reset the state
@@ -758,7 +757,7 @@ function updateOffscreenComponent(
         // using the same cache. Unless the parent changed, since that means
         // there was a refresh.
         if (current !== null) {
-          pushSpawnedCachePool(workInProgress, null);
+          pushTransition(workInProgress, null);
         }
       }
     }
@@ -1329,7 +1328,7 @@ function updateHostRoot(current, workInProgress, renderLanes) {
 
   if (enableCache) {
     const nextCache: Cache = nextState.cache;
-    pushRootCachePool(root);
+    pushRootTransition(root);
     pushCacheProvider(workInProgress, nextCache);
     if (nextCache !== prevState.cache) {
       // The root cache refreshed.
@@ -1910,7 +1909,7 @@ const SUSPENDED_MARKER: SuspenseState = {
 function mountSuspenseOffscreenState(renderLanes: Lanes): OffscreenState {
   return {
     baseLanes: renderLanes,
-    cachePool: getSuspendedCachePool(),
+    cachePool: getSuspendedCache(),
   };
 }
 
@@ -1939,7 +1938,7 @@ function updateSuspenseOffscreenState(
       }
     } else {
       // If there's no previous cache pool, grab the current one.
-      cachePool = getSuspendedCachePool();
+      cachePool = getSuspendedCache();
     }
   }
   return {
@@ -3504,7 +3503,7 @@ function attemptEarlyBailoutIfNoScheduledUpdate(
       if (enableCache) {
         const cache: Cache = current.memoizedState.cache;
         pushCacheProvider(workInProgress, cache);
-        pushRootCachePool(root);
+        pushRootTransition(root);
       }
       if (enableTransitionTracing) {
         workInProgress.memoizedState.transitions = getWorkInProgressTransitions();

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

@@ -229,15 +229,7 @@ import {
   getWorkInProgressTransitions,
 } from './ReactFiberWorkLoop.old';
 import {setWorkInProgressVersion} from './ReactMutableSource.old';
-import {
-  requestCacheFromPool,
-  pushCacheProvider,
-  pushRootCachePool,
-  CacheContext,
-  getSuspendedCachePool,
-  pushSpawnedCachePool,
-  getOffscreenDeferredCachePool,
-} from './ReactFiberCacheComponent.old';
+import {pushCacheProvider, CacheContext} from './ReactFiberCacheComponent.old';
 import {createCapturedValue} from './ReactCapturedValue';
 import {createClassErrorUpdate} from './ReactFiberThrow.old';
 import {completeSuspendedOffscreenHostContainer} from './ReactFiberCompleteWork.old';
@@ -248,6 +240,13 @@ import {
   pushTreeId,
   pushMaterializedTreeId,
 } from './ReactFiberTreeContext.old';
+import {
+  requestCacheFromPool,
+  pushRootTransition,
+  getSuspendedCache,
+  pushTransition,
+  getOffscreenDeferredCache,
+} from './ReactFiberTransition.old';
 
 const ReactCurrentOwner = ReactSharedInternals.ReactCurrentOwner;
 
@@ -652,7 +651,7 @@ function updateOffscreenComponent(
         // push the cache pool even though we're going to bail out
         // because otherwise there'd be a context mismatch
         if (current !== null) {
-          pushSpawnedCachePool(workInProgress, null);
+          pushTransition(workInProgress, null);
         }
       }
       pushRenderLanes(workInProgress, renderLanes);
@@ -666,7 +665,7 @@ function updateOffscreenComponent(
         nextBaseLanes = mergeLanes(prevBaseLanes, renderLanes);
         if (enableCache) {
           // Save the cache pool so we can resume later.
-          spawnedCachePool = getOffscreenDeferredCachePool();
+          spawnedCachePool = getOffscreenDeferredCache();
         }
       } else {
         nextBaseLanes = renderLanes;
@@ -686,7 +685,7 @@ function updateOffscreenComponent(
         // push the cache pool even though we're going to bail out
         // because otherwise there'd be a context mismatch
         if (current !== null) {
-          pushSpawnedCachePool(workInProgress, null);
+          pushTransition(workInProgress, null);
         }
       }
 
@@ -724,7 +723,7 @@ function updateOffscreenComponent(
         // using the same cache. Unless the parent changed, since that means
         // there was a refresh.
         const prevCachePool = prevState !== null ? prevState.cachePool : null;
-        pushSpawnedCachePool(workInProgress, prevCachePool);
+        pushTransition(workInProgress, prevCachePool);
       }
 
       pushRenderLanes(workInProgress, subtreeRenderLanes);
@@ -742,7 +741,7 @@ function updateOffscreenComponent(
         // using the same cache. Unless the parent changed, since that means
         // there was a refresh.
         const prevCachePool = prevState.cachePool;
-        pushSpawnedCachePool(workInProgress, prevCachePool);
+        pushTransition(workInProgress, prevCachePool);
       }
 
       // Since we're not hidden anymore, reset the state
@@ -758,7 +757,7 @@ function updateOffscreenComponent(
         // using the same cache. Unless the parent changed, since that means
         // there was a refresh.
         if (current !== null) {
-          pushSpawnedCachePool(workInProgress, null);
+          pushTransition(workInProgress, null);
         }
       }
     }
@@ -1329,7 +1328,7 @@ function updateHostRoot(current, workInProgress, renderLanes) {
 
   if (enableCache) {
     const nextCache: Cache = nextState.cache;
-    pushRootCachePool(root);
+    pushRootTransition(root);
     pushCacheProvider(workInProgress, nextCache);
     if (nextCache !== prevState.cache) {
       // The root cache refreshed.
@@ -1910,7 +1909,7 @@ const SUSPENDED_MARKER: SuspenseState = {
 function mountSuspenseOffscreenState(renderLanes: Lanes): OffscreenState {
   return {
     baseLanes: renderLanes,
-    cachePool: getSuspendedCachePool(),
+    cachePool: getSuspendedCache(),
   };
 }
 
@@ -1939,7 +1938,7 @@ function updateSuspenseOffscreenState(
       }
     } else {
       // If there's no previous cache pool, grab the current one.
-      cachePool = getSuspendedCachePool();
+      cachePool = getSuspendedCache();
     }
   }
   return {
@@ -3504,7 +3503,7 @@ function attemptEarlyBailoutIfNoScheduledUpdate(
       if (enableCache) {
         const cache: Cache = current.memoizedState.cache;
         pushCacheProvider(workInProgress, cache);
-        pushRootCachePool(root);
+        pushRootTransition(root);
       }
       if (enableTransitionTracing) {
         workInProgress.memoizedState.transitions = getWorkInProgressTransitions();

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

@@ -8,18 +8,12 @@
  */
 
 import type {ReactContext} from 'shared/ReactTypes';
-import type {FiberRoot} from './ReactInternalTypes';
-import type {Lanes} from './ReactFiberLane.new';
-import type {StackCursor} from './ReactFiberStack.new';
 
 import {enableCache} from 'shared/ReactFeatureFlags';
 import {REACT_CONTEXT_TYPE} from 'shared/ReactSymbols';
 
-import {isPrimaryRenderer} from './ReactFiberHostConfig';
-import {createCursor, push, pop} from './ReactFiberStack.new';
 import {pushProvider, popProvider} from './ReactFiberNewContext.new';
 import * as Scheduler from 'scheduler';
-import {getWorkInProgressRoot} from './ReactFiberWorkLoop.new';
 
 export type Cache = {|
   controller: AbortController,
@@ -62,10 +56,6 @@ if (__DEV__ && enableCache) {
   CacheContext._currentRenderer2 = null;
 }
 
-// When retrying a Suspense/Offscreen boundary, we restore the cache that was
-// used during the previous render by placing it here, on the stack.
-const resumedCache: StackCursor<Cache | null> = createCursor(null);
-
 // Creates a new empty Cache instance with a ref-count of 0. The caller is responsible
 // for retaining the cache once it is in use (retainCache), and releasing the cache
 // once it is no longer needed (releaseCache).
@@ -131,135 +121,3 @@ export function popCacheProvider(workInProgress: Fiber, cache: Cache) {
   }
   popProvider(CacheContext, workInProgress);
 }
-
-function peekCacheFromPool(): Cache | null {
-  if (!enableCache) {
-    return (null: any);
-  }
-
-  // Check if the cache pool already has a cache we can use.
-
-  // If we're rendering inside a Suspense boundary that is currently hidden,
-  // we should use the same cache that we used during the previous render, if
-  // one exists.
-  const cacheResumedFromPreviousRender = resumedCache.current;
-  if (cacheResumedFromPreviousRender !== null) {
-    return cacheResumedFromPreviousRender;
-  }
-
-  // Otherwise, check the root's cache pool.
-  const root = (getWorkInProgressRoot(): any);
-  const cacheFromRootCachePool = root.pooledCache;
-
-  return cacheFromRootCachePool;
-}
-
-export function requestCacheFromPool(renderLanes: Lanes): Cache {
-  // Similar to previous function, except if there's not already a cache in the
-  // pool, we allocate a new one.
-  const cacheFromPool = peekCacheFromPool();
-  if (cacheFromPool !== null) {
-    return cacheFromPool;
-  }
-
-  // Create a fresh cache and add it to the root cache pool. A cache can have
-  // multiple owners:
-  // - A cache pool that lives on the FiberRoot. This is where all fresh caches
-  //   are originally created (TODO: except during refreshes, until we implement
-  //   this correctly). The root takes ownership immediately when the cache is
-  //   created. Conceptually, root.pooledCache is an Option<Arc<Cache>> (owned),
-  //   and the return value of this function is a &Arc<Cache> (borrowed).
-  // - One of several fiber types: host root, cache boundary, suspense
-  //   component. These retain and release in the commit phase.
-
-  const root = (getWorkInProgressRoot(): any);
-  const freshCache = createCache();
-  root.pooledCache = freshCache;
-  retainCache(freshCache);
-  if (freshCache !== null) {
-    root.pooledCacheLanes |= renderLanes;
-  }
-  return freshCache;
-}
-
-export function pushRootCachePool(root: FiberRoot) {
-  if (!enableCache) {
-    return;
-  }
-  // Note: This function currently does nothing but I'll leave it here for
-  // code organization purposes in case that changes.
-}
-
-export function popRootCachePool(root: FiberRoot, renderLanes: Lanes) {
-  if (!enableCache) {
-    return;
-  }
-  // Note: This function currently does nothing but I'll leave it here for
-  // code organization purposes in case that changes.
-}
-
-export function pushSpawnedCachePool(
-  offscreenWorkInProgress: Fiber,
-  prevCachePool: SpawnedCachePool | null,
-): void {
-  if (!enableCache) {
-    return;
-  }
-
-  if (prevCachePool === null) {
-    push(resumedCache, resumedCache.current, offscreenWorkInProgress);
-  } else {
-    push(resumedCache, prevCachePool.pool, offscreenWorkInProgress);
-  }
-}
-
-export function popCachePool(workInProgress: Fiber) {
-  if (!enableCache) {
-    return;
-  }
-
-  pop(resumedCache, workInProgress);
-}
-
-export function getSuspendedCachePool(): SpawnedCachePool | null {
-  if (!enableCache) {
-    return null;
-  }
-  // This function is called when a Suspense boundary suspends. It returns the
-  // cache that would have been used to render fresh data during this render,
-  // if there was any, so that we can resume rendering with the same cache when
-  // we receive more data.
-  const cacheFromPool = peekCacheFromPool();
-  if (cacheFromPool === null) {
-    return null;
-  }
-
-  return {
-    // We must also save the parent, so that when we resume we can detect
-    // a refresh.
-    parent: isPrimaryRenderer
-      ? CacheContext._currentValue
-      : CacheContext._currentValue2,
-    pool: cacheFromPool,
-  };
-}
-
-export function getOffscreenDeferredCachePool(): SpawnedCachePool | null {
-  if (!enableCache) {
-    return null;
-  }
-
-  const cacheFromPool = peekCacheFromPool();
-  if (cacheFromPool === null) {
-    return null;
-  }
-
-  return {
-    // We must also store the parent, so that when we resume we can detect
-    // a refresh.
-    parent: isPrimaryRenderer
-      ? CacheContext._currentValue
-      : CacheContext._currentValue2,
-    pool: cacheFromPool,
-  };
-}

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

@@ -8,18 +8,12 @@
  */
 
 import type {ReactContext} from 'shared/ReactTypes';
-import type {FiberRoot} from './ReactInternalTypes';
-import type {Lanes} from './ReactFiberLane.old';
-import type {StackCursor} from './ReactFiberStack.old';
 
 import {enableCache} from 'shared/ReactFeatureFlags';
 import {REACT_CONTEXT_TYPE} from 'shared/ReactSymbols';
 
-import {isPrimaryRenderer} from './ReactFiberHostConfig';
-import {createCursor, push, pop} from './ReactFiberStack.old';
 import {pushProvider, popProvider} from './ReactFiberNewContext.old';
 import * as Scheduler from 'scheduler';
-import {getWorkInProgressRoot} from './ReactFiberWorkLoop.old';
 
 export type Cache = {|
   controller: AbortController,
@@ -62,10 +56,6 @@ if (__DEV__ && enableCache) {
   CacheContext._currentRenderer2 = null;
 }
 
-// When retrying a Suspense/Offscreen boundary, we restore the cache that was
-// used during the previous render by placing it here, on the stack.
-const resumedCache: StackCursor<Cache | null> = createCursor(null);
-
 // Creates a new empty Cache instance with a ref-count of 0. The caller is responsible
 // for retaining the cache once it is in use (retainCache), and releasing the cache
 // once it is no longer needed (releaseCache).
@@ -131,135 +121,3 @@ export function popCacheProvider(workInProgress: Fiber, cache: Cache) {
   }
   popProvider(CacheContext, workInProgress);
 }
-
-function peekCacheFromPool(): Cache | null {
-  if (!enableCache) {
-    return (null: any);
-  }
-
-  // Check if the cache pool already has a cache we can use.
-
-  // If we're rendering inside a Suspense boundary that is currently hidden,
-  // we should use the same cache that we used during the previous render, if
-  // one exists.
-  const cacheResumedFromPreviousRender = resumedCache.current;
-  if (cacheResumedFromPreviousRender !== null) {
-    return cacheResumedFromPreviousRender;
-  }
-
-  // Otherwise, check the root's cache pool.
-  const root = (getWorkInProgressRoot(): any);
-  const cacheFromRootCachePool = root.pooledCache;
-
-  return cacheFromRootCachePool;
-}
-
-export function requestCacheFromPool(renderLanes: Lanes): Cache {
-  // Similar to previous function, except if there's not already a cache in the
-  // pool, we allocate a new one.
-  const cacheFromPool = peekCacheFromPool();
-  if (cacheFromPool !== null) {
-    return cacheFromPool;
-  }
-
-  // Create a fresh cache and add it to the root cache pool. A cache can have
-  // multiple owners:
-  // - A cache pool that lives on the FiberRoot. This is where all fresh caches
-  //   are originally created (TODO: except during refreshes, until we implement
-  //   this correctly). The root takes ownership immediately when the cache is
-  //   created. Conceptually, root.pooledCache is an Option<Arc<Cache>> (owned),
-  //   and the return value of this function is a &Arc<Cache> (borrowed).
-  // - One of several fiber types: host root, cache boundary, suspense
-  //   component. These retain and release in the commit phase.
-
-  const root = (getWorkInProgressRoot(): any);
-  const freshCache = createCache();
-  root.pooledCache = freshCache;
-  retainCache(freshCache);
-  if (freshCache !== null) {
-    root.pooledCacheLanes |= renderLanes;
-  }
-  return freshCache;
-}
-
-export function pushRootCachePool(root: FiberRoot) {
-  if (!enableCache) {
-    return;
-  }
-  // Note: This function currently does nothing but I'll leave it here for
-  // code organization purposes in case that changes.
-}
-
-export function popRootCachePool(root: FiberRoot, renderLanes: Lanes) {
-  if (!enableCache) {
-    return;
-  }
-  // Note: This function currently does nothing but I'll leave it here for
-  // code organization purposes in case that changes.
-}
-
-export function pushSpawnedCachePool(
-  offscreenWorkInProgress: Fiber,
-  prevCachePool: SpawnedCachePool | null,
-): void {
-  if (!enableCache) {
-    return;
-  }
-
-  if (prevCachePool === null) {
-    push(resumedCache, resumedCache.current, offscreenWorkInProgress);
-  } else {
-    push(resumedCache, prevCachePool.pool, offscreenWorkInProgress);
-  }
-}
-
-export function popCachePool(workInProgress: Fiber) {
-  if (!enableCache) {
-    return;
-  }
-
-  pop(resumedCache, workInProgress);
-}
-
-export function getSuspendedCachePool(): SpawnedCachePool | null {
-  if (!enableCache) {
-    return null;
-  }
-  // This function is called when a Suspense boundary suspends. It returns the
-  // cache that would have been used to render fresh data during this render,
-  // if there was any, so that we can resume rendering with the same cache when
-  // we receive more data.
-  const cacheFromPool = peekCacheFromPool();
-  if (cacheFromPool === null) {
-    return null;
-  }
-
-  return {
-    // We must also save the parent, so that when we resume we can detect
-    // a refresh.
-    parent: isPrimaryRenderer
-      ? CacheContext._currentValue
-      : CacheContext._currentValue2,
-    pool: cacheFromPool,
-  };
-}
-
-export function getOffscreenDeferredCachePool(): SpawnedCachePool | null {
-  if (!enableCache) {
-    return null;
-  }
-
-  const cacheFromPool = peekCacheFromPool();
-  if (cacheFromPool === null) {
-    return null;
-  }
-
-  return {
-    // We must also store the parent, so that when we resume we can detect
-    // a refresh.
-    parent: isPrimaryRenderer
-      ? CacheContext._currentValue
-      : CacheContext._currentValue2,
-    pool: cacheFromPool,
-  };
-}

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

@@ -162,12 +162,9 @@ import {
 import {resetChildFibers} from './ReactChildFiber.new';
 import {createScopeInstance} from './ReactFiberScope.new';
 import {transferActualDuration} from './ReactProfilerTimer.new';
-import {
-  popCacheProvider,
-  popRootCachePool,
-  popCachePool,
-} from './ReactFiberCacheComponent.new';
+import {popCacheProvider} from './ReactFiberCacheComponent.new';
 import {popTreeContext} from './ReactFiberTreeContext.new';
+import {popRootTransition, popTransition} from './ReactFiberTransition.new';
 
 function markUpdate(workInProgress: Fiber) {
   // Tag the fiber with an update effect. This turns a Placement into
@@ -864,7 +861,7 @@ function completeWork(
     case HostRoot: {
       const fiberRoot = (workInProgress.stateNode: FiberRoot);
       if (enableCache) {
-        popRootCachePool(fiberRoot, renderLanes);
+        popRootTransition(fiberRoot, renderLanes);
 
         let previousCache: Cache | null = null;
         if (current !== null) {
@@ -1553,7 +1550,7 @@ function completeWork(
           workInProgress.flags |= Passive;
         }
         if (current !== null) {
-          popCachePool(workInProgress);
+          popTransition(workInProgress);
         }
       }
 

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

@@ -162,12 +162,9 @@ import {
 import {resetChildFibers} from './ReactChildFiber.old';
 import {createScopeInstance} from './ReactFiberScope.old';
 import {transferActualDuration} from './ReactProfilerTimer.old';
-import {
-  popCacheProvider,
-  popRootCachePool,
-  popCachePool,
-} from './ReactFiberCacheComponent.old';
+import {popCacheProvider} from './ReactFiberCacheComponent.old';
 import {popTreeContext} from './ReactFiberTreeContext.old';
+import {popRootTransition, popTransition} from './ReactFiberTransition.old';
 
 function markUpdate(workInProgress: Fiber) {
   // Tag the fiber with an update effect. This turns a Placement into
@@ -864,7 +861,7 @@ function completeWork(
     case HostRoot: {
       const fiberRoot = (workInProgress.stateNode: FiberRoot);
       if (enableCache) {
-        popRootCachePool(fiberRoot, renderLanes);
+        popRootTransition(fiberRoot, renderLanes);
 
         let previousCache: Cache | null = null;
         if (current !== null) {
@@ -1553,7 +1550,7 @@ function completeWork(
           workInProgress.flags |= Passive;
         }
         if (current !== null) {
-          popCachePool(workInProgress);
+          popTransition(workInProgress);
         }
       }
 

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

@@ -0,0 +1,154 @@
+/**
+ * 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 {FiberRoot} from './ReactInternalTypes';
+import type {Lanes} from './ReactFiberLane.new';
+import type {StackCursor} from './ReactFiberStack.new';
+import type {Cache, SpawnedCachePool} from './ReactFiberCacheComponent.new';
+
+import {enableCache} from 'shared/ReactFeatureFlags';
+import {isPrimaryRenderer} from './ReactFiberHostConfig';
+import {createCursor, push, pop} from './ReactFiberStack.new';
+import {getWorkInProgressRoot} from './ReactFiberWorkLoop.new';
+import {
+  createCache,
+  retainCache,
+  CacheContext,
+} from './ReactFiberCacheComponent.new';
+
+// When retrying a Suspense/Offscreen boundary, we restore the cache that was
+// used during the previous render by placing it here, on the stack.
+const resumedCache: StackCursor<Cache | null> = createCursor(null);
+
+function peekCacheFromPool(): Cache | null {
+  if (!enableCache) {
+    return (null: any);
+  }
+
+  // Check if the cache pool already has a cache we can use.
+
+  // If we're rendering inside a Suspense boundary that is currently hidden,
+  // we should use the same cache that we used during the previous render, if
+  // one exists.
+  const cacheResumedFromPreviousRender = resumedCache.current;
+  if (cacheResumedFromPreviousRender !== null) {
+    return cacheResumedFromPreviousRender;
+  }
+
+  // Otherwise, check the root's cache pool.
+  const root = (getWorkInProgressRoot(): any);
+  const cacheFromRootCachePool = root.pooledCache;
+
+  return cacheFromRootCachePool;
+}
+
+export function requestCacheFromPool(renderLanes: Lanes): Cache {
+  // Similar to previous function, except if there's not already a cache in the
+  // pool, we allocate a new one.
+  const cacheFromPool = peekCacheFromPool();
+  if (cacheFromPool !== null) {
+    return cacheFromPool;
+  }
+
+  // Create a fresh cache and add it to the root cache pool. A cache can have
+  // multiple owners:
+  // - A cache pool that lives on the FiberRoot. This is where all fresh caches
+  //   are originally created (TODO: except during refreshes, until we implement
+  //   this correctly). The root takes ownership immediately when the cache is
+  //   created. Conceptually, root.pooledCache is an Option<Arc<Cache>> (owned),
+  //   and the return value of this function is a &Arc<Cache> (borrowed).
+  // - One of several fiber types: host root, cache boundary, suspense
+  //   component. These retain and release in the commit phase.
+
+  const root = (getWorkInProgressRoot(): any);
+  const freshCache = createCache();
+  root.pooledCache = freshCache;
+  retainCache(freshCache);
+  if (freshCache !== null) {
+    root.pooledCacheLanes |= renderLanes;
+  }
+  return freshCache;
+}
+
+export function pushRootTransition(root: FiberRoot) {
+  if (enableCache) {
+    return;
+  }
+  // Note: This function currently does nothing but I'll leave it here for
+  // code organization purposes in case that changes.
+}
+
+export function popRootTransition(root: FiberRoot, renderLanes: Lanes) {
+  if (enableCache) {
+    return;
+  }
+  // Note: This function currently does nothing but I'll leave it here for
+  // code organization purposes in case that changes.
+}
+
+export function pushTransition(
+  offscreenWorkInProgress: Fiber,
+  prevCachePool: SpawnedCachePool | null,
+): void {
+  if (enableCache) {
+    if (prevCachePool === null) {
+      push(resumedCache, resumedCache.current, offscreenWorkInProgress);
+    } else {
+      push(resumedCache, prevCachePool.pool, offscreenWorkInProgress);
+    }
+  }
+}
+
+export function popTransition(workInProgress: Fiber) {
+  if (enableCache) {
+    pop(resumedCache, workInProgress);
+  }
+}
+
+export function getSuspendedCache(): SpawnedCachePool | null {
+  if (!enableCache) {
+    return null;
+  }
+  // This function is called when a Suspense boundary suspends. It returns the
+  // cache that would have been used to render fresh data during this render,
+  // if there was any, so that we can resume rendering with the same cache when
+  // we receive more data.
+  const cacheFromPool = peekCacheFromPool();
+  if (cacheFromPool === null) {
+    return null;
+  }
+
+  return {
+    // We must also save the parent, so that when we resume we can detect
+    // a refresh.
+    parent: isPrimaryRenderer
+      ? CacheContext._currentValue
+      : CacheContext._currentValue2,
+    pool: cacheFromPool,
+  };
+}
+
+export function getOffscreenDeferredCache(): SpawnedCachePool | null {
+  if (!enableCache) {
+    return null;
+  }
+
+  const cacheFromPool = peekCacheFromPool();
+  if (cacheFromPool === null) {
+    return null;
+  }
+
+  return {
+    // We must also store the parent, so that when we resume we can detect
+    // a refresh.
+    parent: isPrimaryRenderer
+      ? CacheContext._currentValue
+      : CacheContext._currentValue2,
+    pool: cacheFromPool,
+  };
+}

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

@@ -0,0 +1,154 @@
+/**
+ * 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 {FiberRoot} from './ReactInternalTypes';
+import type {Lanes} from './ReactFiberLane.old';
+import type {StackCursor} from './ReactFiberStack.old';
+import type {Cache, SpawnedCachePool} from './ReactFiberCacheComponent.old';
+
+import {enableCache} from 'shared/ReactFeatureFlags';
+import {isPrimaryRenderer} from './ReactFiberHostConfig';
+import {createCursor, push, pop} from './ReactFiberStack.old';
+import {getWorkInProgressRoot} from './ReactFiberWorkLoop.old';
+import {
+  createCache,
+  retainCache,
+  CacheContext,
+} from './ReactFiberCacheComponent.old';
+
+// When retrying a Suspense/Offscreen boundary, we restore the cache that was
+// used during the previous render by placing it here, on the stack.
+const resumedCache: StackCursor<Cache | null> = createCursor(null);
+
+function peekCacheFromPool(): Cache | null {
+  if (!enableCache) {
+    return (null: any);
+  }
+
+  // Check if the cache pool already has a cache we can use.
+
+  // If we're rendering inside a Suspense boundary that is currently hidden,
+  // we should use the same cache that we used during the previous render, if
+  // one exists.
+  const cacheResumedFromPreviousRender = resumedCache.current;
+  if (cacheResumedFromPreviousRender !== null) {
+    return cacheResumedFromPreviousRender;
+  }
+
+  // Otherwise, check the root's cache pool.
+  const root = (getWorkInProgressRoot(): any);
+  const cacheFromRootCachePool = root.pooledCache;
+
+  return cacheFromRootCachePool;
+}
+
+export function requestCacheFromPool(renderLanes: Lanes): Cache {
+  // Similar to previous function, except if there's not already a cache in the
+  // pool, we allocate a new one.
+  const cacheFromPool = peekCacheFromPool();
+  if (cacheFromPool !== null) {
+    return cacheFromPool;
+  }
+
+  // Create a fresh cache and add it to the root cache pool. A cache can have
+  // multiple owners:
+  // - A cache pool that lives on the FiberRoot. This is where all fresh caches
+  //   are originally created (TODO: except during refreshes, until we implement
+  //   this correctly). The root takes ownership immediately when the cache is
+  //   created. Conceptually, root.pooledCache is an Option<Arc<Cache>> (owned),
+  //   and the return value of this function is a &Arc<Cache> (borrowed).
+  // - One of several fiber types: host root, cache boundary, suspense
+  //   component. These retain and release in the commit phase.
+
+  const root = (getWorkInProgressRoot(): any);
+  const freshCache = createCache();
+  root.pooledCache = freshCache;
+  retainCache(freshCache);
+  if (freshCache !== null) {
+    root.pooledCacheLanes |= renderLanes;
+  }
+  return freshCache;
+}
+
+export function pushRootTransition(root: FiberRoot) {
+  if (enableCache) {
+    return;
+  }
+  // Note: This function currently does nothing but I'll leave it here for
+  // code organization purposes in case that changes.
+}
+
+export function popRootTransition(root: FiberRoot, renderLanes: Lanes) {
+  if (enableCache) {
+    return;
+  }
+  // Note: This function currently does nothing but I'll leave it here for
+  // code organization purposes in case that changes.
+}
+
+export function pushTransition(
+  offscreenWorkInProgress: Fiber,
+  prevCachePool: SpawnedCachePool | null,
+): void {
+  if (enableCache) {
+    if (prevCachePool === null) {
+      push(resumedCache, resumedCache.current, offscreenWorkInProgress);
+    } else {
+      push(resumedCache, prevCachePool.pool, offscreenWorkInProgress);
+    }
+  }
+}
+
+export function popTransition(workInProgress: Fiber) {
+  if (enableCache) {
+    pop(resumedCache, workInProgress);
+  }
+}
+
+export function getSuspendedCache(): SpawnedCachePool | null {
+  if (!enableCache) {
+    return null;
+  }
+  // This function is called when a Suspense boundary suspends. It returns the
+  // cache that would have been used to render fresh data during this render,
+  // if there was any, so that we can resume rendering with the same cache when
+  // we receive more data.
+  const cacheFromPool = peekCacheFromPool();
+  if (cacheFromPool === null) {
+    return null;
+  }
+
+  return {
+    // We must also save the parent, so that when we resume we can detect
+    // a refresh.
+    parent: isPrimaryRenderer
+      ? CacheContext._currentValue
+      : CacheContext._currentValue2,
+    pool: cacheFromPool,
+  };
+}
+
+export function getOffscreenDeferredCache(): SpawnedCachePool | null {
+  if (!enableCache) {
+    return null;
+  }
+
+  const cacheFromPool = peekCacheFromPool();
+  if (cacheFromPool === null) {
+    return null;
+  }
+
+  return {
+    // We must also store the parent, so that when we resume we can detect
+    // a refresh.
+    parent: isPrimaryRenderer
+      ? CacheContext._currentValue
+      : CacheContext._currentValue2,
+    pool: cacheFromPool,
+  };
+}

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

@@ -0,0 +1,154 @@
+/**
+ * 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 {FiberRoot} from './ReactInternalTypes';
+import type {Lanes} from './ReactFiberLane.old';
+import type {StackCursor} from './ReactFiberStack.old';
+import type {Cache, SpawnedCachePool} from './ReactFiberCacheComponent.old';
+
+import {enableCache} from 'shared/ReactFeatureFlags';
+import {isPrimaryRenderer} from './ReactFiberHostConfig';
+import {createCursor, push, pop} from './ReactFiberStack.old';
+import {getWorkInProgressRoot} from './ReactFiberWorkLoop.old';
+import {
+  createCache,
+  retainCache,
+  CacheContext,
+} from './ReactFiberCacheComponent.old';
+
+// When retrying a Suspense/Offscreen boundary, we restore the cache that was
+// used during the previous render by placing it here, on the stack.
+const resumedCache: StackCursor<Cache | null> = createCursor(null);
+
+function peekCacheFromPool(): Cache | null {
+  if (!enableCache) {
+    return (null: any);
+  }
+
+  // Check if the cache pool already has a cache we can use.
+
+  // If we're rendering inside a Suspense boundary that is currently hidden,
+  // we should use the same cache that we used during the previous render, if
+  // one exists.
+  const cacheResumedFromPreviousRender = resumedCache.current;
+  if (cacheResumedFromPreviousRender !== null) {
+    return cacheResumedFromPreviousRender;
+  }
+
+  // Otherwise, check the root's cache pool.
+  const root = (getWorkInProgressRoot(): any);
+  const cacheFromRootCachePool = root.pooledCache;
+
+  return cacheFromRootCachePool;
+}
+
+export function requestCacheFromPool(renderLanes: Lanes): Cache {
+  // Similar to previous function, except if there's not already a cache in the
+  // pool, we allocate a new one.
+  const cacheFromPool = peekCacheFromPool();
+  if (cacheFromPool !== null) {
+    return cacheFromPool;
+  }
+
+  // Create a fresh cache and add it to the root cache pool. A cache can have
+  // multiple owners:
+  // - A cache pool that lives on the FiberRoot. This is where all fresh caches
+  //   are originally created (TODO: except during refreshes, until we implement
+  //   this correctly). The root takes ownership immediately when the cache is
+  //   created. Conceptually, root.pooledCache is an Option<Arc<Cache>> (owned),
+  //   and the return value of this function is a &Arc<Cache> (borrowed).
+  // - One of several fiber types: host root, cache boundary, suspense
+  //   component. These retain and release in the commit phase.
+
+  const root = (getWorkInProgressRoot(): any);
+  const freshCache = createCache();
+  root.pooledCache = freshCache;
+  retainCache(freshCache);
+  if (freshCache !== null) {
+    root.pooledCacheLanes |= renderLanes;
+  }
+  return freshCache;
+}
+
+export function pushRootTransitionPool(root: FiberRoot) {
+  if (enableCache) {
+    return;
+  }
+  // Note: This function currently does nothing but I'll leave it here for
+  // code organization purposes in case that changes.
+}
+
+export function popRootTransitionPool(root: FiberRoot, renderLanes: Lanes) {
+  if (enableCache) {
+    return;
+  }
+  // Note: This function currently does nothing but I'll leave it here for
+  // code organization purposes in case that changes.
+}
+
+export function pushTransitionPool(
+  offscreenWorkInProgress: Fiber,
+  prevCachePool: SpawnedCachePool | null,
+): void {
+  if (enableCache) {
+    if (prevCachePool === null) {
+      push(resumedCache, resumedCache.current, offscreenWorkInProgress);
+    } else {
+      push(resumedCache, prevCachePool.pool, offscreenWorkInProgress);
+    }
+  }
+}
+
+export function popTransitionPool(workInProgress: Fiber) {
+  if (enableCache) {
+    pop(resumedCache, workInProgress);
+  }
+}
+
+export function getSuspendedCachePool(): SpawnedCachePool | null {
+  if (!enableCache) {
+    return null;
+  }
+  // This function is called when a Suspense boundary suspends. It returns the
+  // cache that would have been used to render fresh data during this render,
+  // if there was any, so that we can resume rendering with the same cache when
+  // we receive more data.
+  const cacheFromPool = peekCacheFromPool();
+  if (cacheFromPool === null) {
+    return null;
+  }
+
+  return {
+    // We must also save the parent, so that when we resume we can detect
+    // a refresh.
+    parent: isPrimaryRenderer
+      ? CacheContext._currentValue
+      : CacheContext._currentValue2,
+    pool: cacheFromPool,
+  };
+}
+
+export function getOffscreenDeferredCachePool(): SpawnedCachePool | null {
+  if (!enableCache) {
+    return null;
+  }
+
+  const cacheFromPool = peekCacheFromPool();
+  if (cacheFromPool === null) {
+    return null;
+  }
+
+  return {
+    // We must also store the parent, so that when we resume we can detect
+    // a refresh.
+    parent: isPrimaryRenderer
+      ? CacheContext._currentValue
+      : CacheContext._currentValue2,
+    pool: cacheFromPool,
+  };
+}

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

@@ -44,13 +44,10 @@ import {
 } from './ReactFiberContext.new';
 import {popProvider} from './ReactFiberNewContext.new';
 import {popRenderLanes} from './ReactFiberWorkLoop.new';
-import {
-  popCacheProvider,
-  popRootCachePool,
-  popCachePool,
-} from './ReactFiberCacheComponent.new';
+import {popCacheProvider} from './ReactFiberCacheComponent.new';
 import {transferActualDuration} from './ReactProfilerTimer.new';
 import {popTreeContext} from './ReactFiberTreeContext.new';
+import {popRootTransition, popTransition} from './ReactFiberTransition.new';
 
 function unwindWork(
   current: Fiber | null,
@@ -84,7 +81,7 @@ function unwindWork(
     case HostRoot: {
       if (enableCache) {
         const root: FiberRoot = workInProgress.stateNode;
-        popRootCachePool(root, renderLanes);
+        popRootTransition(root, renderLanes);
 
         const cache: Cache = workInProgress.memoizedState.cache;
         popCacheProvider(workInProgress, cache);
@@ -158,7 +155,7 @@ function unwindWork(
       popRenderLanes(workInProgress);
       if (enableCache) {
         if (current !== null) {
-          popCachePool(workInProgress);
+          popTransition(workInProgress);
         }
       }
       return null;
@@ -194,7 +191,7 @@ function unwindInterruptedWork(
     case HostRoot: {
       if (enableCache) {
         const root: FiberRoot = interruptedWork.stateNode;
-        popRootCachePool(root, renderLanes);
+        popRootTransition(root, renderLanes);
 
         const cache: Cache = interruptedWork.memoizedState.cache;
         popCacheProvider(interruptedWork, cache);
@@ -226,7 +223,7 @@ function unwindInterruptedWork(
       popRenderLanes(interruptedWork);
       if (enableCache) {
         if (current !== null) {
-          popCachePool(interruptedWork);
+          popTransition(interruptedWork);
         }
       }
 

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

@@ -44,13 +44,10 @@ import {
 } from './ReactFiberContext.old';
 import {popProvider} from './ReactFiberNewContext.old';
 import {popRenderLanes} from './ReactFiberWorkLoop.old';
-import {
-  popCacheProvider,
-  popRootCachePool,
-  popCachePool,
-} from './ReactFiberCacheComponent.old';
+import {popCacheProvider} from './ReactFiberCacheComponent.old';
 import {transferActualDuration} from './ReactProfilerTimer.old';
 import {popTreeContext} from './ReactFiberTreeContext.old';
+import {popRootTransition, popTransition} from './ReactFiberTransition.old';
 
 function unwindWork(
   current: Fiber | null,
@@ -84,7 +81,7 @@ function unwindWork(
     case HostRoot: {
       if (enableCache) {
         const root: FiberRoot = workInProgress.stateNode;
-        popRootCachePool(root, renderLanes);
+        popRootTransition(root, renderLanes);
 
         const cache: Cache = workInProgress.memoizedState.cache;
         popCacheProvider(workInProgress, cache);
@@ -158,7 +155,7 @@ function unwindWork(
       popRenderLanes(workInProgress);
       if (enableCache) {
         if (current !== null) {
-          popCachePool(workInProgress);
+          popTransition(workInProgress);
         }
       }
       return null;
@@ -194,7 +191,7 @@ function unwindInterruptedWork(
     case HostRoot: {
       if (enableCache) {
         const root: FiberRoot = interruptedWork.stateNode;
-        popRootCachePool(root, renderLanes);
+        popRootTransition(root, renderLanes);
 
         const cache: Cache = interruptedWork.memoizedState.cache;
         popCacheProvider(interruptedWork, cache);
@@ -226,7 +223,7 @@ function unwindInterruptedWork(
       popRenderLanes(interruptedWork);
       if (enableCache) {
         if (current !== null) {
-          popCachePool(interruptedWork);
+          popTransition(interruptedWork);
         }
       }