wip: rework route data fetching (2)

packages/react-server/src/entry/react-server.tsx

@@ -6,11 +6,14 @@ import {
   generateRouteModuleTree,
   renderRouteMap,
 } from "../features/router/server";
+import type { RouteType } from "../features/router/tree";
 import {
   type LayoutRequest,
+  type RouteDataKey,
   type ServerRouterData,
   createLayoutContentRequest,
   getNewLayoutContentKeys,
+  getRouteMapping,
 } from "../features/router/utils";
 import { runActionContext } from "../features/server-action/context";
 import {
@@ -71,6 +74,37 @@ export const handler: ReactServerHandler = async (ctx) => {
     });
   }
 
+  // render flight
+  {
+    const mapping = getRouteMapping(url.pathname);
+    // TODO: we should let browser tell server about this?
+    // well, we're using `lastPathname` as a minimal information to do that...
+    const skipKeys: RouteDataKey[] = [];
+    if (
+      streamParam?.lastPathname &&
+      !streamParam.revalidate &&
+      !actionResult?.context.revalidate
+    ) {
+      const lastMapping = getRouteMapping(streamParam.lastPathname);
+      // "page" always rerender
+      // "layout" don't return if it's alway in client
+    }
+    const { entries } = await renderRouteMap(router.tree, request);
+
+    const stream = ReactServer.renderToReadableStream<ServerRouterData>(
+      {
+        entries: {} as any,
+        action: actionResult
+          ? objectPick(actionResult, ["data", "error"])
+          : undefined,
+      },
+      createBundlerConfig(),
+      {
+        onError: reactServerOnError,
+      },
+    );
+  }
+
   let layoutRequest = createLayoutContentRequest(url.pathname);
   if (
     streamParam?.lastPathname &&
@@ -114,6 +148,8 @@ async function render({
   // we only need to filter out unnecessary map.
   // then client can map back to each prefix.
   const result = await renderRouteMap(router.tree, request);
+  result.entries;
+  // revalidated
   const entries = {
     layouts: objectPick(
       result.layouts,
@@ -132,6 +168,11 @@ async function render({
   return ReactServer.renderToReadableStream<ServerRouterData>(
     {
       entries,
+      result2: {
+        mapping: {},
+        // TODO: filter out based on
+        entries: result.entries,
+      },
       action: actionResult
         ? objectPick(actionResult, ["data", "error"])
         : undefined,

packages/react-server/src/features/router/client.tsx

@@ -23,6 +23,7 @@ export const LayoutStateContext = React.createContext<LayoutStateContextType>(
   undefined!,
 );
 
+// TODO: name -> id
 export function LayoutContent(props: { name: string }) {
   const ctx = React.useContext(LayoutStateContext);
   const data = React.use(ctx.data);

packages/react-server/src/features/router/server.tsx

@@ -1,6 +1,12 @@
 import React from "react";
 import { type ReactServerErrorContext, createError } from "../../lib/error";
-import { type TreeNode, createFsRouteTree, matchRouteTree } from "./tree";
+import {
+  type TreeNode,
+  createFsRouteTree,
+  matchRouteTree,
+  toRouteId,
+} from "./tree";
+import type { RouteDataEntry, RouteDataKey } from "./utils";
 
 // cf. https://nextjs.org/docs/app/building-your-application/routing#file-conventions
 interface RouteEntry {
@@ -63,6 +69,7 @@ async function renderLayout(
 export async function renderRouteMap(
   tree: RouteModuleNode,
   request: Pick<Request, "url" | "headers">,
+  skipKeys: RouteDataKey[] = [],
 ) {
   const url = serializeUrl(new URL(request.url));
   const baseProps: Omit<BaseProps, "params"> = {
@@ -75,17 +82,24 @@ export async function renderRouteMap(
   const pages: Record<string, React.ReactNode> = {};
   const layouts: Record<string, React.ReactNode> = {};
   const result = matchRouteTree(tree, url.pathname);
+  const entries: RouteDataEntry[] = [];
   for (const m of result.matches) {
-    const props: BaseProps = { ...baseProps, params: m.params };
-    if (m.type === "layout") {
-      layouts[m.prefix] = await renderLayout(m.node, props, m.prefix);
-    } else if (m.type === "page") {
-      pages[m.prefix] = renderPage(m.node, props);
-    } else {
-      m.type satisfies never;
+    if (skipKeys.some((k) => k.type === m.type && k.prefix === m.prefix)) {
+      continue;
     }
+    const props: BaseProps = { ...baseProps, params: m.params };
+    entries.push({
+      type: m.type,
+      prefix: m.prefix,
+      node:
+        m.type === "layout"
+          ? await renderLayout(m.node, props, m.prefix)
+          : m.type === "page"
+            ? renderPage(m.node, props)
+            : null,
+    });
   }
-  return { pages, layouts };
+  return { pages, layouts, entries };
 }
 
 const ThrowNotFound: React.FC = () => {

packages/react-server/src/features/router/tree.ts

@@ -43,17 +43,29 @@ function sortDynamicRoutes<T>(tree: TreeNode<T>) {
 }
 
 type MatchNodeEntry<T> = {
+  type: RouteType;
   prefix: string;
-  type: "layout" | "page";
   node: TreeNode<T>;
   params: Record<string, string>;
 };
 
+export type RouteType = "layout" | "page";
+
+export function toRouteId(type: RouteType, pathname: string) {
+  return `${pathname}#${type}`;
+}
+
+export function parseRouteId(id: string) {
+  const [pathname, type] = id.split("#") as [string, RouteType];
+  return { pathname, type };
+}
+
 type MatchResult<T> = {
   matches: MatchNodeEntry<T>[];
 };
 
 export function matchRouteTree<T>(tree: TreeNode<T>, pathname: string) {
+  // TODO: force non-trailing slash at higher level
   // TODO: more uniform handling of trailing slash
   pathname = normalizePathname(pathname);
   const prefixes = getPathPrefixes(pathname);

packages/react-server/src/features/router/utils.tsx

@@ -1,4 +1,6 @@
+import type React from "react";
 import type { ActionResult } from "../server-action/react-server";
+import { type RouteType, toRouteId } from "./tree";
 
 // TODO: rename
 // TODO
@@ -8,20 +10,58 @@ export type LayoutRequest = Record<
   string,
   {
     type: "page" | "layout";
+    // TODO: rename to prefix
     name: string;
   }
 >;
 
+// pathname -> routeId
+export type RouteMapping = Record<string, string>;
+
+// routeId -> node
+export type RouteEntries = Record<string, React.ReactNode>;
+
+export type RouteDataKey = {
+  type: RouteType;
+  prefix: string;
+};
+
+export type RouteDataEntry = {
+  type: RouteType;
+  prefix: string;
+  node: React.ReactNode;
+};
+
 export type ServerRouterData = {
   action?: Pick<ActionResult, "error" | "data">;
-  entries: {
-    layouts: Record<string, React.ReactNode>;
-    pages: Record<string, React.ReactNode>;
-  };
+  entries: RouteDataEntry[];
+  // entries: {
+  //   layouts: Record<string, React.ReactNode>;
+  //   pages: Record<string, React.ReactNode>;
+  // };
+  // result2?: {
+  //   mapping: RouteMapping;
+  //   entries2: RouteDataEntry[];
+  // };
 };
 
 export const LAYOUT_ROOT_NAME = "__root";
 
+export function getRouteMapping(pathname: string): RouteMapping {
+  const map: RouteMapping = {};
+  map[LAYOUT_ROOT_NAME] = toRouteId("layout", "/");
+  const prefixes = getPathPrefixes(pathname);
+  for (let i = 0; i < prefixes.length; i++) {
+    const prefix = prefixes[i]!;
+    if (i < prefixes.length - 1) {
+      map[prefix] = toRouteId("layout", prefixes[i + 1]!);
+    } else {
+      map[prefix] = toRouteId("page", prefix);
+    }
+  }
+  return map;
+}
+
 export function createLayoutContentRequest(pathname: string): LayoutRequest {
   const prefixes = getPathPrefixes(pathname);
   const map: LayoutRequest = {
@@ -44,6 +84,7 @@ export function createLayoutContentRequest(pathname: string): LayoutRequest {
   return map;
 }
 
+// TODO: this is bad. should return array of LayoutContentTarget
 export function getNewLayoutContentKeys(prev: string, next: string): string[] {
   const prevMap = createLayoutContentRequest(prev);
   const nextMap = createLayoutContentRequest(next);