From cf9a04fb30d4a01158cd0c4f9551c7b499bb1776 Mon Sep 17 00:00:00 2001
From: Hiroshi Ogawa <hi.ogawa.zz@gmail.com>
Date: Fri, 13 Dec 2024 17:37:57 +0900
Subject: [PATCH] refactor: simplify entry ensureChunk

---
 packages/vite/misc/rolldown-runtime.js        | 38 +++++++----
 .../src/node/server/environments/rolldown.ts  | 68 +++++++++++++------
 2 files changed, 71 insertions(+), 35 deletions(-)

diff --git a/packages/vite/misc/rolldown-runtime.js b/packages/vite/misc/rolldown-runtime.js
index 2dd968ae9adf50..2e108ccbb8f756 100644
--- a/packages/vite/misc/rolldown-runtime.js
+++ b/packages/vite/misc/rolldown-runtime.js
@@ -252,22 +252,30 @@ self.__rolldown_runtime = {
       }
     }
   },
-  /** @type {{ chunks: Record<string, { fileName: string, imports: string[] }> }} */
+
+  /** @type {{ chunks: Record<string, { file: string, dependencies: string[] }> }} */
   manifest: {},
-  /**
-   * @param {string} chunkName
-   */
-  async ensureChunk(chunkName) {
-    await this.ensureChunkDeps(chunkName)
-    const file = this.manifest.chunks[chunkName].fileName
-    await import(`/${file}`)
+
+  /** @type {(name: string) => Promise<void>} */
+  async ensureChunk(name) {
+    const entry = this.manifest.chunks[name]
+    await Promise.all(
+      [name, ...entry.dependencies].map((name) => this.loadChunkCached(name)),
+    )
   },
-  /**
-   * @param {string} chunkName
-   */
-  async ensureChunkDeps(chunkName) {
-    for (const file of this.manifest.chunks[chunkName].imports) {
-      await import(`/${file}`)
-    }
+
+  /** @type {Record<string, Promise<void>>} */
+  loadChunkPromises: {},
+
+  /** @type {(name: string) => Promise<void>} */
+  async loadChunkCached(name) {
+    return (this.loadChunkPromises[name] ??= this.loadChunk(name))
+  },
+
+  /** @type {(name: string) => Promise<void>} */
+  async loadChunk(name) {
+    // TODO: use classic script
+    const file = this.manifest.chunks[name].file
+    await import(`/` + file)
   },
 }
diff --git a/packages/vite/src/node/server/environments/rolldown.ts b/packages/vite/src/node/server/environments/rolldown.ts
index 7b94d580a21476..f2a4f933a10317 100644
--- a/packages/vite/src/node/server/environments/rolldown.ts
+++ b/packages/vite/src/node/server/environments/rolldown.ts
@@ -295,7 +295,7 @@ class RolldownEnvironment extends DevEnvironment {
   }
 
   async buildHmr(file: string): Promise<{
-    manifest: ChunkManifest
+    manifest: BuildManifest
     chunk?: rolldown.RolldownOutputChunk
   }> {
     logger.info(`hmr '${file}'`, { timestamp: true })
@@ -468,21 +468,22 @@ self.__rolldown_runtime.manifest = ${JSON.stringify(manifest, null, 2)};
 self.__rolldown_runtime.patch(__rolldown_modules);
 `
           } else {
-            // TODO: avoid top-level-await?
             chunk.code += `
 Object.assign(self.__rolldown_runtime.moduleFactoryMap, __rolldown_modules);
-await self.__rolldown_runtime.ensureChunkDeps(${JSON.stringify(chunk.name)});
 `
-          }
-          if (chunk.isEntry) {
-            assert(chunk.facadeModuleId)
-            const stableId = path.relative(
-              environment.config.root,
-              chunk.facadeModuleId,
-            )
-            chunk.code += `
-self.__rolldown_runtime.require(${JSON.stringify(stableId)});
+            if (chunk.isEntry) {
+              assert(chunk.facadeModuleId)
+              const stableId = path.relative(
+                environment.config.root,
+                chunk.facadeModuleId,
+              )
+              chunk.code += `
+self.__rolldown_runtime.loadChunkPromises[${JSON.stringify(chunk.name)}] = Promise.resolve();
+self.__rolldown_runtime.ensureChunk(${JSON.stringify(chunk.name)}).then(function(){
+  self.__rolldown_runtime.require(${JSON.stringify(stableId)});
+});
 `
+            }
           }
           chunk.code = moveInlineSourcemapToEnd(chunk.code)
         }
@@ -491,20 +492,47 @@ self.__rolldown_runtime.require(${JSON.stringify(stableId)});
   }
 }
 
-export type ChunkManifest = {
-  chunks: Record<string, { fileName: string; imports: string[] }>
+export type BuildManifest = {
+  chunks: Record<string, { file: string; dependencies: string[] }>
 }
 
 function getChunkManifest(
   outputs: (rolldown.RolldownOutputChunk | rolldown.RolldownOutputAsset)[],
-): ChunkManifest {
-  const manifest: ChunkManifest = {
+): BuildManifest {
+  const chunks = outputs.filter((o) => o.type === 'chunk')
+  const fileToChunkName: Record<string, string> = {}
+  for (const chunk of chunks) {
+    fileToChunkName[chunk.fileName] = chunk.name
+  }
+
+  const directDepMap: Record<string, string[]> = {}
+  for (const chunk of chunks) {
+    directDepMap[chunk.name] = chunk.imports.map(
+      (file) => fileToChunkName[file],
+    )
+  }
+
+  function traverse(name: string, adj: Record<string, string[]>): string[] {
+    const visited = new Set<string>()
+    function recurse(name: string) {
+      if (!visited.has(name)) {
+        visited.add(name)
+        for (const dep of adj[name]) {
+          recurse(dep)
+        }
+      }
+    }
+    recurse(name)
+    return [...visited]
+  }
+
+  const manifest: BuildManifest = {
     chunks: {},
   }
-  for (const chunk of outputs) {
-    if (chunk.type === 'chunk') {
-      const { fileName, imports } = chunk
-      manifest.chunks[chunk.name] = { fileName, imports }
+  for (const chunk of chunks) {
+    manifest.chunks[chunk.name] = {
+      file: chunk.fileName,
+      dependencies: traverse(chunk.name, directDepMap),
     }
   }
   return manifest