wip: dispose no longer imported modules

packages/vite/src/client/client.ts

@@ -119,6 +119,18 @@ async function handleMessage(payload: HMRPayload) {
         location.reload()
       }
       break
+    case 'dispose':
+      // After an HMR update, some modules are no longer imported on the page
+      // but they may have left behind side effects that need to be cleaned up
+      // (.e.g style injections)
+      // TODO Trigger their dispose callbacks.
+      payload.paths.forEach((path) => {
+        const fn = disposeMap.get(path)
+        if (fn) {
+          fn(dataMap.get(path))
+        }
+      })
+      break
     case 'error':
       const err = payload.err
       if (enableOverlay) {
@@ -230,7 +242,7 @@ export function updateStyle(id: string, content: string) {
   sheetsMap.set(id, style)
 }
 
-function removeStyle(id: string) {
+export function removeStyle(id: string) {
   let style = sheetsMap.get(id)
   if (style) {
     if (style instanceof CSSStyleSheet) {

packages/vite/src/node/plugins/css.ts

@@ -74,11 +74,15 @@ export function cssPlugin(config: ResolvedConfig, isBuild: boolean): Plugin {
         if (isProxyRequest) {
           debug(`[import] ${chalk.dim(path.relative(config.root, id))}`)
           return [
-            `import { updateStyle } from ${JSON.stringify(CLIENT_PUBLIC_PATH)}`,
+            `import { updateStyle, removeStyle } from ${JSON.stringify(
+              CLIENT_PUBLIC_PATH
+            )}`,
+            `const id = ${JSON.stringify(id)}`,
             `const css = ${JSON.stringify(css)}`,
-            `updateStyle(${JSON.stringify(id)}, css)`,
+            `updateStyle(id, css)`,
             `${modulesCode || `export default css`}`,
-            `import.meta.hot.accept()`
+            `import.meta.hot.accept()`,
+            `import.meta.hot.dispose(() => removeStyle(id))`
           ].join('\n')
         } else {
           debug(`[link] ${chalk.dim(path.relative(config.root, id))}`)

packages/vite/src/node/plugins/rewriteImports.ts → packages/vite/src/node/plugins/imports.ts

@@ -6,20 +6,23 @@ import { init, parse, ImportSpecifier } from 'es-module-lexer'
 import { isCSSRequest } from './css'
 import slash from 'slash'
 import { createDebugger, prettifyUrl, timeFrom } from '../utils'
-import { debugHmr } from '../server/hmr'
+import { debugHmr, handleDisposedModules } from '../server/hmr'
 import { FILE_PREFIX, CLIENT_PUBLIC_PATH } from '../constants'
 import { RollupError } from 'rollup'
 import { FAILED_RESOLVE } from './resolve'
 
 const isDebug = !!process.env.DEBUG
-const debugRewrite = createDebugger('vite:rewrite')
+const debugRewrite = createDebugger('vite:imports')
 
 const skipRE = /\.(map|json)$/
 const canSkip = (id: string) => skipRE.test(id) || isCSSRequest(id)
 
 /**
- * Server-only plugin that rewrites url imports (bare modules, css/asset imports)
- * so that they can be properly handled by the server.
+ * Server-only plugin that lexes, resolves, rewrites and analyzes url imports.
+ *
+ * - Imports are resolved to ensure they exist on disk
+ *
+ * - Lexes HMR accept calls and updates import relationships in the module graph
  *
  * - Bare module imports are resolved (by @rollup-plugin/node-resolve) to
  * absolute file paths, e.g.
@@ -43,7 +46,7 @@ const canSkip = (id: string) => skipRE.test(id) || isCSSRequest(id)
  *     import './style.css.js'
  *     ```
  */
-export function rewritePlugin(config: ResolvedConfig): Plugin {
+export function importAnalysisPlugin(config: ResolvedConfig): Plugin {
   return {
     name: 'vite:imports',
     async transform(source, importer) {
@@ -89,11 +92,6 @@ export function rewritePlugin(config: ResolvedConfig): Plugin {
       // since we are already in the transform phase of the importer, it must
       // have been loaded so its entry is guaranteed in the module graph.
       const importerModule = moduleGraph.getModuleById(importer)!
-
-      if (!importerModule) {
-        debugger
-      }
-
       const importedUrls = new Set<string>()
       const acceptedUrls = new Set<string>()
       const toAbsoluteUrl = (url: string) =>
@@ -227,13 +225,17 @@ export function rewritePlugin(config: ResolvedConfig): Plugin {
       }
 
       // update the module graph for HMR analysis
-      await moduleGraph.updateModuleInfo(
+      const noLongerImportedDeps = await moduleGraph.updateModuleInfo(
         importerModule,
         importedUrls,
         new Set([...acceptedUrls].map(toAbsoluteUrl)),
         isSelfAccepting
       )
 
+      if (hasHMR && noLongerImportedDeps) {
+        handleDisposedModules(noLongerImportedDeps, (this as any).server)
+      }
+
       isDebug &&
         debugRewrite(
           `${timeFrom(rewriteStart, timeSpentResolving)} ${prettyImporter}`

packages/vite/src/node/plugins/index.ts

@@ -4,7 +4,7 @@ import jsonPlugin from '@rollup/plugin-json'
 import { nodeResolve } from '@rollup/plugin-node-resolve'
 import { resolvePlugin, supportedExts } from './resolve'
 import { esbuildPlugin } from './esbuild'
-import { rewritePlugin } from './rewriteImports'
+import { importAnalysisPlugin } from './imports'
 import { cssPlugin } from './css'
 import { assetPlugin } from './asset'
 import { clientInjectionsPlugin } from './clientInjections'
@@ -33,6 +33,8 @@ export function resolvePlugins(
     assetPlugin(config, isBuild),
     ...postPlugins,
     // internal server-only plugins are always applied after everything else
-    ...(isBuild ? [] : [clientInjectionsPlugin(config), rewritePlugin(config)])
+    ...(isBuild
+      ? []
+      : [clientInjectionsPlugin(config), importAnalysisPlugin(config)])
   ].filter(Boolean) as Plugin[]
 }

packages/vite/src/node/server/hmr.ts

@@ -18,10 +18,13 @@ export interface HmrOptions {
   overlay?: boolean
 }
 
-export function handleHMRUpdate(file: string, context: ViteDevServer): any {
+export function handleHMRUpdate(
+  file: string,
+  { ws, config, moduleGraph }: ViteDevServer
+): any {
   debugHmr(`[file change] ${chalk.dim(file)}`)
 
-  if (file === context.config.configPath) {
+  if (file === config.configPath) {
     // TODO auto restart server
     return
   }
@@ -33,14 +36,14 @@ export function handleHMRUpdate(file: string, context: ViteDevServer): any {
 
   // html files and the client itself cannot be hot updated.
   if (file.endsWith('.html') || file.startsWith(CLIENT_DIR)) {
-    context.ws.send({
+    ws.send({
       type: 'full-reload',
-      path: '/' + slash(path.relative(context.config.root, file))
+      path: '/' + slash(path.relative(config.root, file))
     })
     return
   }
 
-  const mods = context.moduleGraph.getModulesByFile(file)
+  const mods = moduleGraph.getModulesByFile(file)
   if (!mods) {
     // loaded but not in the module graph, probably not js
     debugHmr(`[no module entry] ${chalk.dim(file)}`)
@@ -55,7 +58,7 @@ export function handleHMRUpdate(file: string, context: ViteDevServer): any {
     const hasDeadEnd = propagateUpdate(mod, timestamp, boundaries)
     if (hasDeadEnd) {
       debugHmr(`[full reload] ${chalk.dim(file)}`)
-      context.ws.send({
+      ws.send({
         type: 'full-reload'
       })
       return
@@ -75,19 +78,35 @@ export function handleHMRUpdate(file: string, context: ViteDevServer): any {
     )
   }
 
-  context.ws.send({
+  ws.send({
     type: 'update',
     updates
   })
 }
 
+export function handleDisposedModules(
+  mods: Set<ModuleNode>,
+  { ws }: ViteDevServer
+) {
+  // update the disposed modules' hmr timestamp
+  // since if it's re-imported, it should re-apply side effects
+  // and without the timestamp the browser will not re-import it!
+  const t = Date.now()
+  mods.forEach((mod) => {
+    mod.lastHMRTimestamp = t
+  })
+  ws.send({
+    type: 'dispose',
+    paths: [...mods].map((m) => m.url)
+  })
+}
+
 function propagateUpdate(
   node: ModuleNode,
   timestamp: number,
   boundaries: Set<ModuleNode>,
   currentChain: ModuleNode[] = [node]
 ): boolean /* hasDeadEnd */ {
-  debugger
   if (node.isSelfAccepting) {
     boundaries.add(node)
     // mark current propagation chain dirty.

packages/vite/src/node/server/moduleGraph.ts

@@ -17,6 +17,7 @@ export class ModuleNode {
   file: string | null = null
   type: 'js' | 'css'
   importers = new Set<ModuleNode>()
+  importedModules = new Set<ModuleNode>()
   acceptedHmrDeps = new Set<ModuleNode>()
   isSelfAccepting = false
   transformResult: TransformResult | null = null
@@ -61,33 +62,43 @@ export class ModuleGraph {
     }
   }
 
+  /**
+   * Update the module graph based on a module's updated imports information
+   * If there are dependencies that no longer have any importers, they are
+   * returned as a Set.
+   */
   async updateModuleInfo(
     mod: ModuleNode,
     importedUrls: Set<string>,
     acceptedUrls: Set<string>,
     isSelfAccepting: boolean
-  ) {
+  ): Promise<Set<ModuleNode> | undefined> {
     mod.isSelfAccepting = isSelfAccepting
-    const prevDeps = mod.acceptedHmrDeps
-    const newDeps = (mod.acceptedHmrDeps = new Set())
+    const prevImports = mod.importedModules
+    const nextImports = (mod.importedModules = new Set())
+    let noLongerImported: Set<ModuleNode> | undefined
+    // update import graph
     for (const depUrl of importedUrls) {
       const dep = await this.ensureEntry(depUrl)
       dep.importers.add(mod)
-    }
-    for (const depUrl of acceptedUrls) {
-      const dep = await this.ensureEntry(depUrl)
-      dep.importers.add(mod)
-      newDeps.add(dep)
+      nextImports.add(dep)
     }
     // remove the importer from deps that were imported but no longer are.
-    prevDeps.forEach((dep) => {
-      if (!newDeps.has(dep)) {
+    prevImports.forEach((dep) => {
+      if (!nextImports.has(dep)) {
         dep.importers.delete(mod)
         if (!dep.importers.size) {
-          console.log(`module no longer has importers: ${mod.url}`)
+          // dependency no longer imported
+          ;(noLongerImported || (noLongerImported = new Set())).add(dep)
         }
       }
     })
+    // update accepted hmr deps
+    const newDeps = (mod.acceptedHmrDeps = new Set())
+    for (const depUrl of acceptedUrls) {
+      newDeps.add(await this.ensureEntry(depUrl))
+    }
+    return noLongerImported
   }
 
   async ensureEntry(rawUrl: string) {

packages/vite/types/hmrPayload.d.ts

@@ -5,6 +5,7 @@ export type HMRPayload =
   | StyleRemovePayload
   | CustomPayload
   | ErrorPayload
+  | DisposePayload
 
 export interface ConnectedPayload {
   type: 'connected'
@@ -22,6 +23,11 @@ export interface Update {
   timestamp: number
 }
 
+export interface DisposePayload {
+  type: 'dispose'
+  paths: string[]
+}
+
 export interface StyleRemovePayload {
   type: 'css-remove'
   path: string