fix(rsc): expose only "use server" as server functions (#752)

Author: hi-ogawa

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

@@ -40,6 +40,7 @@ import { TestImportMetaGlob } from './import-meta-glob/server'
 import { TestAssetsServer } from './assets/server'
 import { TestHmrSwitchServer } from './hmr-switch/server'
 import { TestHmrSwitchClient } from './hmr-switch/client'
+import { TestTreeShakeServer } from './tree-shake/server'
 
 export function Root(props: { url: URL }) {
   return (
@@ -93,6 +94,7 @@ export function Root(props: { url: URL }) {
         <TestCssQueries />
         <TestImportMetaGlob />
         <TestAssetsServer />
+        <TestTreeShakeServer />
       </body>
     </html>
   )

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

@@ -0,0 +1,17 @@
+export function TestTreeShakeServer() {
+  return (
+    <form
+      action={async () => {
+        'use server'
+        console.log('test-tree-shake-server')
+      }}
+    >
+      <button>test-tree-shake-server</button>
+    </form>
+  )
+}
+
+// this should not be exported as server functions
+export function __unused_server_export__() {
+  console.log('__unused_server_export__')
+}

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

@@ -25,12 +25,13 @@ export default defineConfig({
         fileName !== '__server_secret.txt',
     }),
     {
-      name: 'test-client-reference-tree-shaking',
+      name: 'test-tree-shake',
       enforce: 'post',
       writeBundle(_options, bundle) {
         for (const chunk of Object.values(bundle)) {
           if (chunk.type === 'chunk') {
             assert(!chunk.code.includes('__unused_client_reference__'))
+            assert(!chunk.code.includes('__unused_server_export__'))
           }
         }
       },

packages/plugin-rsc/src/plugin.ts

@@ -59,9 +59,8 @@ type ClientReferenceMeta = {
 type ServerRerferenceMeta = {
   importId: string
   referenceKey: string
-  // TODO: expose only "use server" exports
-  // TODO: tree shake unused exports
-  // exportNames: string[]
+  // TODO: tree shake unused server functions
+  exportNames: string[]
 }
 
 const PKG_NAME = '@vitejs/plugin-rsc'
@@ -1458,6 +1457,7 @@ function vitePluginUseServer(
           manager.serverReferenceMetaMap[id] = {
             importId: id,
             referenceKey: getNormalizedId(),
+            exportNames: 'names' in result ? result.names : result.exportNames,
           }
           const importSource = resolvePackage(`${PKG_NAME}/react/rsc`)
           output.prepend(`import * as $$ReactServer from "${importSource}";\n`)
@@ -1499,6 +1499,7 @@ function vitePluginUseServer(
           manager.serverReferenceMetaMap[id] = {
             importId: id,
             referenceKey: getNormalizedId(),
+            exportNames: result.exportNames,
           }
           const name =
             this.environment.name === browserEnvironmentName ? 'browser' : 'ssr'
@@ -1519,7 +1520,15 @@ function vitePluginUseServer(
       for (const meta of Object.values(manager.serverReferenceMetaMap)) {
         const key = JSON.stringify(meta.referenceKey)
         const id = JSON.stringify(meta.importId)
-        code += `${key}: () => import(${id}),`
+        const exports = meta.exportNames
+          .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 }

packages/plugin-rsc/src/transforms/hoist.test.ts

@@ -35,6 +35,16 @@ describe(transformHoistInlineDirective, () => {
     return output.toString()
   }
 
+  async function testTransformNames(input: string) {
+    const ast = await parseAstAsync(input)
+    const result = transformHoistInlineDirective(input, ast, {
+      runtime: (value, name) =>
+        `$$register(${value}, "<id>", ${JSON.stringify(name)})`,
+      directive: 'use server',
+    })
+    return result.names
+  }
+
   it('none', async () => {
     const input = `
 const x = "x";
@@ -104,6 +114,14 @@ export default function w() {
     expect(await testTransform(input, { encode: true })).toBe(
       await testTransform(input),
     )
+
+    expect(await testTransformNames(input)).toMatchInlineSnapshot(`
+      [
+        "$$hoist_0_f",
+        "$$hoist_1_h",
+        "$$hoist_2_w",
+      ]
+    `)
   })
 
   it('closure', async () => {

packages/plugin-rsc/src/transforms/wrap-export.test.ts

@@ -23,6 +23,16 @@ async function testTransform(
   return output.hasChanged() && output.toString()
 }
 
+async function testTransformNames(input: string) {
+  const ast = await parseAstAsync(input)
+  const result = transformWrapExport(input, ast, {
+    runtime: (value, name) =>
+      `$$wrap(${value}, "<id>", ${JSON.stringify(name)})`,
+    ignoreExportAllDeclaration: true,
+  })
+  return result.exportNames
+}
+
 describe(transformWrapExport, () => {
   test('basic', async () => {
     const input = `
@@ -52,6 +62,16 @@ export class Cls {};
       export { $$wrap_$$default as default };
       "
     `)
+
+    expect(await testTransformNames(input)).toMatchInlineSnapshot(`
+      [
+        "Arrow",
+        "default",
+        "Fn",
+        "AsyncFn",
+        "Cls",
+      ]
+    `)
   })
 
   test('preserve reference', async () => {

packages/plugin-rsc/src/transforms/wrap-export.ts

@@ -37,6 +37,7 @@ export function transformWrapExport(
     end: number,
     exports: { name: string; meta: ExportMeta }[],
   ) {
+    exportNames.push(...exports.map((e) => e.name))
     // update code and move to preserve `registerServerReference` position
     // e.g.
     // input
@@ -64,6 +65,7 @@ export function transformWrapExport(
   }
 
   function wrapExport(name: string, exportName: string, meta: ExportMeta = {}) {
+    exportNames.push(exportName)
     if (!filter(exportName, meta)) {
       toAppend.push(`export { ${name} as ${exportName} }`)
       return