Fix prefetch for new router

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

@@ -15,7 +15,7 @@ import type {
 import type { FlightRouterState, FlightData } from '../../server/app-render'
 import {
   ACTION_NAVIGATE,
-  // ACTION_PREFETCH,
+  ACTION_PREFETCH,
   ACTION_RELOAD,
   ACTION_RESTORE,
   ACTION_SERVER_PATCH,
@@ -97,7 +97,7 @@ function ErrorOverlay({ children }: PropsWithChildren<{}>): ReactElement {
 let initialParallelRoutes: CacheNode['parallelRoutes'] =
   typeof window === 'undefined' ? null! : new Map()
 
-// const prefetched = new Set<string>()
+const prefetched = new Set<string>()
 
 /**
  * The global router that wraps the application components.
@@ -208,32 +208,34 @@ export default function AppRouter({
 
     const routerInstance: AppRouterInstance = {
       // TODO-APP: implement prefetching of flight
-      prefetch: async (_href) => {
+      prefetch: async (href) => {
         // If prefetch has already been triggered, don't trigger it again.
-        // if (prefetched.has(href)) {
-        //   return
-        // }
-        // prefetched.add(href)
-        // const url = new URL(href, location.origin)
-        // try {
-        //   // TODO-APP: handle case where history.state is not the new router history entry
-        //   const serverResponse = await fetchServerResponse(
-        //     url,
-        //     // initialTree is used when history.state.tree is missing because the history state is set in `useEffect` below, it being missing means this is the hydration case.
-        //     window.history.state?.tree || initialTree,
-        //     true
-        //   )
-        //   // @ts-ignore startTransition exists
-        //   React.startTransition(() => {
-        //     dispatch({
-        //       type: ACTION_PREFETCH,
-        //       url,
-        //       serverResponse,
-        //     })
-        //   })
-        // } catch (err) {
-        //   console.error('PREFETCH ERROR', err)
-        // }
+        if (prefetched.has(href)) {
+          return
+        }
+        prefetched.add(href)
+        const url = new URL(href, location.origin)
+        try {
+          const routerTree = window.history.state?.tree || initialTree
+          // TODO-APP: handle case where history.state is not the new router history entry
+          const serverResponse = await fetchServerResponse(
+            url,
+            // initialTree is used when history.state.tree is missing because the history state is set in `useEffect` below, it being missing means this is the hydration case.
+            routerTree,
+            true
+          )
+          // @ts-ignore startTransition exists
+          React.startTransition(() => {
+            dispatch({
+              type: ACTION_PREFETCH,
+              url,
+              tree: routerTree,
+              serverResponse,
+            })
+          })
+        } catch (err) {
+          console.error('PREFETCH ERROR', err)
+        }
       },
       replace: (href, options = {}) => {
         // @ts-ignore startTransition exists
@@ -266,7 +268,7 @@ export default function AppRouter({
     }
 
     return routerInstance
-  }, [dispatch /*, initialTree*/])
+  }, [dispatch, initialTree])
 
   useEffect(() => {
     // When mpaNavigation flag is set do a hard navigation to the new url.

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

@@ -29,27 +29,7 @@ import {
 import { fetchServerResponse } from './app-router.client'
 import { createInfinitePromise } from './infinite-promise'
 
-// import { matchSegment } from './match-segments'
-
-/**
- * Check if every segment in array a and b matches
- */
-// function equalSegmentPaths(a: Segment[], b: Segment[]) {
-//   // Comparing length is a fast path.
-//   return a.length === b.length && a.every((val, i) => matchSegment(val, b[i]))
-// }
-
-/**
- * Check if flightDataPath matches layoutSegmentPath
- */
-// function segmentPathMatches(
-//   flightDataPath: FlightDataPath,
-//   layoutSegmentPath: FlightSegmentPath
-// ): boolean {
-//   // The last three items are the current segment, tree, and subTreeData
-//   const pathToLayout = flightDataPath.slice(0, -3)
-//   return equalSegmentPaths(layoutSegmentPath, pathToLayout)
-// }
+import { matchSegment } from './match-segments'
 
 /**
  * Add refetch marker to router state at the point of the current layout segment.
@@ -63,7 +43,7 @@ function walkAddRefetch(
     const [segment, parallelRouteKey] = segmentPathToWalk
     const isLast = segmentPathToWalk.length === 2
 
-    if (treeToRecreate[0] === segment) {
+    if (matchSegment(treeToRecreate[0], segment)) {
       if (treeToRecreate[1].hasOwnProperty(parallelRouteKey)) {
         if (isLast) {
           const subTree = walkAddRefetch(

packages/next/client/components/reducer.ts

@@ -590,6 +590,7 @@ interface ServerPatchAction {
 interface PrefetchAction {
   type: typeof ACTION_PREFETCH
   url: URL
+  tree: FlightRouterState
   serverResponse: Awaited<ReturnType<typeof fetchServerResponse>>
 }
 
@@ -627,7 +628,7 @@ type AppRouterState = {
     string,
     {
       flightSegmentPath: FlightSegmentPath
-      treePatch: FlightRouterState
+      tree: FlightRouterState
       canonicalUrlOverride: URL | undefined
     }
   >
@@ -691,16 +692,11 @@ function clientReducer(
       const prefetchValues = state.prefetchCache.get(href)
       if (prefetchValues) {
         // The one before last item is the router state tree patch
-        const { flightSegmentPath, treePatch, canonicalUrlOverride } =
-          prefetchValues
-
-        // Create new tree based on the flightSegmentPath and router state patch
-        const newTree = applyRouterStatePatchToTree(
-          // TODO-APP: remove ''
-          ['', ...flightSegmentPath],
-          state.tree,
-          treePatch
-        )
+        const {
+          flightSegmentPath,
+          tree: newTree,
+          canonicalUrlOverride,
+        } = prefetchValues
 
         if (newTree !== null) {
           mutable.previousTree = state.tree
@@ -1130,11 +1126,26 @@ function clientReducer(
         fillCacheWithPrefetchedSubTreeData(state.cache, flightDataPath)
       }
 
+      const flightSegmentPath = flightDataPath.slice(0, -2)
+
+      const newTree = applyRouterStatePatchToTree(
+        // TODO-APP: remove ''
+        ['', ...flightSegmentPath],
+        state.tree,
+        treePatch
+      )
+
+      // Patch did not apply correctly
+      if (newTree === null) {
+        return state
+      }
+
       // Create new tree based on the flightSegmentPath and router state patch
       state.prefetchCache.set(href, {
         // Path without the last segment, router state, and the subTreeData
-        flightSegmentPath: flightDataPath.slice(0, -2),
-        treePatch,
+        flightSegmentPath,
+        // Create new tree based on the flightSegmentPath and router state patch
+        tree: newTree,
         canonicalUrlOverride,
       })
 

packages/next/server/app-render.tsx

@@ -1125,12 +1125,21 @@ export async function renderToHTMLOrFlight(
        * Use router state to decide at what common layout to render the page.
        * This can either be the common layout between two pages or a specific place to start rendering from using the "refetch" marker in the tree.
        */
-      const walkTreeWithFlightRouterState = async (
-        loaderTreeToFilter: LoaderTree,
-        parentParams: { [key: string]: string | string[] },
-        flightRouterState?: FlightRouterState,
+      const walkTreeWithFlightRouterState = async ({
+        createSegmentPath,
+        loaderTreeToFilter,
+        parentParams,
+        isFirst,
+        flightRouterState,
+        parentRendered,
+      }: {
+        createSegmentPath: CreateSegmentPath
+        loaderTreeToFilter: LoaderTree
+        parentParams: { [key: string]: string | string[] }
+        isFirst: boolean
+        flightRouterState?: FlightRouterState
         parentRendered?: boolean
-      ): Promise<FlightDataPath> => {
+      }): Promise<FlightDataPath> => {
         const [segment, parallelRoutes] = loaderTreeToFilter
         const parallelRoutesKeys = Object.keys(parallelRoutes)
 
@@ -1175,10 +1184,12 @@ export async function renderToHTMLOrFlight(
                     await createComponentTree(
                       // This ensures flightRouterPath is valid and filters down the tree
                       {
-                        createSegmentPath: (child) => child,
+                        createSegmentPath: (child) => {
+                          return createSegmentPath(child)
+                        },
                         loaderTree: loaderTreeToFilter,
                         parentParams: currentParams,
-                        firstItem: true,
+                        firstItem: isFirst,
                       }
                     )
                   ).Component
@@ -1189,12 +1200,22 @@ export async function renderToHTMLOrFlight(
         // Walk through all parallel routes.
         for (const parallelRouteKey of parallelRoutesKeys) {
           const parallelRoute = parallelRoutes[parallelRouteKey]
-          const path = await walkTreeWithFlightRouterState(
-            parallelRoute,
-            currentParams,
-            flightRouterState && flightRouterState[1][parallelRouteKey],
-            parentRendered || renderComponentsOnThisLevel
-          )
+
+          const currentSegmentPath: FlightSegmentPath = isFirst
+            ? [parallelRouteKey]
+            : [actualSegment, parallelRouteKey]
+
+          const path = await walkTreeWithFlightRouterState({
+            createSegmentPath: (child) => {
+              return createSegmentPath([...currentSegmentPath, ...child])
+            },
+            loaderTreeToFilter: parallelRoute,
+            parentParams: currentParams,
+            flightRouterState:
+              flightRouterState && flightRouterState[1][parallelRouteKey],
+            parentRendered: parentRendered || renderComponentsOnThisLevel,
+            isFirst: false,
+          })
 
           if (typeof path[path.length - 1] !== 'string') {
             return [actualSegment, parallelRouteKey, ...path]
@@ -1209,11 +1230,13 @@ export async function renderToHTMLOrFlight(
       const flightData: FlightData = [
         // TODO-APP: change walk to output without ''
         (
-          await walkTreeWithFlightRouterState(
-            loaderTree,
-            {},
-            providedFlightRouterState
-          )
+          await walkTreeWithFlightRouterState({
+            createSegmentPath: (child) => child,
+            loaderTreeToFilter: loaderTree,
+            parentParams: {},
+            flightRouterState: providedFlightRouterState,
+            isFirst: true,
+          })
         ).slice(1),
       ]
 

test/e2e/app-dir/app/app/nested-navigation/CategoryNav.js

@@ -0,0 +1,28 @@
+'client'
+
+import { TabNavItem } from './TabNavItem'
+import { useSelectedLayoutSegment } from 'next/dist/client/components/hooks-client'
+
+const CategoryNav = ({ categories }) => {
+  const selectedLayoutSegment = useSelectedLayoutSegment()
+
+  return (
+    <div style={{ display: 'flex' }}>
+      <TabNavItem href="/nested-navigation" isActive={!selectedLayoutSegment}>
+        Home
+      </TabNavItem>
+
+      {categories.map((item) => (
+        <TabNavItem
+          key={item.slug}
+          href={`/nested-navigation/${item.slug}`}
+          isActive={item.slug === selectedLayoutSegment}
+        >
+          {item.name}
+        </TabNavItem>
+      ))}
+    </div>
+  )
+}
+
+export default CategoryNav

test/e2e/app-dir/app/app/nested-navigation/TabNavItem.js

@@ -0,0 +1,9 @@
+import Link from 'next/link'
+
+export const TabNavItem = ({ children, href }) => {
+  return (
+    <Link href={href}>
+      <a style={{ margin: '10px', display: 'block' }}>{children}</a>
+    </Link>
+  )
+}

test/e2e/app-dir/app/app/nested-navigation/[categorySlug]/SubCategoryNav.js

@@ -0,0 +1,31 @@
+'client'
+
+import { TabNavItem } from '../TabNavItem'
+import { useSelectedLayoutSegment } from 'next/dist/client/components/hooks-client'
+
+const SubCategoryNav = ({ category }) => {
+  const selectedLayoutSegment = useSelectedLayoutSegment()
+
+  return (
+    <div style={{ display: 'flex' }}>
+      <TabNavItem
+        href={`/nested-navigation/${category.slug}`}
+        isActive={!selectedLayoutSegment}
+      >
+        All
+      </TabNavItem>
+
+      {category.items.map((item) => (
+        <TabNavItem
+          key={item.slug}
+          href={`/nested-navigation/${category.slug}/${item.slug}`}
+          isActive={item.slug === selectedLayoutSegment}
+        >
+          {item.name}
+        </TabNavItem>
+      ))}
+    </div>
+  )
+}
+
+export default SubCategoryNav

test/e2e/app-dir/app/app/nested-navigation/[categorySlug]/[subCategorySlug]/page.js

@@ -0,0 +1,11 @@
+import { experimental_use as use } from 'react'
+import { fetchSubCategory } from '../../getCategories'
+
+export default function Page({ params }) {
+  const category = use(
+    fetchSubCategory(params.categorySlug, params.subCategorySlug)
+  )
+  if (!category) return null
+
+  return <h1 id={category.name.toLowerCase()}>{category.name}</h1>
+}

test/e2e/app-dir/app/app/nested-navigation/[categorySlug]/layout.js

@@ -0,0 +1,14 @@
+import { experimental_use as use } from 'react'
+import { fetchCategoryBySlug } from '../getCategories'
+import SubCategoryNav from './SubCategoryNav'
+
+export default function Layout({ children, params }) {
+  const category = use(fetchCategoryBySlug(params.categorySlug))
+  if (!category) return null
+  return (
+    <>
+      <SubCategoryNav category={category} />
+      {children}
+    </>
+  )
+}

test/e2e/app-dir/app/app/nested-navigation/[categorySlug]/page.js

@@ -0,0 +1,9 @@
+import { experimental_use as use } from 'react'
+import { fetchCategoryBySlug } from '../getCategories'
+
+export default function Page({ params }) {
+  const category = use(fetchCategoryBySlug(params.categorySlug))
+  if (!category) return null
+
+  return <h1 id={`all-${category.name.toLowerCase()}`}>All {category.name}</h1>
+}