fix(dev-server): excessive hmr invalidations

.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

packages/compiler/src/compiler.ts

@@ -2,7 +2,11 @@ import assert from 'assert';
 import { join, isAbsolute } from 'path';
 import type { Adapter } from '@vanilla-extract/css';
 import { transformCss } from '@vanilla-extract/css/transformCss';
-import type { ModuleNode, UserConfig as ViteUserConfig } from 'vite';
+import type {
+  ModuleNode,
+  ViteDevServer,
+  UserConfig as ViteUserConfig,
+} from 'vite';
 
 import {
   cssFileFilter,
@@ -94,10 +98,13 @@ const createViteServer = async ({
   identifiers,
   viteConfig,
   enableFileWatcher = true,
+  onFileChange,
 }: Required<
   Pick<CreateCompilerOptions, 'root' | 'identifiers' | 'viteConfig'>
 > &
-  Pick<CreateCompilerOptions, 'enableFileWatcher'>) => {
+  Pick<CreateCompilerOptions, 'enableFileWatcher'> & {
+    onFileChange?: (filePath: string) => void;
+  }) => {
   const pkg = getPackageInfo(root);
   const vite = await import('vite');
 
@@ -200,6 +207,7 @@ const createViteServer = async ({
   if (enableFileWatcher) {
     server.watcher.on('change', (filePath) => {
       runner.moduleCache.invalidateDepTree([filePath]);
+      onFileChange?.(filePath);
     });
   }
 
@@ -241,6 +249,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 {
@@ -279,6 +291,14 @@ export const createCompiler = ({
     'viteConfig cannot be used with viteResolve or vitePlugins',
   );
 
+  const processVanillaFileCache = new Map<
+    string,
+    {
+      lastInvalidationTimestamp: number;
+      result: ProcessedVanillaFile;
+    }
+  >();
+
   const vitePromise = createViteServer({
     root,
     identifiers,
@@ -287,16 +307,16 @@ export const createCompiler = ({
       plugins: vitePlugins,
     },
     enableFileWatcher,
+    onFileChange: (changedFilePath) => {
+      const normalizedPath = normalizePath(changedFilePath);
+      for (const [cacheKey, cached] of processVanillaFileCache.entries()) {
+        if (cached.result.watchFiles.has(normalizedPath)) {
+          processVanillaFileCache.delete(cacheKey);
+        }
+      }
+    },
   });
 
-  const processVanillaFileCache = new Map<
-    string,
-    {
-      lastInvalidationTimestamp: number;
-      result: ProcessedVanillaFile;
-    }
-  >();
-
   const cssCache = new NormalizedMap<{ css: string }>(root);
   const classRegistrationsByModuleId = new NormalizedMap<{
     localClassNames: Set<string>;
@@ -394,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);
@@ -511,5 +534,62 @@ 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, server);
+    },
   };
 };
+
+function _findImporterTree(
+  moduleNode: ModuleNode,
+  transformedModules: Set<string>,
+  server: ViteDevServer,
+  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,
+      server,
+      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,38 @@ 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>();
+
+    for (const mod of importerChain) {
+      if (mod.id && cssFileFilter.test(mod.id)) {
+        const virtualModules = moduleGraph.getModulesByFile(
+          fileIdToVirtualId(mod.id),
+        );
 
-        // Vite uses this timestamp to add `?t=` query string automatically for HMR.
-        module.lastHMRTimestamp =
-          module.lastInvalidationTimestamp || Date.now();
+        for (const virtualModule of virtualModules ?? []) {
+          moduleGraph.invalidateModule(virtualModule, seen, timestamp, true);
+        }
+      } else {
+        moduleGraph.invalidateModule(mod, seen, timestamp, true);
       }
     }
-  }
+  };
 
   return [
     {
@@ -138,6 +157,10 @@ export function vanillaExtractPlugin({
       name: PLUGIN_NAMESPACE,
       configureServer(_server) {
         server = _server;
+
+        server.watcher.on('unlink', (file) => {
+          transformedModules.delete(file);
+        });
       },
       config(_userConfig, _configEnv) {
         configEnv = _configEnv;
@@ -151,7 +174,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 +241,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 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 }) {
+        const importerChain = await compiler?.findImporterTree(
+          normalizePath(file),
+          transformedModules,
+        );
+
+        if (!importerChain || importerChain.size === 0) {
+          return;
+        }
+
+        invalidateImporterChain({
+          importerChain,
+          server,
+          timestamp,
+        });
+
+        // return empty array to tell Vite we've handled HMR ourselves,
+        // preventing Vite's default propagation which causes double invalidation.
+        return [];
       },
       resolveId(source) {
         const [validId, query] = source.split('?');