refactor: add serverReferenceMetaMap

Author: hi-ogawa

packages/plugin-rsc/src/plugin.ts

@@ -82,7 +82,6 @@ class RscPluginManager {
   rscBundle!: Rollup.OutputBundle
   buildAssetsManifest: AssetsManifest | undefined
   isScanBuild: boolean = false
-  serverReferences: Record<string, string> = {}
   clientReferenceMetaMap: Record<string, ClientReferenceMeta> = {}
   serverReferenceMetaMap: Record<string, ServerRerferenceMeta> = {}
   serverResourcesMetaMap: Record<string, { key: string }> = {}
@@ -1436,7 +1435,7 @@ function vitePluginUseServer(
           )
           const enableEncryption =
             useServerPluginOptions.enableActionEncryption ?? true
-          const { output } = transformServerActionServer_(code, ast, {
+          const result = transformServerActionServer_(code, ast, {
             runtime: (value, name) =>
               `$$ReactServer.registerServerReference(${value}, ${JSON.stringify(
                 getNormalizedId(),
@@ -1451,8 +1450,14 @@ function vitePluginUseServer(
                   `await __vite_rsc_encryption_runtime.decryptActionBoundArgs(${value})`
               : undefined,
           })
+          if (!result) return
+          const output = result.output
           if (!output.hasChanged()) return
-          manager.serverReferences[getNormalizedId()] = id
+          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`)
           if (enableEncryption) {
@@ -1487,9 +1492,14 @@ function vitePluginUseServer(
             directive: 'use server',
             rejectNonAsyncFunction: true,
           })
+          if (!result) return
           const output = result?.output
           if (!output?.hasChanged()) return
-          manager.serverReferences[getNormalizedId()] = id
+          manager.serverReferenceMetaMap[id] = {
+            importId: id,
+            referenceKey: getNormalizedId(),
+            exportNames: result.exportNames,
+          }
           const name =
             this.environment.name === browserEnvironmentName ? 'browser' : 'ssr'
           const importSource = resolvePackage(`${PKG_NAME}/react/${name}`)
@@ -1505,7 +1515,19 @@ function vitePluginUseServer(
       if (this.environment.mode === 'dev') {
         return { code: `export {}`, map: null }
       }
-      const code = generateDynamicImportCode(manager.serverReferences)
+      let code = ''
+      for (const meta of Object.values(manager.serverReferenceMetaMap)) {
+        const key = JSON.stringify(meta.referenceKey)
+        const id = JSON.stringify(meta.importId)
+        const exports = meta.exportNames.join(',')
+        code += `
+          ${key}: async () => {
+            const {${exports}} = await import(${id});
+            return {${exports}};
+          },
+        `
+      }
+      code = `export default {${code}};\n`
       return { code, map: null }
     }),
   ]
@@ -1550,16 +1572,6 @@ function createVirtualPlugin(name: string, load: Plugin['load']) {
   } satisfies Plugin
 }
 
-function generateDynamicImportCode(map: Record<string, string>) {
-  let code = Object.entries(map)
-    .map(
-      ([key, id]) =>
-        `${JSON.stringify(key)}: () => import(${JSON.stringify(id)}),`,
-    )
-    .join('\n')
-  return `export default {${code}};\n`
-}
-
 class RuntimeAsset {
   runtime: string
   constructor(value: string) {