feat(react-server): generate mapping of server route to client assets

packages/react-server/package.json

@@ -66,6 +66,7 @@
   "dependencies": {
     "@hiogawa/transforms": "workspace:*",
     "@tanstack/history": "^1.15.13",
+    "fast-glob": "^3.3.2",
     "use-sync-external-store": "^1.2.0"
   },
   "devDependencies": {

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

@@ -0,0 +1,110 @@
+import path from "node:path";
+import { objectMapValues, tinyassert, uniq } from "@hiogawa/utils";
+import FastGlob from "fast-glob";
+import type { Plugin, Rollup } from "vite";
+import type { PluginStateManager } from "../../plugin";
+import { type CustomModuleMeta, createVirtualPlugin } from "../../plugin/utils";
+
+export function routeManifestPluginServer({
+  manager,
+}: { manager: PluginStateManager }): Plugin[] {
+  return [
+    {
+      name: "server-route-manifest",
+      apply: "build",
+      async buildEnd() {
+        if (manager.buildType === "rsc") {
+          const routeFiles = await FastGlob(
+            "./src/routes/**/(page|layout|error).(js|jsx|ts|tsx)",
+          );
+          for (const routeFile of routeFiles) {
+            const absFile = path.join(manager.config.root, routeFile);
+            const deps = collectModuleDeps(absFile, this);
+            let ids: string[] = [];
+            for (const id of deps) {
+              const info = this.getModuleInfo(id);
+              tinyassert(info);
+              const meta = info.meta as CustomModuleMeta;
+              if (meta.$$rsc?.type === "client") {
+                ids.push(id);
+              }
+            }
+            manager.routeToClientReferences[routeFile] = ids;
+          }
+        }
+      },
+    },
+  ];
+}
+
+export function routeManifestPluginClient({
+  manager,
+}: { manager: PluginStateManager }): Plugin[] {
+  return [
+    {
+      name: routeManifestPluginClient.name + ":bundle",
+      apply: "build",
+      generateBundle(_options, bundle) {
+        if (manager.buildType === "client") {
+          const facadeModuleDeps: Record<string, string[]> = {};
+          for (const [k, v] of Object.entries(bundle)) {
+            if (v.type === "chunk" && v.facadeModuleId) {
+              facadeModuleDeps[v.facadeModuleId] = collectAssetDeps(k, bundle);
+            }
+          }
+          manager.routeToClientAssets = objectMapValues(
+            manager.routeToClientReferences,
+            (ids) =>
+              uniq(
+                ids.flatMap((id) => {
+                  const deps = facadeModuleDeps[id];
+                  tinyassert(deps);
+                  return deps;
+                }),
+              ),
+          );
+          console.log(manager.routeToClientAssets);
+        }
+      },
+    },
+    createVirtualPlugin("route-manifest", () => {
+      return `export default "todo"`;
+    }),
+  ];
+}
+
+function collectModuleDeps(id: string, ctx: Rollup.PluginContext) {
+  const visited = new Set<string>();
+  const recurse = (id: string) => {
+    if (visited.has(id)) {
+      return;
+    }
+    visited.add(id);
+    const info = ctx.getModuleInfo(id);
+    tinyassert(info);
+    for (const imported of info.importedIds) {
+      recurse(imported);
+    }
+  };
+  recurse(id);
+  return [...visited];
+}
+
+function collectAssetDeps(fileName: string, bundle: Rollup.OutputBundle) {
+  const visited = new Set<string>();
+
+  function recurse(k: string) {
+    if (visited.has(k)) return;
+    visited.add(k);
+    const v = bundle[k];
+    tinyassert(v);
+    if (v.type === "chunk") {
+      for (const k2 of v.imports) {
+        recurse(k2);
+      }
+    }
+  }
+
+  recurse(fileName);
+  return [...visited];
+}

packages/react-server/src/features/use-client/plugin.ts

@@ -13,6 +13,7 @@ import {
 } from "vite";
 import type { PluginStateManager } from "../../plugin";
 import {
+  type CustomModuleMeta,
   USE_CLIENT,
   USE_CLIENT_RE,
   createVirtualPlugin,
@@ -147,10 +148,19 @@ export function vitePluginServerUseClient({
       );
       manager.rscUseClientIds.add(id);
       if (manager.buildType === "scan") {
-        // only collect references without transform during the scan
+        // to discover server references imported only by client
+        // we keep code as is and continue crawling
         return;
       }
-      return { code: output.toString(), map: output.generateMap() };
+      return {
+        code: output.toString(),
+        map: output.generateMap(),
+        meta: {
+          $$rsc: {
+            type: "client",
+          },
+        } satisfies CustomModuleMeta,
+      };
     },
 
     /**

packages/react-server/src/plugin/index.ts

@@ -5,6 +5,7 @@ import {
   type InlineConfig,
   type Plugin,
   type PluginOption,
+  type ResolvedConfig,
   type ViteDevServer,
   build,
   createLogger,
@@ -14,6 +15,10 @@ import {
   SERVER_CSS_PROXY,
   vitePluginServerAssets,
 } from "../features/assets/plugin";
+import {
+  routeManifestPluginClient,
+  routeManifestPluginServer,
+} from "../features/router/plugin";
 import {
   vitePluginClientUseServer,
   vitePluginServerUseServer,
@@ -51,10 +56,14 @@ export type { PluginStateManager };
 
 class PluginStateManager {
   server?: ViteDevServer;
+  config!: ResolvedConfig;
   configEnv!: ConfigEnv;
 
   buildType?: "scan" | "rsc" | "client" | "ssr";
 
+  routeToClientReferences: Record<string, string[]> = {};
+  routeToClientAssets: Record<string, string[]> = {};
+
   // expose "use client" node modules to client via virtual modules
   // to avoid dual package due to deps optimization hash during dev
   nodeModules = {
@@ -134,6 +143,8 @@ export function vitePluginReactServer(options?: {
         runtimePath: RUNTIME_REACT_SERVER_PATH,
       }),
 
+      routeManifestPluginServer({ manager }),
+
       // this virtual is not necessary anymore but has been used in the past
       // to extend user's react-server entry like ENTRY_CLIENT_WRAPPER
       createVirtualPlugin(
@@ -219,6 +230,9 @@ export function vitePluginReactServer(options?: {
         },
       };
     },
+    configResolved(config) {
+      manager.config = config;
+    },
     async configureServer(server) {
       manager.server = server;
     },
@@ -285,6 +299,7 @@ export function vitePluginReactServer(options?: {
         await build(reactServerViteConfig);
         console.log("▶▶▶ REACT SERVER BUILD (server) [2/4]");
         manager.buildType = "rsc";
+        manager.rscUseClientIds.clear();
         await build(reactServerViteConfig);
         console.log("▶▶▶ REACT SERVER BUILD (browser) [3/4]");
         manager.buildType = "client";
@@ -315,6 +330,7 @@ export function vitePluginReactServer(options?: {
     }),
     ...vitePluginClientUseClient({ manager }),
     ...vitePluginServerAssets({ manager }),
+    ...routeManifestPluginClient({ manager }),
     createVirtualPlugin(ENTRY_CLIENT_WRAPPER.slice("virtual:".length), () => {
       // dev
       if (!manager.buildType) {

packages/react-server/src/plugin/utils.tsx

@@ -8,6 +8,13 @@ export function invalidateModule(server: ViteDevServer, id: string) {
   }
 }
 
+// cf. https://github.com/vercel/next.js/blob/5ae286ffd664e5c76841ed64f6e2da85a0835922/packages/next/src/build/webpack/loaders/get-module-build-info.ts#L8
+export type CustomModuleMeta = {
+  $$rsc?: {
+    type: "client";
+  };
+};
+
 // TODO: configurable?
 export const ENTRY_CLIENT = "/src/entry-client";
 export const ENTRY_REACT_SERVER = "/src/entry-react-server";