@@ -742,6 +742,7 @@ export function markRootFinished(root: FiberRoot, remainingLanes: Lanes) {
742742 const entanglements = root . entanglements ;
743743 const eventTimes = root . eventTimes ;
744744 const expirationTimes = root . expirationTimes ;
745+ const pooledCache = root . pooledCache ;
745746
746747 // Clear the lanes that no longer have pending work
747748 let lanes = noLongerPendingLanes ;
@@ -753,15 +754,31 @@ export function markRootFinished(root: FiberRoot, remainingLanes: Lanes) {
753754 eventTimes [ index ] = NoTimestamp ;
754755 expirationTimes [ index ] = NoTimestamp ;
755756
756- lanes &= ~ lane ;
757- }
757+ if ( enableCache ) {
758+ // Subsequent loads in this lane should use a fresh cache.
759+ // TODO: If a cache is no longer associated with any lane, we should issue
760+ // an abort signal.
761+ const caches = root . caches ;
762+ if ( caches !== null ) {
763+ if ( remainingLanes === 0 ) {
764+ // Fast path. Clear all caches at once.
765+ root . caches = createLaneMap ( null ) ;
766+ root . pooledCache = null ;
767+ } else {
768+ const cache = caches [ index ] ;
769+ if ( cache !== null ) {
770+ caches [ index ] = null ;
771+ if ( cache === pooledCache ) {
772+ // The pooled cache is now part of the committed tree. We'll now
773+ // clear it so that the next transition gets a fresh cache.
774+ root. pooledCache = null ;
775+ }
776+ }
777+ }
778+ }
779+ }
758780
759- if ( enableCache ) {
760- // Clear the pooled cache so subsequent updates get fresh data.
761- // TODO: This is very naive and only works if the shell of a cache boundary
762- // doesn't suspend. The next, key feature is to preserve caches across
763- // multiple attempts (suspend -> ping) to render a new tree.
764- root. pooledCache = null ;
781+ lanes &= ~ lane ;
765782 }
766783}
767784
@@ -785,12 +802,62 @@ export function requestFreshCache(root: FiberRoot, renderLanes: Lanes): Cache {
785802 return ( null : any ) ;
786803 }
787804
788- // Check if there's a pooled cache. This is really just a batching heuristic
789- // so that two transitions that happen in a similar timeframe can share the
790- // same cache.
791- const pooledCache = root.pooledCache;
792- if (pooledCache !== null) {
793- return pooledCache ;
805+ // 1. Check if the currently rendering lanes already have a pending cache
806+ // associated with them. If so, use this cache. If for some reason two or
807+ // more lanes have different caches, pick the highest priority one.
808+ // 2. Otherwise, check the root's `pooledCache`. This the oldest cache
809+ // that has not yet been committed. This is really just a batching
810+ // heuristic so that two transitions that happen in a similar timeframe can
811+ // share the same cache. If it exists, use this cache.
812+ // 3. If there's no pooled cache, create a fresh cache. This is now the
813+ // pooled cache.
814+
815+ let caches = root . caches ;
816+
817+ // TODO: There should be a primary render lane, and we should use whatever
818+ // cache is associated with that one.
819+ if ( caches === null ) {
820+ caches = root . caches = createLaneMap ( null ) ;
821+ } else {
822+ let lanes = renderLanes ;
823+ while ( lanes > 0 ) {
824+ const lane = getHighestPriorityLanes ( lanes ) ;
825+ const index = laneToIndex ( lane ) ;
826+ const inProgressCache : Cache | null = caches [ index ] ;
827+ if ( inProgressCache !== null ) {
828+ // This render lane already has a cache associated with it. Reuse it.
829+
830+ // If the other render lanes are not already associated with a cache,
831+ // associate them with this one.
832+ let otherRenderLanes = renderLanes & ~ lane ;
833+ while ( otherRenderLanes > 0 ) {
834+ const otherIndex = pickArbitraryLaneIndex ( otherRenderLanes ) ;
835+ const otherLane = 1 << otherIndex ;
836+ // We shouldn't overwrite a cache that already exists, since that could
837+ // lead to dropped requests or data, i.e. if the current render suspends.
838+ if ( caches [ otherIndex ] === null ) {
839+ caches [ otherIndex ] = inProgressCache ;
840+ }
841+ otherRenderLanes &= ~ otherLane ;
842+ }
843+ return inProgressCache ;
844+ }
845+ lanes &= ~ lane ;
846+ }
847+ // There are no in-progress caches associated with the current render. Check
848+ // if there's a pooled cache.
849+ const pooledCache = root.pooledCache;
850+ if (pooledCache !== null) {
851+ // Associate the pooled cache with each of the render lanes.
852+ lanes = renderLanes ;
853+ while ( lanes > 0 ) {
854+ const index = pickArbitraryLaneIndex ( lanes ) ;
855+ const lane = 1 << index ;
856+ caches [ index ] = pooledCache ;
857+ lanes &= ~ lane ;
858+ }
859+ return pooledCache;
860+ }
794861 }
795862
796863 // Create a fresh cache.
@@ -801,8 +868,61 @@ export function requestFreshCache(root: FiberRoot, renderLanes: Lanes): Cache {
801868
802869 // This is now the pooled cache.
803870 root.pooledCache = freshCache;
871+
872+ // Associate the new cache with each of the render lanes.
873+ let lanes = renderLanes;
874+ while (lanes > 0 ) {
875+ const index = pickArbitraryLaneIndex ( lanes ) ;
876+ const lane = 1 << index ;
877+ caches [ index ] = freshCache ;
878+ lanes &= ~ lane ;
879+ }
880+
804881 return freshCache;
805882}
883+
884+ export function getWorkInProgressCache (
885+ root : FiberRoot ,
886+ renderLanes : Lanes ,
887+ ) : Cache | null {
888+ if ( enableCache ) {
889+ // TODO: There should be a primary render lane, and we should use whatever
890+ // cache is associated with that one.
891+ const caches = root . caches ;
892+ if ( caches !== null ) {
893+ let lanes = renderLanes ;
894+ while ( lanes > 0 ) {
895+ const lane = getHighestPriorityLanes ( lanes ) ;
896+ const index = laneToIndex ( lane ) ;
897+ const inProgressCache : Cache | null = caches [ index ] ;
898+ if ( inProgressCache !== null ) {
899+ return inProgressCache ;
900+ }
901+ lanes &= ~ lane ;
902+ }
903+ }
904+ }
905+ return null;
906+ }
907+
908+ export function transferCacheToSpawnedLane (
909+ root : FiberRoot ,
910+ cache : Cache ,
911+ lane : Lane ,
912+ ) {
913+ const index = laneToIndex ( lane ) ;
914+ let caches = root . caches ;
915+ if ( caches !== null ) {
916+ const existingCache : Cache | null = caches [ index ] ;
917+ if ( existingCache === null ) {
918+ caches [ index ] = cache ;
919+ }
920+ } else {
921+ caches = root . caches = createLaneMap ( null ) ;
922+ caches [ index ] = cache ;
923+ }
924+ }
925+
806926export function getBumpedLaneForHydration (
807927 root : FiberRoot ,
808928 renderLanes : Lanes ,
0 commit comments