Handle head.js on client-side navigation

Implements the last piece to make client-side navigation with head.js work, actually providing the head value and writing it into the router cache.

Builds on earlier changes in vercel#42791 and vercel#42629.

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

@@ -109,6 +109,37 @@ type AppRouterProps = {
   assetPrefix: string
 }
 
+function findHeadInCache(
+  cache: CacheNode,
+  parallelRoutes: FlightRouterState[1]
+): React.ReactNode {
+  const isLastItem = Object.keys(parallelRoutes).length === 0
+  if (isLastItem) {
+    return cache.head
+  }
+  for (const key in parallelRoutes) {
+    const [segment, childParallelRoutes] = parallelRoutes[key]
+    const childSegmentMap = cache.parallelRoutes.get(key)
+    if (!childSegmentMap) {
+      continue
+    }
+
+    const cacheKey = Array.isArray(segment) ? segment[1] : segment
+
+    const cacheNode = childSegmentMap.get(cacheKey)
+    if (!cacheNode) {
+      continue
+    }
+
+    const item = findHeadInCache(cacheNode, childParallelRoutes)
+    if (item) {
+      return item
+    }
+  }
+
+  return undefined
+}
+
 /**
  * The global router that wraps the application components.
  */
@@ -146,6 +177,10 @@ function Router({
     sync,
   ] = useReducerWithReduxDevtools(reducer, initialState)
 
+  const head = useMemo(() => {
+    return findHeadInCache(cache, tree[1])
+  }, [cache, tree])
+
   useEffect(() => {
     // Ensure initialParallelRoutes is cleaned up from memory once it's used.
     initialParallelRoutes = null!
@@ -378,7 +413,7 @@ function Router({
             >
               {HotReloader ? (
                 <HotReloader assetPrefix={assetPrefix}>
-                  {initialHead}
+                  {head || initialHead}
                   {cache.subTreeData}
                 </HotReloader>
               ) : (

packages/next/client/components/reducer.ts

@@ -164,7 +164,6 @@ function fillCacheWithNewSubTreeData(
   const existingChildCacheNode = existingChildSegmentMap.get(segmentForCache)
   let childCacheNode = childSegmentMap.get(segmentForCache)
 
-  // In case of last segment start the fetch at this level and don't copy further down.
   if (isLastEntry) {
     if (
       !childCacheNode ||
@@ -193,7 +192,7 @@ function fillCacheWithNewSubTreeData(
         childCacheNode,
         existingChildCacheNode,
         flightDataPath[2],
-        /* flightDataPath[4] */ undefined
+        flightDataPath[4]
       )
 
       childSegmentMap.set(segmentForCache, childCacheNode)
@@ -309,12 +308,20 @@ function fillCacheWithPrefetchedSubTreeData(
 
   if (isLastEntry) {
     if (!existingChildCacheNode) {
-      existingChildSegmentMap.set(segmentForCache, {
+      const childCacheNode: CacheNode = {
         status: CacheStates.READY,
         data: null,
         subTreeData: flightDataPath[3],
         parallelRoutes: new Map(),
-      })
+      }
+
+      fillLazyItemsTillLeafWithHead(
+        childCacheNode,
+        existingChildCacheNode,
+        flightDataPath[2],
+        flightDataPath[4]
+      )
+      existingChildSegmentMap.set(segmentForCache, childCacheNode)
     }
 
     return

packages/next/server/app-render.tsx

@@ -498,10 +498,11 @@ export type FlightDataPath =
   // Looks somewhat like this
   | [
       // Holds full path to the segment.
-      ...FlightSegmentPath,
+      ...FlightSegmentPath[],
       /* segment of the rendered slice: */ Segment,
       /* treePatch */ FlightRouterState,
-      /* subTreeData: */ React.ReactNode | null // Can be null during prefetch if there's no loading component
+      /* subTreeData: */ React.ReactNode | null, // Can be null during prefetch if there's no loading component
+      /* head */ React.ReactNode | null
     ]
 
 /**
@@ -1363,13 +1364,15 @@ export async function renderToHTMLOrFlight(
         isFirst,
         flightRouterState,
         parentRendered,
+        rscPayloadHead,
       }: {
         createSegmentPath: CreateSegmentPath
         loaderTreeToFilter: LoaderTree
         parentParams: { [key: string]: string | string[] }
         isFirst: boolean
         flightRouterState?: FlightRouterState
         parentRendered?: boolean
+        rscPayloadHead: React.ReactNode
       }): Promise<FlightDataPath> => {
         const [segment, parallelRoutes] = loaderTreeToFilter
         const parallelRoutesKeys = Object.keys(parallelRoutes)
@@ -1426,7 +1429,7 @@ export async function renderToHTMLOrFlight(
                   ).Component
                 ),
             isPrefetch && !Boolean(loaderTreeToFilter[2].loading) ? null : (
-              <>{null}</> // TODO: change this to head tags.
+              <>{rscPayloadHead}</>
             ),
           ]
         }
@@ -1449,6 +1452,7 @@ export async function renderToHTMLOrFlight(
               flightRouterState && flightRouterState[1][parallelRouteKey],
             parentRendered: parentRendered || renderComponentsOnThisLevel,
             isFirst: false,
+            rscPayloadHead,
           })
 
           if (typeof path[path.length - 1] !== 'string') {
@@ -1459,6 +1463,7 @@ export async function renderToHTMLOrFlight(
         return [actualSegment]
       }
 
+      const rscPayloadHead = await resolveHead(loaderTree, {})
       // Flight data that is going to be passed to the browser.
       // Currently a single item array but in the future multiple patches might be combined in a single request.
       const flightData: FlightData = [
@@ -1470,6 +1475,7 @@ export async function renderToHTMLOrFlight(
             parentParams: {},
             flightRouterState: providedFlightRouterState,
             isFirst: true,
+            rscPayloadHead,
           })
         ).slice(1),
       ]