remove dependency on updateQueue for React Cache

Author: lunaruan

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

@@ -662,10 +662,13 @@ function updateOffscreenComponent(
         nextBaseLanes = mergeLanes(prevBaseLanes, renderLanes);
         if (enableCache) {
           // Save the cache pool so we can resume later.
+          const prevCachePool = prevState.cachePool;
+          if (prevCachePool !== null) {
+            // push the cache pool even though we're going to bail out
+            // because otherwise there'd be a context mismatch
+            restoreSpawnedCachePool(workInProgress, prevCachePool);
+          }
           spawnedCachePool = getOffscreenDeferredCachePool();
-          // We don't need to push to the cache pool because we're about to
-          // bail out. There won't be a context mismatch because we only pop
-          // the cache pool if `updateQueue` is non-null.
         }
       } else {
         nextBaseLanes = renderLanes;
@@ -757,12 +760,6 @@ function updateOffscreenComponent(
     pushRenderLanes(workInProgress, subtreeRenderLanes);
   }
 
-  if (enableCache) {
-    // If we have a cache pool from a previous render attempt, then this will be
-    // non-null. We use this to infer whether to push/pop the cache context.
-    workInProgress.updateQueue = spawnedCachePool;
-  }
-
   if (enablePersistentOffscreenHostContainer && supportsPersistence) {
     // In persistent mode, the offscreen children are wrapped in a host node.
     // TODO: Optimize this to use the OffscreenComponent fiber instead of

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

@@ -28,7 +28,7 @@ import type {
 } from './ReactFiberSuspenseComponent.new';
 import type {SuspenseContext} from './ReactFiberSuspenseContext.new';
 import type {OffscreenState} from './ReactFiberOffscreenComponent';
-import type {Cache, SpawnedCachePool} from './ReactFiberCacheComponent.new';
+import type {Cache} from './ReactFiberCacheComponent.new';
 import {
   enableClientRenderFallbackOnHydrationMismatch,
   enableSuspenseAvoidThisFallback,
@@ -1543,8 +1543,14 @@ function completeWork(
           // Run passive effects to retain/release the cache.
           workInProgress.flags |= Passive;
         }
-        const spawnedCachePool: SpawnedCachePool | null = (workInProgress.updateQueue: any);
-        if (spawnedCachePool !== null) {
+        let prevState: OffscreenState | null = null;
+        if (
+          workInProgress.alternate !== null &&
+          workInProgress.alternate.memoizedState !== null
+        ) {
+          prevState = workInProgress.alternate.memoizedState;
+        }
+        if (prevState !== null && prevState.cachePool !== null) {
           popCachePool(workInProgress);
         }
       }

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

@@ -11,7 +11,8 @@ import type {ReactContext} from 'shared/ReactTypes';
 import type {Fiber, FiberRoot} from './ReactInternalTypes';
 import type {Lanes} from './ReactFiberLane.new';
 import type {SuspenseState} from './ReactFiberSuspenseComponent.new';
-import type {Cache, SpawnedCachePool} from './ReactFiberCacheComponent.new';
+import type {Cache} from './ReactFiberCacheComponent.new';
+import type {OffscreenState} from './ReactFiberOffscreenComponent';
 
 import {resetWorkInProgressVersions as resetMutableSourceWorkInProgressVersions} from './ReactMutableSource.new';
 import {
@@ -151,9 +152,15 @@ function unwindWork(workInProgress: Fiber, renderLanes: Lanes) {
     case OffscreenComponent:
     case LegacyHiddenComponent:
       popRenderLanes(workInProgress);
+      let prevState: OffscreenState | null = null;
       if (enableCache) {
-        const spawnedCachePool: SpawnedCachePool | null = (workInProgress.updateQueue: any);
-        if (spawnedCachePool !== null) {
+        if (
+          workInProgress.alternate !== null &&
+          workInProgress.alternate.memoizedState !== null
+        ) {
+          prevState = workInProgress.alternate.memoizedState;
+        }
+        if (prevState !== null && prevState.cachePool !== null) {
           popCachePool(workInProgress);
         }
       }
@@ -217,8 +224,14 @@ function unwindInterruptedWork(interruptedWork: Fiber, renderLanes: Lanes) {
     case LegacyHiddenComponent:
       popRenderLanes(interruptedWork);
       if (enableCache) {
-        const spawnedCachePool: SpawnedCachePool | null = (interruptedWork.updateQueue: any);
-        if (spawnedCachePool !== null) {
+        let prevState: OffscreenState | null = null;
+        if (
+          interruptedWork.alternate !== null &&
+          interruptedWork.alternate.memoizedState !== null
+        ) {
+          prevState = interruptedWork.alternate.memoizedState;
+        }
+        if (prevState !== null && prevState.cachePool !== null) {
           popCachePool(interruptedWork);
         }
       }