[PPR Nav] Fix flash of loading state during back/forward

A popstate navigation reads data from the local cache. It does not issue
new network requests (unless the cache entries have been evicted). So,
when navigating with back/forward, we should not switch back to the PPR
loading state. We should render the full, cached dynamic data
immediately.

To implement this, on a popstate navigation, we update the cache to drop
the prefetch data for any segment whose dynamic data was already
received. We clone the entire cache node tree and set the `prefetchRsc`
field to `null` to prevent it from being rendered. (We can't mutate the
node in place because Cache Node is a concurrent data structure.)

Tehnically, what we're actually checking is whether the dynamic network
response was received. But since it's a streaming response, this does
not mean that all the dynamic data has fully streamed in. It just means
that _some_ of the dynamic data was received. But as a heuristic, we
assume that the rest dynamic data will stream in quickly, so it's still
better to skip the prefetch state.

packages/next/src/client/components/router-reducer/ppr-navigations.ts

@@ -727,6 +727,67 @@ function abortPendingCacheNode(
   }
 }
 
+export function updateCacheNodeOnPopstateRestoration(
+  oldCacheNode: CacheNode,
+  routerState: FlightRouterState
+) {
+  // A popstate navigation reads data from the local cache. It does not issue
+  // new network requests (unless the cache entries have been evicted). So, we
+  // update the cache to drop the prefetch  data for any segment whose dynamic
+  // data was already received. This prevents an unnecessary flash back to PPR
+  // state during a back/forward navigation.
+  //
+  // This function clones the entire cache node tree and sets the `prefetchRsc`
+  // field to `null` to prevent it from being rendered. We can't mutate the node
+  // in place because this is a concurrent data structure.
+
+  const routerStateChildren = routerState[1]
+  const oldParallelRoutes = oldCacheNode.parallelRoutes
+  const newParallelRoutes = new Map(oldParallelRoutes)
+  for (let parallelRouteKey in routerStateChildren) {
+    const routerStateChild: FlightRouterState =
+      routerStateChildren[parallelRouteKey]
+    const segmentChild = routerStateChild[0]
+    const segmentKeyChild = createRouterCacheKey(segmentChild)
+    const oldSegmentMapChild = oldParallelRoutes.get(parallelRouteKey)
+    if (oldSegmentMapChild !== undefined) {
+      const oldCacheNodeChild = oldSegmentMapChild.get(segmentKeyChild)
+      if (oldCacheNodeChild !== undefined) {
+        const newCacheNodeChild = updateCacheNodeOnPopstateRestoration(
+          oldCacheNodeChild,
+          routerStateChild
+        )
+        const newSegmentMapChild = new Map(oldSegmentMapChild)
+        newSegmentMapChild.set(segmentKeyChild, newCacheNodeChild)
+        newParallelRoutes.set(parallelRouteKey, newSegmentMapChild)
+      }
+    }
+  }
+
+  // Only show prefetched data if the dynamic data is still pending.
+  //
+  // Tehnically, what we're actually checking is whether the dynamic network
+  // response was received. But since it's a streaming response, this does not
+  // mean that all the dynamic data has fully streamed in. It just means that
+  // _some_ of the dynamic data was received. But as a heuristic, we assume that
+  // the rest dynamic data will stream in quickly, so it's still better to skip
+  // the prefetch state.
+  const rsc = oldCacheNode.rsc
+  const shouldUsePrefetch = isDeferredRsc(rsc) && rsc.status === 'pending'
+
+  return {
+    lazyData: null,
+    rsc,
+    head: oldCacheNode.head,
+
+    prefetchHead: shouldUsePrefetch ? oldCacheNode.prefetchHead : null,
+    prefetchRsc: shouldUsePrefetch ? oldCacheNode.prefetchRsc : null,
+
+    // These are the cloned children we computed above
+    parallelRoutes: newParallelRoutes,
+  }
+}
+
 const DEFERRED = Symbol()
 
 type PendingDeferredRsc = Promise<React.ReactNode> & {

packages/next/src/client/components/router-reducer/reducers/restore-reducer.ts

@@ -5,6 +5,7 @@ import type {
   RestoreAction,
 } from '../router-reducer-types'
 import { extractPathFromFlightRouterState } from '../compute-changed-path'
+import { updateCacheNodeOnPopstateRestoration } from '../ppr-navigations'
 
 export function restoreReducer(
   state: ReadonlyReducerState,
@@ -13,6 +14,15 @@ export function restoreReducer(
   const { url, tree } = action
   const href = createHrefFromUrl(url)
 
+  const oldCache = state.cache
+  const newCache = process.env.__NEXT_PPR
+    ? // When PPR is enabled, we update the cache to drop the prefetch
+      // data for any segment whose dynamic data was already received. This
+      // prevents an unnecessary flash back to PPR state during a
+      // back/forward navigation.
+      updateCacheNodeOnPopstateRestoration(oldCache, tree)
+    : oldCache
+
   return {
     buildId: state.buildId,
     // Set canonical url
@@ -24,7 +34,7 @@ export function restoreReducer(
       preserveCustomHistoryState: true,
     },
     focusAndScrollRef: state.focusAndScrollRef,
-    cache: state.cache,
+    cache: newCache,
     prefetchCache: state.prefetchCache,
     // Restore provided tree
     tree: tree,