Add heuristic for optimistic push (#38599)

Co-authored-by: kodiakhq[bot] <49736102+kodiakhq[bot]@users.noreply.github.com>

packages/next/client/components/reducer.ts

@@ -149,6 +149,43 @@ const fillCacheWithDataProperty = (
   )
 }
 
+const canOptimisticallyRender = (
+  segments: string[],
+  flightRouterState: FlightRouterState
+): boolean => {
+  const segment = segments[0]
+  const isLastSegment = segments.length === 1
+
+  const [existingSegment, existingParallelRoutes, , , loadingMarker] =
+    flightRouterState
+
+  const hasLoading = loadingMarker === 'loading'
+
+  // If the tree path holds at least one loading.js it will be optimistic
+  if (hasLoading) {
+    return true
+  }
+
+  // Above already catches the last segment case where `hasLoading` is true, so in this case it would always be `false`.
+  if (isLastSegment) {
+    return false
+  }
+
+  // If the segments mismatch we can't resolve deeper into the tree
+  const segmentMatches = matchSegment(existingSegment, segment)
+
+  // If the existingParallelRoutes does not have a `children` parallelRouteKey we can't resolve deeper into the tree
+  if (!segmentMatches || !existingParallelRoutes.children) {
+    return hasLoading
+  }
+
+  // Resolve deeper in the tree as the current level did not have a loading marker
+  return canOptimisticallyRender(
+    segments.slice(1),
+    existingParallelRoutes.children
+  )
+}
+
 const createOptimisticTree = (
   segments: string[],
   flightRouterState: FlightRouterState | null,
@@ -163,11 +200,12 @@ const createOptimisticTree = (
   const segment = segments[0]
   const isLastSegment = segments.length === 1
 
-  const shouldRefetchThisLevel =
-    !flightRouterState || segment !== flightRouterState[0]
+  const segmentMatches =
+    existingSegment !== null && matchSegment(existingSegment, segment)
+  const shouldRefetchThisLevel = !flightRouterState || !segmentMatches
 
   let parallelRoutes: FlightRouterState[1] = {}
-  if (existingSegment !== null && matchSegment(existingSegment, segment)) {
+  if (existingSegment !== null && segmentMatches) {
     parallelRoutes = existingParallelRoutes
   }
 
@@ -200,6 +238,11 @@ const createOptimisticTree = (
     result[2] = href
   }
 
+  // Copy the loading flag from existing tree
+  if (flightRouterState && flightRouterState[4]) {
+    result[4] = flightRouterState[4]
+  }
+
   return result
 }
 
@@ -215,7 +258,7 @@ const walkTreeWithFlightDataPath = (
     const tree: FlightRouterState = [...treePatch]
 
     if (url) {
-      tree.push(url)
+      tree[2] = url
     }
 
     return tree
@@ -246,7 +289,12 @@ const walkTreeWithFlightDataPath = (
   ]
 
   if (url) {
-    tree.push(url)
+    tree[2] = url
+  }
+
+  // Copy loading flag
+  if (flightSegmentPath[4]) {
+    tree[4] = flightSegmentPath[4]
   }
 
   return tree
@@ -353,7 +401,7 @@ export function reducer(
       }
 
       // TODO-APP: flag on the tree of which part of the tree for if there is a loading boundary
-      const isOptimistic = false
+      const isOptimistic = canOptimisticallyRender(segments, state.tree)
 
       if (isOptimistic) {
         // Build optimistic tree
@@ -368,6 +416,7 @@ export function reducer(
 
         // Fill in the cache with blank that holds the `data` field.
         // TODO-APP: segments.slice(1) strips '', we can get rid of '' altogether.
+        cache.subTreeData = state.cache.subTreeData
         const res = fillCacheWithDataProperty(
           cache,
           state.cache,

packages/next/server/app-render.tsx

@@ -264,7 +264,8 @@ export type FlightRouterState = [
   segment: Segment,
   parallelRoutes: { [parallelRouterKey: string]: FlightRouterState },
   url?: string,
-  refresh?: 'refetch'
+  refresh?: 'refetch',
+  loading?: 'loading'
 ]
 
 export type FlightSegmentPath =
@@ -508,7 +509,9 @@ export async function renderToHTML(
   const createFlightRouterStateFromLoaderTree = ([
     segment,
     parallelRoutes,
+    { loading },
   ]: LoaderTree): FlightRouterState => {
+    const hasLoading = Boolean(loading)
     const dynamicParam = getDynamicParamFromSegment(segment)
 
     const segmentTree: FlightRouterState = [
@@ -529,6 +532,10 @@ export async function renderToHTML(
         {} as FlightRouterState[1]
       )
     }
+
+    if (hasLoading) {
+      segmentTree[4] = 'loading'
+    }
     return segmentTree
   }