Add prefetch to new router

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

@@ -1,4 +1,5 @@
-import React, { useEffect } from 'react'
+import type { PropsWithChildren, ReactElement, ReactNode } from 'react'
+import React, { useEffect, useMemo, useCallback } from 'react'
 import { createFromReadableStream } from 'next/dist/compiled/react-server-dom-webpack'
 import {
   AppRouterContext,
@@ -12,6 +13,7 @@ import type {
 import type { FlightRouterState, FlightData } from '../../server/app-render'
 import {
   ACTION_NAVIGATE,
+  ACTION_PREFETCH,
   ACTION_RELOAD,
   ACTION_RESTORE,
   ACTION_SERVER_PATCH,
@@ -23,13 +25,15 @@ import {
   PathnameContext,
   // LayoutSegmentsContext,
 } from './hooks-client-context'
+import { useReducerWithReduxDevtools } from './use-reducer-with-devtools'
 
 /**
  * Fetch the flight data for the provided url. Takes in the current router state to decide what to render server-side.
  */
 function fetchFlight(
   url: URL,
-  flightRouterState: FlightRouterState
+  flightRouterState: FlightRouterState,
+  prefetch?: true
 ): ReadableStream {
   const flightUrl = new URL(url)
   const searchParams = flightUrl.searchParams
@@ -40,6 +44,9 @@ function fetchFlight(
     '__flight_router_state_tree__',
     JSON.stringify(flightRouterState)
   )
+  if (prefetch) {
+    searchParams.append('__flight_prefetch__', '1')
+  }
 
   // TODO-APP: Verify that TransformStream is supported.
   const { readable, writable } = new TransformStream()
@@ -56,18 +63,17 @@ function fetchFlight(
  */
 export function fetchServerResponse(
   url: URL,
-  flightRouterState: FlightRouterState
+  flightRouterState: FlightRouterState,
+  prefetch?: true
 ): { readRoot: () => FlightData } {
   // Handle the `fetch` readable stream that can be read using `readRoot`.
-  return createFromReadableStream(fetchFlight(url, flightRouterState))
+  return createFromReadableStream(fetchFlight(url, flightRouterState, prefetch))
 }
 
 /**
  * Renders development error overlay when NODE_ENV is development.
  */
-function ErrorOverlay({
-  children,
-}: React.PropsWithChildren<{}>): React.ReactElement {
+function ErrorOverlay({ children }: PropsWithChildren<{}>): ReactElement {
   if (process.env.NODE_ENV === 'production') {
     return <>{children}</>
   } else {
@@ -83,6 +89,8 @@ function ErrorOverlay({
 let initialParallelRoutes: CacheNode['parallelRoutes'] =
   typeof window === 'undefined' ? null! : new Map()
 
+const prefetched = new Set<string>()
+
 /**
  * The global router that wraps the application components.
  */
@@ -94,34 +102,41 @@ export default function AppRouter({
 }: {
   initialTree: FlightRouterState
   initialCanonicalUrl: string
-  children: React.ReactNode
-  hotReloader?: React.ReactNode
+  children: ReactNode
+  hotReloader?: ReactNode
 }) {
-  const [{ tree, cache, pushRef, focusAndScrollRef, canonicalUrl }, dispatch] =
-    React.useReducer(reducer, {
+  const initialState = useMemo(() => {
+    return {
       tree: initialTree,
       cache: {
         data: null,
         subTreeData: children,
         parallelRoutes:
           typeof window === 'undefined' ? new Map() : initialParallelRoutes,
       },
+      prefetchCache: new Map(),
       pushRef: { pendingPush: false, mpaNavigation: false },
       focusAndScrollRef: { apply: false },
       canonicalUrl:
         initialCanonicalUrl +
         // Hash is read as the initial value for canonicalUrl in the browser
         // This is safe to do as canonicalUrl can't be rendered, it's only used to control the history updates the useEffect further down.
         (typeof window !== 'undefined' ? window.location.hash : ''),
-    })
+    }
+  }, [children, initialCanonicalUrl, initialTree])
+  const [
+    { tree, cache, prefetchCache, pushRef, focusAndScrollRef, canonicalUrl },
+    dispatch,
+    sync,
+  ] = useReducerWithReduxDevtools(reducer, initialState)
 
   useEffect(() => {
     // Ensure initialParallelRoutes is cleaned up from memory once it's used.
     initialParallelRoutes = null!
   }, [])
 
   // Add memoized pathname/query for useSearchParams and usePathname.
-  const { searchParams, pathname } = React.useMemo(() => {
+  const { searchParams, pathname } = useMemo(() => {
     const url = new URL(
       canonicalUrl,
       typeof window === 'undefined' ? 'http://n' : window.location.href
@@ -138,7 +153,7 @@ export default function AppRouter({
   /**
    * Server response that only patches the cache and tree.
    */
-  const changeByServerResponse = React.useCallback(
+  const changeByServerResponse = useCallback(
     (previousTree: FlightRouterState, flightData: FlightData) => {
       dispatch({
         type: ACTION_SERVER_PATCH,
@@ -149,24 +164,25 @@ export default function AppRouter({
           subTreeData: null,
           parallelRoutes: new Map(),
         },
+        mutable: {},
       })
     },
-    []
+    [dispatch]
   )
 
   /**
    * The app router that is exposed through `useRouter`. It's only concerned with dispatching actions to the reducer, does not hold state.
    */
-  const appRouter = React.useMemo<AppRouterInstance>(() => {
+  const appRouter = useMemo<AppRouterInstance>(() => {
     const navigate = (
       href: string,
-      cacheType: 'hard' | 'soft',
-      navigateType: 'push' | 'replace'
+      navigateType: 'push' | 'replace',
+      forceOptimisticNavigation: boolean
     ) => {
       return dispatch({
         type: ACTION_NAVIGATE,
         url: new URL(href, location.origin),
-        cacheType,
+        forceOptimisticNavigation,
         navigateType,
         cache: {
           data: null,
@@ -179,29 +195,47 @@ export default function AppRouter({
 
     const routerInstance: AppRouterInstance = {
       // TODO-APP: implement prefetching of flight
-      prefetch: (_href) => Promise.resolve(),
-      replace: (href) => {
-        // @ts-ignore startTransition exists
-        React.startTransition(() => {
-          navigate(href, 'hard', 'replace')
-        })
-      },
-      softReplace: (href) => {
-        // @ts-ignore startTransition exists
-        React.startTransition(() => {
-          navigate(href, 'soft', 'replace')
-        })
+      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)
+        // TODO-APP: handle case where history.state is not the new router history entry
+        const r = 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
+        )
+        try {
+          r.readRoot()
+        } catch (e) {
+          await e
+          const flightData = r.readRoot()
+          // @ts-ignore startTransition exists
+          React.startTransition(() => {
+            dispatch({
+              type: ACTION_PREFETCH,
+              url,
+              flightData,
+            })
+          })
+        }
       },
-      softPush: (href) => {
+      replace: (href, options = {}) => {
         // @ts-ignore startTransition exists
         React.startTransition(() => {
-          navigate(href, 'soft', 'push')
+          navigate(href, 'replace', Boolean(options.forceOptimisticNavigation))
         })
       },
-      push: (href) => {
+      push: (href, options = {}) => {
         // @ts-ignore startTransition exists
         React.startTransition(() => {
-          navigate(href, 'hard', 'push')
+          navigate(href, 'push', Boolean(options.forceOptimisticNavigation))
         })
       },
       reload: () => {
@@ -211,7 +245,6 @@ export default function AppRouter({
             type: ACTION_RELOAD,
 
             // TODO-APP: revisit if this needs to be passed.
-            url: new URL(window.location.href),
             cache: {
               data: null,
               subTreeData: null,
@@ -224,7 +257,7 @@ export default function AppRouter({
     }
 
     return routerInstance
-  }, [])
+  }, [dispatch, initialTree])
 
   useEffect(() => {
     // When mpaNavigation flag is set do a hard navigation to the new url.
@@ -245,47 +278,52 @@ export default function AppRouter({
     } else {
       window.history.replaceState(historyState, '', canonicalUrl)
     }
-  }, [tree, pushRef, canonicalUrl])
+
+    sync()
+  }, [tree, pushRef, canonicalUrl, sync])
 
   // Add `window.nd` for debugging purposes.
   // This is not meant for use in applications as concurrent rendering will affect the cache/tree/router.
   if (typeof window !== 'undefined') {
     // @ts-ignore this is for debugging
-    window.nd = { router: appRouter, cache, tree }
+    window.nd = { router: appRouter, cache, prefetchCache, tree }
   }
 
   /**
    * Handle popstate event, this is used to handle back/forward in the browser.
    * By default dispatches ACTION_RESTORE, however if the history entry was not pushed/replaced by app-router it will reload the page.
    * That case can happen when the old router injected the history entry.
    */
-  const onPopState = React.useCallback(({ state }: PopStateEvent) => {
-    if (!state) {
-      // TODO-APP: this case only happens when pushState/replaceState was called outside of Next.js. It should probably reload the page in this case.
-      return
-    }
+  const onPopState = useCallback(
+    ({ state }: PopStateEvent) => {
+      if (!state) {
+        // TODO-APP: this case only happens when pushState/replaceState was called outside of Next.js. It should probably reload the page in this case.
+        return
+      }
 
-    // TODO-APP: this case happens when pushState/replaceState was called outside of Next.js or when the history entry was pushed by the old router.
-    // It reloads the page in this case but we might have to revisit this as the old router ignores it.
-    if (!state.__NA) {
-      window.location.reload()
-      return
-    }
+      // TODO-APP: this case happens when pushState/replaceState was called outside of Next.js or when the history entry was pushed by the old router.
+      // It reloads the page in this case but we might have to revisit this as the old router ignores it.
+      if (!state.__NA) {
+        window.location.reload()
+        return
+      }
 
-    // @ts-ignore useTransition exists
-    // TODO-APP: Ideally the back button should not use startTransition as it should apply the updates synchronously
-    // Without startTransition works if the cache is there for this path
-    React.startTransition(() => {
-      dispatch({
-        type: ACTION_RESTORE,
-        url: new URL(window.location.href),
-        tree: state.tree,
+      // @ts-ignore useTransition exists
+      // TODO-APP: Ideally the back button should not use startTransition as it should apply the updates synchronously
+      // Without startTransition works if the cache is there for this path
+      React.startTransition(() => {
+        dispatch({
+          type: ACTION_RESTORE,
+          url: new URL(window.location.href),
+          tree: state.tree,
+        })
       })
-    })
-  }, [])
+    },
+    [dispatch]
+  )
 
   // Register popstate event to call onPopstate.
-  React.useEffect(() => {
+  useEffect(() => {
     window.addEventListener('popstate', onPopState)
     return () => {
       window.removeEventListener('popstate', onPopState)