fix(vite): refactor module runner to module graph

packages/vite/src/module-runner/hmrHandler.ts

@@ -54,7 +54,7 @@ export async function handleHotPayload(
 
       hmrClient.logger.debug(`program reload`)
       await hmrClient.notifyListeners('vite:beforeFullReload', payload)
-      runner.moduleCache.clear()
+      runner.moduleGraph.clear()
 
       for (const id of clearEntrypoints) {
         await runner.import(id)
@@ -122,9 +122,9 @@ class Queue {
 
 function getModulesByFile(runner: ModuleRunner, file: string) {
   const modules: string[] = []
-  for (const [id, mod] of runner.moduleCache.entries()) {
+  for (const [_, mod] of runner.moduleGraph.idToModuleMap.entries()) {
     if (mod.meta && 'file' in mod.meta && mod.meta.file === file) {
-      modules.push(id)
+      modules.push(mod.url)
     }
   }
   return modules
@@ -139,9 +139,12 @@ function getModulesEntrypoints(
   for (const moduleId of modules) {
     if (visited.has(moduleId)) continue
     visited.add(moduleId)
-    const module = runner.moduleCache.getByModuleId(moduleId)
+    const module = runner.moduleGraph.getModuleById(moduleId)
+    if (!module) {
+      continue
+    }
     if (module.importers && !module.importers.size) {
-      entrypoints.add(moduleId)
+      entrypoints.add(module.url)
       continue
     }
     for (const importer of module.importers || []) {
@@ -155,9 +158,9 @@ function findAllEntrypoints(
   runner: ModuleRunner,
   entrypoints = new Set<string>(),
 ): Set<string> {
-  for (const [id, mod] of runner.moduleCache.entries()) {
+  for (const [_, mod] of runner.moduleGraph.idToModuleMap.entries()) {
     if (mod.importers && !mod.importers.size) {
-      entrypoints.add(id)
+      entrypoints.add(mod.url)
     }
   }
   return entrypoints

packages/vite/src/module-runner/index.ts

@@ -1,6 +1,6 @@
 // this file should re-export only things that don't rely on Node.js or other runner features
 
-export { ModuleCacheMap } from './moduleCache'
+export { ModuleRunnerGraph, type ModuleRunnerNode } from './moduleCache'
 export { ModuleRunner } from './runner'
 export { ESModulesEvaluator } from './esmEvaluator'
 export { RemoteRunnerTransport } from './runnerTransport'
@@ -10,7 +10,6 @@ export type { HMRLogger, HMRConnection } from '../shared/hmr'
 export type {
   ModuleEvaluator,
   ModuleRunnerContext,
-  ModuleCache,
   FetchResult,
   FetchFunction,
   FetchFunctionOptions,

packages/vite/src/module-runner/moduleCache.ts

@@ -1,139 +1,110 @@
-import { isWindows, slash, withTrailingSlash } from '../shared/utils'
+import {
+  cleanUrl,
+  isWindows,
+  slash,
+  unwrapId,
+  withTrailingSlash,
+} from '../shared/utils'
 import { SOURCEMAPPING_URL } from '../shared/constants'
-import { decodeBase64 } from './utils'
+import { decodeBase64, posixResolve } from './utils'
 import { DecodedMap } from './sourcemap/decoder'
-import type { ModuleCache } from './types'
+import type { ResolvedResult } from './types'
 
 const MODULE_RUNNER_SOURCEMAPPING_REGEXP = new RegExp(
   `//# ${SOURCEMAPPING_URL}=data:application/json;base64,(.+)`,
 )
 
-export class ModuleCacheMap extends Map<string, ModuleCache> {
-  private root: string
-
-  constructor(root: string, entries?: [string, ModuleCache][]) {
-    super(entries)
-    this.root = withTrailingSlash(root)
-  }
-
-  normalize(fsPath: string): string {
-    return normalizeModuleId(fsPath, this.root)
+export class ModuleRunnerNode {
+  public importers = new Set<string>()
+  public imports = new Set<string>()
+  public lastInvalidationTimestamp = 0
+  public evaluated = false
+  public meta: ResolvedResult | undefined
+  public promise: Promise<any> | undefined
+  public exports: any | undefined
+  public file: string
+  public map: DecodedMap | undefined
+
+  constructor(
+    public id: string,
+    public url: string,
+  ) {
+    this.file = cleanUrl(id)
   }
+}
 
-  /**
-   * Assign partial data to the map
-   */
-  update(fsPath: string, mod: ModuleCache): this {
-    fsPath = this.normalize(fsPath)
-    if (!super.has(fsPath)) this.setByModuleId(fsPath, mod)
-    else Object.assign(super.get(fsPath)!, mod)
-    return this
-  }
+export class ModuleRunnerGraph {
+  private root: string
 
-  setByModuleId(modulePath: string, mod: ModuleCache): this {
-    return super.set(modulePath, mod)
-  }
+  public idToModuleMap = new Map<string, ModuleRunnerNode>()
+  public fileToModuleMap = new Map<string, ModuleRunnerNode[]>()
 
-  override set(fsPath: string, mod: ModuleCache): this {
-    return this.setByModuleId(this.normalize(fsPath), mod)
+  constructor(root: string) {
+    this.root = withTrailingSlash(root)
   }
 
-  getByModuleId(modulePath: string): ModuleCache {
-    if (!super.has(modulePath)) this.setByModuleId(modulePath, {})
-
-    const mod = super.get(modulePath)!
-    if (!mod.imports) {
-      Object.assign(mod, {
-        imports: new Set(),
-        importers: new Set(),
-        timestamp: 0,
-      })
-    }
-    return mod
+  public getModuleById(id: string): ModuleRunnerNode | undefined {
+    return this.idToModuleMap.get(id)
   }
 
-  override get(fsPath: string): ModuleCache {
-    return this.getByModuleId(this.normalize(fsPath))
+  public getModulesByFile(file: string): ModuleRunnerNode[] {
+    return this.fileToModuleMap.get(file) || []
   }
 
-  deleteByModuleId(modulePath: string): boolean {
-    return super.delete(modulePath)
+  public getModuleByUrl(url: string): ModuleRunnerNode | undefined {
+    url = unwrapId(url)
+    if (url.startsWith('/')) {
+      const id = posixResolve(this.root, url.slice(1))
+      return this.idToModuleMap.get(id)
+    }
+    return this.idToModuleMap.get(url)
   }
 
-  override delete(fsPath: string): boolean {
-    return this.deleteByModuleId(this.normalize(fsPath))
-  }
+  public ensureModule(id: string, url: string): ModuleRunnerNode {
+    id = normalizeModuleId(id)
+    if (this.idToModuleMap.has(id)) {
+      return this.idToModuleMap.get(id)!
+    }
+    const moduleNode = new ModuleRunnerNode(id, url)
+    this.idToModuleMap.set(id, moduleNode)
 
-  invalidateUrl(id: string): void {
-    const module = this.get(id)
-    this.invalidateModule(module)
+    const fileModules = this.fileToModuleMap.get(moduleNode.file) || []
+    fileModules.push(moduleNode)
+    this.fileToModuleMap.set(moduleNode.file, fileModules)
+    return moduleNode
   }
 
-  invalidateModule(module: ModuleCache): void {
-    module.evaluated = false
-    module.meta = undefined
-    module.map = undefined
-    module.promise = undefined
-    module.exports = undefined
+  public invalidateModule(node: ModuleRunnerNode): void {
+    node.evaluated = false
+    node.meta = undefined
+    node.map = undefined
+    node.promise = undefined
+    node.exports = undefined
     // remove imports in case they are changed,
     // don't remove the importers because otherwise it will be empty after evaluation
     // this can create a bug when file was removed but it still triggers full-reload
     // we are fine with the bug for now because it's not a common case
-    module.imports?.clear()
-  }
-
-  /**
-   * Invalidate modules that dependent on the given modules, up to the main entry
-   */
-  invalidateDepTree(
-    ids: string[] | Set<string>,
-    invalidated = new Set<string>(),
-  ): Set<string> {
-    for (const _id of ids) {
-      const id = this.normalize(_id)
-      if (invalidated.has(id)) continue
-      invalidated.add(id)
-      const mod = super.get(id)
-      if (mod?.importers) this.invalidateDepTree(mod.importers, invalidated)
-      this.invalidateUrl(id)
-    }
-    return invalidated
-  }
-
-  /**
-   * Invalidate dependency modules of the given modules, down to the bottom-level dependencies
-   */
-  invalidateSubDepTree(
-    ids: string[] | Set<string>,
-    invalidated = new Set<string>(),
-  ): Set<string> {
-    for (const _id of ids) {
-      const id = this.normalize(_id)
-      if (invalidated.has(id)) continue
-      invalidated.add(id)
-      const subIds = Array.from(super.entries())
-        .filter(([, mod]) => mod.importers?.has(id))
-        .map(([key]) => key)
-      if (subIds.length) {
-        this.invalidateSubDepTree(subIds, invalidated)
-      }
-      super.delete(id)
-    }
-    return invalidated
+    node.imports?.clear()
   }
 
-  getSourceMap(moduleId: string): null | DecodedMap {
-    const mod = this.get(moduleId)
+  getModuleSourceMapById(id: string): null | DecodedMap {
+    const mod = this.getModuleById(id)
+    if (!mod) return null
     if (mod.map) return mod.map
     if (!mod.meta || !('code' in mod.meta)) return null
     const mapString = MODULE_RUNNER_SOURCEMAPPING_REGEXP.exec(
       mod.meta.code,
     )?.[1]
     if (!mapString) return null
-    const baseFile = mod.meta.file || moduleId.split('?')[0]
+    const baseFile = mod.file
     mod.map = new DecodedMap(JSON.parse(decodeBase64(mapString)), baseFile)
     return mod.map
   }
+
+  public clear(): void {
+    this.idToModuleMap.clear()
+    this.fileToModuleMap.clear()
+  }
 }
 
 // unique id that is not available as "$bare_import" like "test"
@@ -146,20 +117,15 @@ const prefixedBuiltins = new Set(['node:test', 'node:sqlite'])
 // /root/id.js -> /id.js
 // C:/root/id.js -> /id.js
 // C:\root\id.js -> /id.js
-function normalizeModuleId(file: string, root: string): string {
+function normalizeModuleId(file: string): string {
   if (prefixedBuiltins.has(file)) return file
 
   // unix style, but Windows path still starts with the drive letter to check the root
-  let unixFile = slash(file)
+  const unixFile = slash(file)
     .replace(/^\/@fs\//, isWindows ? '' : '/')
     .replace(/^node:/, '')
     .replace(/^\/+/, '/')
 
-  if (unixFile.startsWith(root)) {
-    // keep slash
-    unixFile = unixFile.slice(root.length - 1)
-  }
-
   // if it's not in the root, keep it as a path, not a URL
   return unixFile.replace(/^file:\//, '/')
 }