feat(rsc): ability to merge client reference chunks (#766)

Co-authored-by: Claude <[email protected]>

Author: hi-ogawa

packages/plugin-rsc/e2e/basic.test.ts

@@ -9,6 +9,8 @@ import {
   waitForHydration,
 } from './helper'
 import { x } from 'tinyexec'
+import { normalizePath, type Rollup } from 'vite'
+import path from 'node:path'
 
 test.describe('dev-default', () => {
   const f = useFixture({ root: 'examples/basic', mode: 'dev' })
@@ -45,6 +47,17 @@ test.describe('dev-initial', () => {
 test.describe('build-default', () => {
   const f = useFixture({ root: 'examples/basic', mode: 'build' })
   defineTest(f)
+
+  test('custom client chunk', async () => {
+    const { chunks }: { chunks: Rollup.OutputChunk[] } = JSON.parse(
+      f.createEditor('dist/client/.vite/test.json').read(),
+    )
+    const chunk = chunks.find((c) => c.name === 'custom-chunk')
+    const expected = [1, 2, 3].map((i) =>
+      normalizePath(path.join(f.root, `src/routes/chunk/client${i}.tsx`)),
+    )
+    expect(chunk?.moduleIds).toEqual(expect.arrayContaining(expected))
+  })
 })
 
 test.describe('dev-non-optimized-cjs', () => {
@@ -143,7 +156,7 @@ test.describe('build-stable-chunks', () => {
       .sort()
     expect(newChunks).toEqual([
       'src/framework/entry.browser.tsx',
-      'src/routes/client.tsx',
+      'virtual:vite-rsc/client-references/group/src/routes/client.tsx',
     ])
     expect(oldChunks).toEqual(newChunks)
   })

packages/plugin-rsc/examples/basic/src/routes/chunk/client1.tsx

@@ -0,0 +1,9 @@
+'use client'
+
+export function TestClientChunk1() {
+  return <span>test-chunk1</span>
+}
+
+export function TestClientChunkConflict() {
+  return <span>test-chunk-conflict1</span>
+}

packages/plugin-rsc/examples/basic/src/routes/chunk/client2.tsx

@@ -0,0 +1,5 @@
+'use client'
+
+export function TestClientChunk2() {
+  return <span>test-chunk2</span>
+}

packages/plugin-rsc/examples/basic/src/routes/chunk/client3.tsx

@@ -0,0 +1,9 @@
+'use client'
+
+export function TestClientChunk3() {
+  return <span>test-chunk3</span>
+}
+
+export function TestClientChunkConflict() {
+  return <span>test-chunk-conflict3</span>
+}

packages/plugin-rsc/examples/basic/src/routes/chunk/server.tsx

@@ -0,0 +1,21 @@
+import {
+  TestClientChunk1,
+  TestClientChunkConflict as TestClientChunkConflict1,
+} from './client1'
+import { TestClientChunk2 } from './client2'
+import {
+  TestClientChunk3,
+  TestClientChunkConflict as TestClientChunkConflict3,
+} from './client3'
+
+export function TestClientChunkServer() {
+  return (
+    <div data-testid="test-client-chunk">
+      <TestClientChunk1 />|
+      <TestClientChunkConflict1 />|
+      <TestClientChunk2 />|
+      <TestClientChunk3 />|
+      <TestClientChunkConflict3 />
+    </div>
+  )
+}

packages/plugin-rsc/examples/basic/src/routes/root.tsx

@@ -41,6 +41,7 @@ import { TestAssetsServer } from './assets/server'
 import { TestHmrSwitchServer } from './hmr-switch/server'
 import { TestHmrSwitchClient } from './hmr-switch/client'
 import { TestTreeShakeServer } from './tree-shake/server'
+import { TestClientChunkServer } from './chunk/server'
 
 export function Root(props: { url: URL }) {
   return (
@@ -95,6 +96,7 @@ export function Root(props: { url: URL }) {
         <TestImportMetaGlob />
         <TestAssetsServer />
         <TestTreeShakeServer />
+        <TestClientChunkServer />
       </body>
     </html>
   )

packages/plugin-rsc/examples/basic/vite.config.ts

@@ -11,6 +11,7 @@ import {
 } from 'vite'
 // import inspect from 'vite-plugin-inspect'
 import path from 'node:path'
+import fs from 'node:fs'
 import { fileURLToPath } from 'node:url'
 
 export default defineConfig({
@@ -30,6 +31,11 @@ export default defineConfig({
       rscCssTransform: false,
       copyServerAssetsToClient: (fileName) =>
         fileName !== '__server_secret.txt',
+      clientChunks(id) {
+        if (id.includes('/src/routes/chunk/')) {
+          return 'custom-chunk'
+        }
+      },
     }),
     {
       name: 'test-tree-shake',
@@ -43,6 +49,23 @@ export default defineConfig({
         }
       },
     },
+    {
+      // dump entire bundle to analyze build output for e2e
+      name: 'test-metadata',
+      enforce: 'post',
+      writeBundle(options, bundle) {
+        const chunks: Rollup.OutputChunk[] = []
+        for (const chunk of Object.values(bundle)) {
+          if (chunk.type === 'chunk') {
+            chunks.push(chunk)
+          }
+        }
+        fs.writeFileSync(
+          path.join(options.dir!, '.vite/test.json'),
+          JSON.stringify({ chunks }, null, 2),
+        )
+      },
+    },
     {
       name: 'test-server-assets-security',
       buildStart() {

packages/plugin-rsc/src/plugin.ts

@@ -93,6 +93,8 @@ class RscPluginManager {
   buildAssetsManifest: AssetsManifest | undefined
   isScanBuild: boolean = false
   clientReferenceMetaMap: Record<string, ClientReferenceMeta> = {}
+  clientReferenceGroups: Record</* group name*/ string, ClientReferenceMeta[]> =
+    {}
   serverReferenceMetaMap: Record<string, ServerRerferenceMeta> = {}
   serverResourcesMetaMap: Record<string, { key: string }> = {}
 
@@ -176,6 +178,17 @@ export type RscPluginOptions = {
     ssr?: string
     rsc?: string
   }
+
+  /**
+   * Custom chunking strategy for client reference modules.
+   *
+   * This function allows you to group multiple client components into
+   * custom chunks instead of having each module in its own chunk.
+   *
+   * @param id - The absolute path of the client module
+   * @returns The chunk name to group this module with, or undefined to use default behavior
+   */
+  clientChunks?: (id: string) => string | undefined
 }
 
 /** @experimental */
@@ -977,7 +990,7 @@ function scanBuildStripPlugin({
 function vitePluginUseClient(
   useClientPluginOptions: Pick<
     RscPluginOptions,
-    'keepUseCientProxy' | 'environment'
+    'keepUseCientProxy' | 'environment' | 'clientChunks'
   >,
   manager: RscPluginManager,
 ): Plugin[] {
@@ -1107,28 +1120,77 @@ function vitePluginUseClient(
         return { code: output.toString(), map: { mappings: '' } }
       },
     },
-    createVirtualPlugin('vite-rsc/client-references', function () {
-      if (this.environment.mode === 'dev') {
-        return { code: `export default {}`, map: null }
-      }
-      let code = ''
-      for (const meta of Object.values(manager.clientReferenceMetaMap)) {
-        // vite/rollup can apply tree-shaking to dynamic import of this form
-        const key = JSON.stringify(meta.referenceKey)
-        const id = JSON.stringify(meta.importId)
-        const exports = meta.renderedExports
-          .map((name) => (name === 'default' ? 'default: _default' : name))
-          .sort()
-        code += `
-          ${key}: async () => {
-            const {${exports}} = await import(${id});
-            return {${exports}};
-          },
-        `
-      }
-      code = `export default {${code}};\n`
-      return { code, map: null }
-    }),
+    {
+      name: 'rsc:use-client/build-references',
+      resolveId(source) {
+        if (source.startsWith('virtual:vite-rsc/client-references')) {
+          return '\0' + source
+        }
+      },
+      load(id) {
+        if (id === '\0virtual:vite-rsc/client-references') {
+          // not used during dev
+          if (this.environment.mode === 'dev') {
+            return { code: `export default {}`, map: null }
+          }
+          // no custom chunking needed for scan
+          if (manager.isScanBuild) {
+            let code = ``
+            for (const meta of Object.values(manager.clientReferenceMetaMap)) {
+              code += `import ${JSON.stringify(meta.importId)};\n`
+            }
+            return { code, map: null }
+          }
+          let code = ''
+          // group client reference modules by `clientChunks` option
+          manager.clientReferenceGroups = {}
+          for (const meta of Object.values(manager.clientReferenceMetaMap)) {
+            const name =
+              useClientPluginOptions.clientChunks?.(meta.importId) ||
+              // use original module id as name by default
+              normalizePath(path.relative(manager.config.root, meta.importId))
+            const group = (manager.clientReferenceGroups[name] ??= [])
+            group.push(meta)
+          }
+          for (const [name, metas] of Object.entries(
+            manager.clientReferenceGroups,
+          )) {
+            const groupVirtual = `virtual:vite-rsc/client-references/group/${name}`
+            for (const meta of metas) {
+              code += `\
+                ${JSON.stringify(meta.referenceKey)}: async () => {
+                  const m = await import(${JSON.stringify(groupVirtual)});
+                  return m.export_${meta.referenceKey};
+                },
+              `
+            }
+          }
+          code = `export default {${code}};\n`
+          return { code, map: null }
+        }
+        // re-export client reference modules from each group
+        if (id.startsWith('\0virtual:vite-rsc/client-references/group/')) {
+          const name = id.slice(
+            '\0virtual:vite-rsc/client-references/group/'.length,
+          )
+          const metas = manager.clientReferenceGroups[name]
+          assert(metas, `unknown client reference group: ${name}`)
+          let code = ``
+          for (const meta of metas) {
+            // pick only renderedExports to tree-shake unused client references
+            const exports = meta.renderedExports
+              .map((name) => `${name}: import_${meta.referenceKey}.${name},\n`)
+              .sort()
+              .join('')
+            code += `
+              import * as import_${meta.referenceKey} from ${JSON.stringify(meta.importId)};
+              export const export_${meta.referenceKey} = {${exports}};
+            `
+          }
+          return { code, map: null }
+        }
+      },
+    },
     {
       name: 'rsc:virtual-client-in-server-package',
       async load(id) {