From b8f816188f9d473fa30a9b2d6271781b00717453 Mon Sep 17 00:00:00 2001
From: Alexandre Stahmer <alexandre.stahmer@gmail.com>
Date: Wed, 7 Dec 2022 11:41:57 +0100
Subject: [PATCH] fix: hmr (from #925) + SSR (rootRelativeId) adapted from:
 https://github.com/vanilla-extract-css/vanilla-extract/pull/925/files

---
 packages/vite-plugin/src/index.ts | 48 +++++++++++++++++++++++++------
 1 file changed, 40 insertions(+), 8 deletions(-)

diff --git a/packages/vite-plugin/src/index.ts b/packages/vite-plugin/src/index.ts
index 6183efa01..f5d3746c7 100644
--- a/packages/vite-plugin/src/index.ts
+++ b/packages/vite-plugin/src/index.ts
@@ -20,6 +20,7 @@ const styleUpdateEvent = (fileId: string) =>
 
 const virtualExtCss = '.vanilla.css';
 const virtualExtJs = '.vanilla.js';
+const virtualRE = /.vanilla.(css|js)$/;
 
 export interface VanillaExtractPluginOptions
   extends Pick<
@@ -88,6 +89,30 @@ export function vanillaExtractPlugin({
         forceEmitCssInSsrBuild = true;
       }
     },
+    // Re-parse .css.ts files when they change
+    async handleHotUpdate({ file, modules }) {
+      if (!cssFileFilter.test(file)) return;
+      try {
+        const virtuals: any[] = [];
+        const invalidate = (type: string) => {
+          const found = server.moduleGraph.getModulesByFile(`${file}${type}`);
+          found?.forEach((m) => {
+            virtuals.push(m);
+            return server.moduleGraph.invalidateModule(m);
+          });
+        };
+        invalidate(virtualExtCss);
+        invalidate(virtualExtJs);
+        // load new CSS
+        await server.ssrLoadModule(file);
+        return [...modules, ...virtuals];
+      } catch (e) {
+        // eslint-disable-next-line no-console
+        console.error(e);
+        throw e;
+      }
+    },
+    // Convert .vanilla.(js|css) URLs to their absolute version
     resolveId(source) {
       const [validId, query] = source.split('?');
       if (!validId.endsWith(virtualExtCss) && !validId.endsWith(virtualExtJs)) {
@@ -103,18 +128,25 @@ export function vanillaExtractPlugin({
       // There should always be an entry in the `cssMap` here.
       // The only valid scenario for a missing one is if someone had written
       // a file in their app using the .vanilla.js/.vanilla.css extension
-      if (cssMap.has(absoluteId)) {
-        // Keep the original query string for HMR.
-        return absoluteId + (query ? `?${query}` : '');
-      }
+      // Keep the original query string for HMR.
+      return absoluteId + (query ? `?${query}` : '');
     },
-    load(id) {
+    // Provide virtual CSS content
+    async load(id) {
       const [validId] = id.split('?');
 
-      if (!cssMap.has(validId)) {
+      if (!virtualRE.test(validId)) {
         return;
       }
 
+      if (!cssMap.has(validId)) {
+        // Try to parse the parent
+        const parentId = validId.replace(virtualRE, '');
+        await server.ssrLoadModule(parentId);
+        // Now we should have the CSS
+        if (!cssMap.has(validId)) return;
+      }
+
       const css = cssMap.get(validId);
 
       if (typeof css !== 'string') {
@@ -193,7 +225,7 @@ export function vanillaExtractPlugin({
         serializeVanillaModule,
         serializeVirtualCssPath: async ({ fileScope, source }) => {
           const rootRelativeId = `${fileScope.filePath}${
-            config.command === 'build' || (ssr && forceEmitCssInSsrBuild)
+            config.command === 'build' || forceEmitCssInSsrBuild
               ? virtualExtCss
               : virtualExtJs
           }`;
@@ -216,7 +248,7 @@ export function vanillaExtractPlugin({
           if (
             server &&
             cssMap.has(absoluteId) &&
-            cssMap.get(absoluteId) !== source
+            cssMap.get(absoluteId) !== cssSource
           ) {
             const { moduleGraph } = server;
             const [module] = Array.from(