Generate per-segment prefetch responses

Follow up to vercel#71113.

This generates a prefetch response for route segments, so the client
can request them individually, and reuse shared layouts between
separate page prefetches.

Since this is exclusive to PPR, segment prefetches are always static
and do not reach the application layer. They are served directly from
cache. (This part was implemented in the previous PR.)

The prerendered segment data is generated (and revalidated)
simultaneously with the prerender of the entire page, to ensure
consistency between parent and child segments. Since we always rebuild
the whole page, we can save an extra render of each segment by reusing
the Flight response that was used to generate the initial HTML.

We do this by decoding the page data using a server-side Flight
consumer, picking out an individual segment's data, then re-encoding it
as its own Flight response. We have to do this once per segment — not
ideal, but for now it's the only way to make sure the side effects from
the original Flight stream (e.g. Float preloads) are transferred to the
per-segment streams.

This PR does not yet include any updates to the client router.

packages/next/src/server/app-render/app-render.tsx

@@ -19,7 +19,10 @@ import type { RequestStore } from '../app-render/work-unit-async-storage.externa
 import type { NextParsedUrlQuery } from '../request-meta'
 import type { LoaderTree } from '../lib/app-dir-module'
 import type { AppPageModule } from '../route-modules/app-page/module'
-import type { ClientReferenceManifest } from '../../build/webpack/plugins/flight-manifest-plugin'
+import type {
+  ClientReferenceManifest,
+  ManifestNode,
+} from '../../build/webpack/plugins/flight-manifest-plugin'
 import type { DeepReadonly } from '../../shared/lib/deep-readonly'
 import type { BaseNextRequest, BaseNextResponse } from '../base-http'
 import type { IncomingHttpHeaders } from 'http'
@@ -92,7 +95,10 @@ import { makeGetServerInsertedHTML } from './make-get-server-inserted-html'
 import { walkTreeWithFlightRouterState } from './walk-tree-with-flight-router-state'
 import { createComponentTree } from './create-component-tree'
 import { getAssetQueryString } from './get-asset-query-string'
-import { setReferenceManifestsSingleton } from './encryption-utils'
+import {
+  setReferenceManifestsSingleton,
+  getServerModuleMap,
+} from './encryption-utils'
 import {
   DynamicState,
   type PostponedState,
@@ -1253,32 +1259,6 @@ async function renderToHTMLOrFlightImpl(
       }
     }
 
-    // Per-segment prefetch data
-    //
-    // All of the segments for a page are generated simultaneously, including
-    // during revalidations. This is to ensure consistency, because it's
-    // possible for a mismatch between a layout and page segment can cause the
-    // client to error during rendering. We want to preserve the ability of the
-    // client to recover from such a mismatch by re-requesting all the segments
-    // to get a consistent view of the page.
-    //
-    // TODO (Per Segment Prefetching): This is placeholder data. Populate with
-    // the actual data generated during prerender.
-    if (renderOpts.experimental.isRoutePPREnabled === true) {
-      const placeholder = Buffer.from(
-        'TODO (Per Segment Prefetching): Not yet implemented\n'
-      )
-      metadata.segmentFlightData = new Map([
-        // Root segment
-        ['/', placeholder],
-        ['/blog', placeholder],
-        // TODO: Update the client to use the same encoding for segment paths that
-        // we use here, so we don't have to convert between them. Needs to be
-        // filesystem safe.
-        ['/blog/[post]-1-d', placeholder],
-      ])
-    }
-
     return new RenderResult(await streamToString(response.stream), options)
   } else {
     // We're rendering dynamically
@@ -2761,7 +2741,14 @@ async function prerenderToStream(
           tracingMetadata: tracingMetadata,
         })
 
-        metadata.flightData = await streamToBuffer(reactServerResult.asStream())
+        const flightData = await streamToBuffer(reactServerResult.asStream())
+        metadata.flightData = flightData
+        metadata.segmentFlightData = await collectSegmentData(
+          finalAttemptRSCPayload,
+          flightData,
+          ComponentMod,
+          renderOpts
+        )
 
         if (serverIsDynamic || clientIsDynamic) {
           if (postponed != null) {
@@ -3212,9 +3199,16 @@ async function prerenderToStream(
         // const reactServerResult =
         //   await createReactServerPrerenderResultFromRender(reactServerStream!)
 
-        metadata.flightData = await streamToBuffer(
+        const flightData = await streamToBuffer(
           serverPrerenderStreamResult.asStream()
         )
+        metadata.flightData = flightData
+        metadata.segmentFlightData = await collectSegmentData(
+          finalServerPayload,
+          flightData,
+          ComponentMod,
+          renderOpts
+        )
 
         const getServerInsertedHTML = makeGetServerInsertedHTML({
           polyfills,
@@ -3337,6 +3331,12 @@ async function prerenderToStream(
 
       if (shouldGenerateStaticFlightData(workStore)) {
         metadata.flightData = flightData
+        metadata.segmentFlightData = await collectSegmentData(
+          RSCPayload,
+          flightData,
+          ComponentMod,
+          renderOpts
+        )
       }
 
       /**
@@ -3516,7 +3516,14 @@ async function prerenderToStream(
       )
 
       if (shouldGenerateStaticFlightData(workStore)) {
-        metadata.flightData = await streamToBuffer(reactServerResult.asStream())
+        const flightData = await streamToBuffer(reactServerResult.asStream())
+        metadata.flightData = flightData
+        metadata.segmentFlightData = await collectSegmentData(
+          RSCPayload,
+          flightData,
+          ComponentMod,
+          renderOpts
+        )
       }
 
       const getServerInsertedHTML = makeGetServerInsertedHTML({
@@ -3668,9 +3675,16 @@ async function prerenderToStream(
       })
 
       if (shouldGenerateStaticFlightData(workStore)) {
-        metadata.flightData = await streamToBuffer(
+        const flightData = await streamToBuffer(
           reactServerPrerenderResult.asStream()
         )
+        metadata.flightData = flightData
+        metadata.segmentFlightData = await collectSegmentData(
+          errorRSCPayload,
+          flightData,
+          ComponentMod,
+          renderOpts
+        )
       }
 
       const validateRootLayout = renderOpts.dev
@@ -3797,3 +3811,63 @@ const getGlobalErrorStyles = async (
 
   return globalErrorStyles
 }
+
+async function collectSegmentData(
+  rscPayload: InitialRSCPayload,
+  fullPageDataBuffer: Buffer,
+  ComponentMod: AppPageModule,
+  renderOpts: RenderOpts
+): Promise<Map<string, Buffer> | undefined> {
+  // Per-segment prefetch data
+  //
+  // All of the segments for a page are generated simultaneously, including
+  // during revalidations. This is to ensure consistency, because it's
+  // possible for a mismatch between a layout and page segment can cause the
+  // client to error during rendering. We want to preserve the ability of the
+  // client to recover from such a mismatch by re-requesting all the segments
+  // to get a consistent view of the page.
+  //
+  // For performance, we reuse the Flight output that was created when
+  // generating the initial page HTML. The Flight stream for the whole page is
+  // decomposed into a separate stream per segment.
+
+  const clientReferenceManifest = renderOpts.clientReferenceManifest
+  if (
+    !clientReferenceManifest ||
+    renderOpts.experimental.isRoutePPREnabled !== true
+  ) {
+    return
+  }
+
+  // FlightDataPath is an unsound type, hence the additional checks.
+  const flightDataPaths = rscPayload.f
+  if (flightDataPaths.length !== 1 && flightDataPaths[0].length !== 3) {
+    console.error(
+      'Internal Next.js error: InitialRSCPayload does not match the expected ' +
+        'shape for a prerendered page during segment prefetch generation.'
+    )
+    return
+  }
+  const routeTree: FlightRouterState = flightDataPaths[0][0]
+
+  // Manifest passed to the Flight client for reading the full-page Flight
+  // stream. Based off similar code in use-cache-wrapper.ts.
+  const isEdgeRuntime = process.env.NEXT_RUNTIME === 'edge'
+  const serverConsumerManifest = {
+    // moduleLoading must be null because we don't want to trigger preloads of ClientReferences
+    // to be added to the consumer. Instead, we'll wait for any ClientReference to be emitted
+    // which themselves will handle the preloading.
+    moduleLoading: null,
+    moduleMap: isEdgeRuntime
+      ? clientReferenceManifest.edgeRscModuleMapping
+      : clientReferenceManifest.rscModuleMapping,
+    serverModuleMap: getServerModuleMap(),
+  }
+
+  return await ComponentMod.collectSegmentData(
+    routeTree,
+    fullPageDataBuffer,
+    clientReferenceManifest.clientModules as ManifestNode,
+    serverConsumerManifest
+  )
+}