vite: Reduce unnecessary HMR invalidations (#1701)

Author: askoufis

.changeset/new-radios-swim.md

@@ -0,0 +1,5 @@
+---
+'@vanilla-extract/compiler': minor
+---
+
+Expose `compiler.findImporterTree` API

.changeset/weak-apricots-know.md

@@ -0,0 +1,5 @@
+---
+'@vanilla-extract/vite-plugin': patch
+---
+
+Reduce unnecessary HMR invalidations

fixtures/low-level/src/index.ts

@@ -8,3 +8,9 @@ document.body.innerHTML = `
   </div>
 </div>
 `;
+
+// @ts-expect-error Vite env not defined
+if (import.meta.hot) {
+  // @ts-expect-error Vite env not defined
+  import.meta.hot.accept();
+}

packages/compiler/src/compiler.ts

@@ -246,6 +246,10 @@ export interface Compiler {
   getCssForFile(virtualCssFilePath: string): { filePath: string; css: string };
   close(): Promise<void>;
   getAllCss(): string;
+  findImporterTree(
+    filePath: string,
+    transformedVeModules: Set<string>,
+  ): Promise<Set<ModuleNode>>;
 }
 
 interface ProcessedVanillaFile {
@@ -410,7 +414,10 @@ export const createCompiler = ({
         await lock(async () => {
           runner.cssAdapter = cssAdapter;
 
-          const fileExports = await runner.executeFile(filePath);
+          const fileExports = (await runner.executeFile(filePath)) as Record<
+            string,
+            unknown
+          >;
 
           const moduleId = normalizePath(filePath);
           const moduleNode = server.moduleGraph.getModuleById(moduleId);
@@ -558,5 +565,56 @@ export const createCompiler = ({
 
       return allCss;
     },
+    /**
+     * Returns an importer tree based off the compiler's module graph. We can't use the
+     * consuming Vite dev server's module graph as it ends up modified by the `transform` hook to a
+     * point where we can't reconstruct the original importer chain.
+     */
+    async findImporterTree(filePath, transformedModules) {
+      const { server } = await vitePromise;
+
+      // The compiler's module graph is always a subset of the consuming Vite dev server's module
+      // graph, so this early exit will be hit for any modules that aren't involved in compiling VE
+      // modules
+      const moduleNode = server.moduleGraph.getModuleById(
+        normalizePath(filePath),
+      );
+      if (!moduleNode) {
+        return new Set();
+      }
+
+      return _findImporterTree(moduleNode, transformedModules);
+    },
   };
 };
+
+function _findImporterTree(
+  moduleNode: ModuleNode,
+  transformedModules: Set<string>,
+  visited = new Set<string>(),
+): Set<ModuleNode> {
+  const result = new Set<ModuleNode>();
+  if (!moduleNode.id || visited.has(moduleNode.id)) {
+    return result;
+  }
+
+  // Include the starting module in the tree
+  result.add(moduleNode);
+  visited.add(moduleNode.id);
+
+  // Stop if we hit a transformed module as this is a VE module boundary that we don't
+  // need to invalidate past
+  if (transformedModules.has(moduleNode.id)) {
+    return result;
+  }
+
+  for (const importer of moduleNode.importers) {
+    const chain = _findImporterTree(importer, transformedModules, visited);
+
+    for (const mod of chain) {
+      result.add(mod);
+    }
+  }
+
+  return result;
+}

packages/vite-plugin/src/index.ts

@@ -4,10 +4,11 @@ import type {
   Plugin,
   ResolvedConfig,
   ConfigEnv,
-  ViteDevServer,
   PluginOption,
   TransformResult,
   UserConfig,
+  ModuleNode,
+  ViteDevServer,
 } from 'vite';
 import { type Compiler, createCompiler } from '@vanilla-extract/compiler';
 import {
@@ -70,6 +71,8 @@ export function vanillaExtractPlugin({
   let isBuild: boolean;
   const vitePromise = import('vite');
 
+  const transformedModules = new Set<string>();
+
   const getIdentOption = () =>
     identifiers ?? (config.mode === 'production' ? 'short' : 'debug');
   const getAbsoluteId = (filePath: string) => {
@@ -92,22 +95,43 @@ export function vanillaExtractPlugin({
     return normalizePath(resolvedId);
   };
 
-  function invalidateModule(absoluteId: string) {
-    if (!server) return;
-
+  /**
+   * Custom invalidation function that takes a chain of importers to invalidate. If an importer is a
+   * VE module, its virtual CSS is invalidated instead. Otherwise, the module is invalidated
+   * normally.
+   */
+  const invalidateImporterChain = ({
+    importerChain,
+    server,
+    timestamp,
+  }: {
+    importerChain: Set<ModuleNode>;
+    server: ViteDevServer;
+    timestamp: number;
+  }) => {
     const { moduleGraph } = server;
-    const modules = moduleGraph.getModulesByFile(absoluteId);
 
-    if (modules) {
-      for (const module of modules) {
-        moduleGraph.invalidateModule(module);
+    const seen = new Set<ModuleNode>();
 
-        // Vite uses this timestamp to add `?t=` query string automatically for HMR.
-        module.lastHMRTimestamp =
-          module.lastInvalidationTimestamp || Date.now();
+    for (const mod of importerChain) {
+      if (mod.id && cssFileFilter.test(mod.id)) {
+        const virtualModules = moduleGraph.getModulesByFile(
+          fileIdToVirtualId(mod.id),
+        );
+
+        for (const virtualModule of virtualModules ?? []) {
+          moduleGraph.invalidateModule(virtualModule, seen, timestamp, true);
+        }
+      } else if (mod.id) {
+        // `mod` is from the compiler's internal Vite server, so look up the
+        // corresponding module in the consuming server's graph by ID
+        const serverMod = moduleGraph.getModuleById(mod.id);
+        if (serverMod) {
+          moduleGraph.invalidateModule(serverMod, seen, timestamp, true);
+        }
       }
     }
-  }
+  };
 
   return [
     {
@@ -138,6 +162,10 @@ export function vanillaExtractPlugin({
       name: PLUGIN_NAMESPACE,
       configureServer(_server) {
         server = _server;
+
+        server.watcher.on('unlink', (file) => {
+          transformedModules.delete(normalizePath(file));
+        });
       },
       config(_userConfig, _configEnv) {
         configEnv = _configEnv;
@@ -151,7 +179,7 @@ export function vanillaExtractPlugin({
           },
         };
       },
-      async configResolved(_resolvedConfig) {
+      configResolved(_resolvedConfig) {
         config = _resolvedConfig;
         isBuild = config.command === 'build' && !config.build.watch;
         packageName = getPackageInfo(config.root).name;
@@ -218,52 +246,75 @@ export function vanillaExtractPlugin({
         }
 
         const identOption = getIdentOption();
+        const normalizedId = normalizePath(validId);
 
         if (unstable_mode === 'transform') {
+          transformedModules.add(normalizedId);
+
           return transform({
             source: code,
-            filePath: normalizePath(validId),
+            filePath: normalizedId,
             rootPath: config.root,
             packageName,
             identOption,
           });
         }
 
-        if (compiler) {
-          const absoluteId = getAbsoluteId(validId);
+        if (!compiler) {
+          return null;
+        }
 
-          const { source, watchFiles } = await compiler.processVanillaFile(
-            absoluteId,
-            { outputCss: true },
-          );
-          const result: TransformResult = {
-            code: source,
-            map: { mappings: '' },
-          };
+        const absoluteId = getAbsoluteId(validId);
 
-          // We don't need to watch files or invalidate modules in build mode or during SSR
-          if (isBuild || options?.ssr) {
-            return result;
-          }
+        const { source, watchFiles } = await compiler.processVanillaFile(
+          absoluteId,
+          { outputCss: true },
+        );
 
-          for (const file of watchFiles) {
-            if (
-              !file.includes('node_modules') &&
-              normalizePath(file) !== absoluteId
-            ) {
-              this.addWatchFile(file);
-            }
-
-            // We have to invalidate the virtual module & deps, not the real one we just transformed
-            // The deps have to be invalidated in case one of them changing was the trigger causing
-            // the current transformation
-            if (cssFileFilter.test(file)) {
-              invalidateModule(fileIdToVirtualId(file));
-            }
-          }
+        transformedModules.add(normalizedId);
+
+        const result: TransformResult = {
+          code: source,
+          map: { mappings: '' },
+        };
 
+        // We don't need to watch files or invalidate modules in build mode or during SSR
+        if (isBuild || options?.ssr) {
           return result;
         }
+
+        for (const file of watchFiles) {
+          if (
+            !file.includes('node_modules') &&
+            normalizePath(file) !== absoluteId
+          ) {
+            this.addWatchFile(file);
+          }
+        }
+
+        return result;
+      },
+      // The compiler's module graph is always a subset of the consuming Vite dev server's module
+      // graph, so this early exit will be hit for any modules that aren't related to VE modules.
+      async handleHotUpdate({ file, server, timestamp }) {
+        if (!compiler) {
+          return;
+        }
+
+        const importerChain = await compiler.findImporterTree(
+          normalizePath(file),
+          transformedModules,
+        );
+
+        if (importerChain.size === 0) {
+          return;
+        }
+
+        invalidateImporterChain({
+          importerChain,
+          server,
+          timestamp,
+        });
       },
       resolveId(source) {
         const [validId, query] = source.split('?');