Add lazy initialize of router cache nodes (#42629)

packages/next/client/components/app-router.tsx

@@ -7,6 +7,7 @@ import {
   AppRouterContext,
   LayoutRouterContext,
   GlobalLayoutRouterContext,
+  CacheStates,
 } from '../../shared/lib/app-router-context'
 import type {
   CacheNode,
@@ -121,11 +122,12 @@ function Router({
     return {
       tree: initialTree,
       cache: {
+        status: CacheStates.READY,
         data: null,
         subTreeData: children,
         parallelRoutes:
           typeof window === 'undefined' ? new Map() : initialParallelRoutes,
-      },
+      } as CacheNode,
       prefetchCache: new Map(),
       pushRef: { pendingPush: false, mpaNavigation: false },
       focusAndScrollRef: { apply: false },
@@ -176,6 +178,7 @@ function Router({
         previousTree,
         overrideCanonicalUrl,
         cache: {
+          status: CacheStates.LAZYINITIALIZED,
           data: null,
           subTreeData: null,
           parallelRoutes: new Map(),
@@ -201,6 +204,7 @@ function Router({
         forceOptimisticNavigation,
         navigateType,
         cache: {
+          status: CacheStates.LAZYINITIALIZED,
           data: null,
           subTreeData: null,
           parallelRoutes: new Map(),
@@ -262,6 +266,7 @@ function Router({
 
             // TODO-APP: revisit if this needs to be passed.
             cache: {
+              status: CacheStates.LAZYINITIALIZED,
               data: null,
               subTreeData: null,
               parallelRoutes: new Map(),

packages/next/client/components/layout-router.tsx

@@ -16,6 +16,7 @@ import type {
 } from '../../server/app-render'
 import type { ErrorComponent } from './error-boundary'
 import {
+  CacheStates,
   LayoutRouterContext,
   GlobalLayoutRouterContext,
   TemplateContext,
@@ -143,21 +144,29 @@ export function InnerLayoutRouter({
   if (
     childProp &&
     // TODO-APP: verify if this can be null based on user code
-    childProp.current !== null &&
-    !childNode /*&&
-    !childProp.partial*/
+    childProp.current !== null
   ) {
-    // Add the segment's subTreeData to the cache.
-    // This writes to the cache when there is no item in the cache yet. It never *overwrites* existing cache items which is why it's safe in concurrent mode.
-    childNodes.set(path, {
-      data: null,
-      subTreeData: childProp.current,
-      parallelRoutes: new Map(),
-    })
-    // Mutates the prop in order to clean up the memory associated with the subTreeData as it is now part of the cache.
-    childProp.current = null
-    // In the above case childNode was set on childNodes, so we have to get it from the cacheNodes again.
-    childNode = childNodes.get(path)
+    if (childNode && childNode.status === CacheStates.LAZYINITIALIZED) {
+      // @ts-expect-error TODO-APP: handle changing of the type
+      childNode.status = CacheStates.READY
+      // @ts-expect-error TODO-APP: handle changing of the type
+      childNode.subTreeData = childProp.current
+      // Mutates the prop in order to clean up the memory associated with the subTreeData as it is now part of the cache.
+      childProp.current = null
+    } else {
+      // Add the segment's subTreeData to the cache.
+      // This writes to the cache when there is no item in the cache yet. It never *overwrites* existing cache items which is why it's safe in concurrent mode.
+      childNodes.set(path, {
+        status: CacheStates.READY,
+        data: null,
+        subTreeData: childProp.current,
+        parallelRoutes: new Map(),
+      })
+      // Mutates the prop in order to clean up the memory associated with the subTreeData as it is now part of the cache.
+      childProp.current = null
+      // In the above case childNode was set on childNodes, so we have to get it from the cacheNodes again.
+      childNode = childNodes.get(path)
+    }
   }
 
   // When childNode is not available during rendering client-side we need to fetch it from the server.
@@ -172,6 +181,7 @@ export function InnerLayoutRouter({
      * Flight data fetch kicked off during render and put into the cache.
      */
     childNodes.set(path, {
+      status: CacheStates.DATAFETCH,
       data: fetchServerResponse(new URL(url, location.origin), refetchTree),
       subTreeData: null,
       parallelRoutes: new Map(),

packages/next/client/components/reducer.ts

@@ -1,4 +1,4 @@
-import type { CacheNode } from '../../shared/lib/app-router-context'
+import { CacheNode, CacheStates } from '../../shared/lib/app-router-context'
 import type {
   FlightRouterState,
   FlightData,
@@ -55,7 +55,7 @@ function invalidateCacheByRouterState(
   newCache: CacheNode,
   existingCache: CacheNode,
   routerState: FlightRouterState
-) {
+): void {
   // Remove segment that we got data for so that it is filled in during rendering of subTreeData.
   for (const key in routerState[1]) {
     const segmentForParallelRoute = routerState[1][key][0]
@@ -72,6 +72,65 @@ function invalidateCacheByRouterState(
   }
 }
 
+function fillLazyItemsTillLeafWithHead(
+  newCache: CacheNode,
+  existingCache: CacheNode | undefined,
+  routerState: FlightRouterState,
+  head: React.ReactNode
+): void {
+  const isLastSegment = Object.keys(routerState[1]).length === 0
+  if (isLastSegment) {
+    newCache.head = head
+    return
+  }
+  // Remove segment that we got data for so that it is filled in during rendering of subTreeData.
+  for (const key in routerState[1]) {
+    const parallelRouteState = routerState[1][key]
+    const segmentForParallelRoute = parallelRouteState[0]
+    const cacheKey = Array.isArray(segmentForParallelRoute)
+      ? segmentForParallelRoute[1]
+      : segmentForParallelRoute
+    if (existingCache) {
+      const existingParallelRoutesCacheNode =
+        existingCache.parallelRoutes.get(key)
+      if (existingParallelRoutesCacheNode) {
+        let parallelRouteCacheNode = new Map(existingParallelRoutesCacheNode)
+        parallelRouteCacheNode.delete(cacheKey)
+        const newCacheNode: CacheNode = {
+          status: CacheStates.LAZYINITIALIZED,
+          data: null,
+          subTreeData: null,
+          parallelRoutes: new Map(),
+        }
+        parallelRouteCacheNode.set(cacheKey, newCacheNode)
+        fillLazyItemsTillLeafWithHead(
+          newCacheNode,
+          undefined,
+          parallelRouteState,
+          head
+        )
+
+        newCache.parallelRoutes.set(key, parallelRouteCacheNode)
+        continue
+      }
+    }
+
+    const newCacheNode: CacheNode = {
+      status: CacheStates.LAZYINITIALIZED,
+      data: null,
+      subTreeData: null,
+      parallelRoutes: new Map(),
+    }
+    newCache.parallelRoutes.set(key, new Map([[cacheKey, newCacheNode]]))
+    fillLazyItemsTillLeafWithHead(
+      newCacheNode,
+      undefined,
+      parallelRouteState,
+      head
+    )
+  }
+}
+
 /**
  * Fill cache with subTreeData based on flightDataPath
  */
@@ -111,6 +170,7 @@ function fillCacheWithNewSubTreeData(
       childCacheNode === existingChildCacheNode
     ) {
       childCacheNode = {
+        status: CacheStates.READY,
         data: null,
         subTreeData: flightDataPath[3],
         // Ensure segments other than the one we got data for are preserved.
@@ -127,6 +187,13 @@ function fillCacheWithNewSubTreeData(
         )
       }
 
+      fillLazyItemsTillLeafWithHead(
+        childCacheNode,
+        existingChildCacheNode,
+        flightDataPath[2],
+        /* flightDataPath[4] */ undefined
+      )
+
       childSegmentMap.set(segmentForCache, childCacheNode)
     }
     return
@@ -140,10 +207,11 @@ function fillCacheWithNewSubTreeData(
 
   if (childCacheNode === existingChildCacheNode) {
     childCacheNode = {
+      status: childCacheNode.status,
       data: childCacheNode.data,
       subTreeData: childCacheNode.subTreeData,
       parallelRoutes: new Map(childCacheNode.parallelRoutes),
-    }
+    } as CacheNode
     childSegmentMap.set(segmentForCache, childCacheNode)
   }
 
@@ -199,10 +267,11 @@ function invalidateCacheBelowFlightSegmentPath(
 
   if (childCacheNode === existingChildCacheNode) {
     childCacheNode = {
+      status: childCacheNode.status,
       data: childCacheNode.data,
       subTreeData: childCacheNode.subTreeData,
       parallelRoutes: new Map(childCacheNode.parallelRoutes),
-    }
+    } as CacheNode
     childSegmentMap.set(segmentForCache, childCacheNode)
   }
 
@@ -239,6 +308,7 @@ function fillCacheWithPrefetchedSubTreeData(
   if (isLastEntry) {
     if (!existingChildCacheNode) {
       existingChildSegmentMap.set(segmentForCache, {
+        status: CacheStates.READY,
         data: null,
         subTreeData: flightDataPath[3],
         parallelRoutes: new Map(),
@@ -300,6 +370,7 @@ function fillCacheWithDataProperty(
       childCacheNode === existingChildCacheNode
     ) {
       childSegmentMap.set(segment, {
+        status: CacheStates.DATAFETCH,
         data: fetchResponse(),
         subTreeData: null,
         parallelRoutes: new Map(),
@@ -312,6 +383,7 @@ function fillCacheWithDataProperty(
     // Start fetch in the place where the existing cache doesn't have the data yet.
     if (!childCacheNode) {
       childSegmentMap.set(segment, {
+        status: CacheStates.DATAFETCH,
         data: fetchResponse(),
         subTreeData: null,
         parallelRoutes: new Map(),
@@ -322,10 +394,11 @@ function fillCacheWithDataProperty(
 
   if (childCacheNode === existingChildCacheNode) {
     childCacheNode = {
+      status: childCacheNode.status,
       data: childCacheNode.data,
       subTreeData: childCacheNode.subTreeData,
       parallelRoutes: new Map(childCacheNode.parallelRoutes),
-    }
+    } as CacheNode
     childSegmentMap.set(segment, childCacheNode)
   }
 

packages/next/shared/lib/app-router-context.ts

@@ -6,26 +6,61 @@ import type { FlightRouterState, FlightData } from '../../server/app-render'
 
 export type ChildSegmentMap = Map<string, CacheNode>
 
+// eslint-disable-next-line no-shadow
+export enum CacheStates {
+  LAZYINITIALIZED = 'LAZYINITIALIZED',
+  DATAFETCH = 'DATAFETCH',
+  READY = 'READY',
+}
+
 /**
  * Cache node used in app-router / layout-router.
  */
-export type CacheNode = {
-  /**
-   * In-flight request for this node.
-   */
-  data: ReturnType<
-    typeof import('../../client/components/app-router').fetchServerResponse
-  > | null
-  /**
-   * React Component for this node.
-   */
-  subTreeData: React.ReactNode | null
-  /**
-   * Child parallel routes.
-   */
-  parallelRoutes: Map<string, ChildSegmentMap>
-}
-
+export type CacheNode =
+  | {
+      status: CacheStates.DATAFETCH
+      /**
+       * In-flight request for this node.
+       */
+      data: ReturnType<
+        typeof import('../../client/components/app-router').fetchServerResponse
+      > | null
+      head?: React.ReactNode
+      /**
+       * React Component for this node.
+       */
+      subTreeData: null
+      /**
+       * Child parallel routes.
+       */
+      parallelRoutes: Map<string, ChildSegmentMap>
+    }
+  | {
+      status: CacheStates.READY
+      /**
+       * In-flight request for this node.
+       */
+      data: null
+      head?: React.ReactNode
+      /**
+       * React Component for this node.
+       */
+      subTreeData: React.ReactNode
+      /**
+       * Child parallel routes.
+       */
+      parallelRoutes: Map<string, ChildSegmentMap>
+    }
+  | {
+      status: CacheStates.LAZYINITIALIZED
+      data: null
+      head?: React.ReactNode
+      subTreeData: null
+      /**
+       * Child parallel routes.
+       */
+      parallelRoutes: Map<string, ChildSegmentMap>
+    }
 interface NavigateOptions {
   forceOptimisticNavigation?: boolean
 }