From 3e8844fa871fa477026375db6d921beb4b23b0dc Mon Sep 17 00:00:00 2001
From: Matthew Phillips <matthew@matthewphillips.info>
Date: Wed, 2 Feb 2022 11:35:13 -0500
Subject: [PATCH] Fix support for scss in static build (#2522)

* Fix support for scss in static build

* Adds a changeset

* Pass the normalizedID to transformWithVite
---
 .changeset/brown-dancers-perform.md           |  5 ++
 examples/fast-build/src/pages/index.astro     |  4 +-
 examples/fast-build/src/styles/_global.scss   |  1 +
 .../astro/src/vite-plugin-astro/compile.ts    | 44 +++++++-----
 packages/astro/src/vite-plugin-astro/hmr.ts   | 68 +++++++++++++++++++
 packages/astro/src/vite-plugin-astro/index.ts | 12 +++-
 6 files changed, 114 insertions(+), 20 deletions(-)
 create mode 100644 .changeset/brown-dancers-perform.md
 create mode 100644 examples/fast-build/src/styles/_global.scss
 create mode 100644 packages/astro/src/vite-plugin-astro/hmr.ts
diff --git a/.changeset/brown-dancers-perform.md b/.changeset/brown-dancers-perform.md
new file mode 100644
index 000000000..23591a3d5
--- /dev/null
+++ b/.changeset/brown-dancers-perform.md
@@ -0,0 +1,5 @@
+---
+'astro': patch
+---
+
+Fix for CSS superset support and HMR in the static build
diff --git a/examples/fast-build/src/pages/index.astro b/examples/fast-build/src/pages/index.astro
index 4d13ae5a5..0b7e7ff21 100644
--- a/examples/fast-build/src/pages/index.astro
+++ b/examples/fast-build/src/pages/index.astro
@@ -17,9 +17,9 @@ import ExternalHoisted from '../components/ExternalHoisted.astro';
 			}
 		</style>
 		<style lang="scss">
-			$color: purple;
+			@import "../styles/_global.scss";
 			h2 {
-				color: purple;
+				color: $color;
 			}
 		</style>
 		<style define:vars={{ color: 'blue' }}>
diff --git a/examples/fast-build/src/styles/_global.scss b/examples/fast-build/src/styles/_global.scss
new file mode 100644
index 000000000..27620a746
--- /dev/null
+++ b/examples/fast-build/src/styles/_global.scss
@@ -0,0 +1 @@
+$color: tan;
diff --git a/packages/astro/src/vite-plugin-astro/compile.ts b/packages/astro/src/vite-plugin-astro/compile.ts
index 2059a6680..616ec7aec 100644
--- a/packages/astro/src/vite-plugin-astro/compile.ts
+++ b/packages/astro/src/vite-plugin-astro/compile.ts
@@ -8,7 +8,7 @@ import { fileURLToPath } from 'url';
 import { transform } from '@astrojs/compiler';
 import { transformWithVite } from './styles.js';
 
-type CompilationCache = Map<string, TransformResult>;
+type CompilationCache = Map<string, CompileResult>;
 
 const configCache = new WeakMap<AstroConfig, CompilationCache>();
 
@@ -26,7 +26,9 @@ function isSSR(options: undefined | boolean | { ssr: boolean }): boolean {
 	return false;
 }
 
-async function compile(config: AstroConfig, filename: string, source: string, viteTransform: TransformHook, opts: boolean | undefined) {
+type CompileResult = TransformResult & { rawCSSDeps: Set<string> };
+
+async function compile(config: AstroConfig, filename: string, source: string, viteTransform: TransformHook, opts: boolean | undefined): Promise<CompileResult> {
 	// pages and layouts should be transformed as full documents (implicit <head> <body> etc)
 	// everything else is treated as a fragment
 	const filenameURL = new URL(`file://${filename}`);
@@ -34,6 +36,7 @@ async function compile(config: AstroConfig, filename: string, source: string, vi
 	const isPage = normalizedID.startsWith(fileURLToPath(config.pages)) || normalizedID.startsWith(fileURLToPath(config.layouts));
 	const pathname = filenameURL.pathname.substr(config.projectRoot.pathname.length - 1);
 
+	let rawCSSDeps = new Set<string>();
 	let cssTransformError: Error | undefined;
 
 	// Transform from `.astro` to valid `.ts`
@@ -51,21 +54,20 @@ async function compile(config: AstroConfig, filename: string, source: string, vi
 		// TODO add experimental flag here
 		preprocessStyle: async (value: string, attrs: Record<string, string>) => {
 			const lang = `.${attrs?.lang || 'css'}`.toLowerCase();
+
 			try {
-				let prefix = '';
-				// In the static build, strip away at-imports so that they can be resolved
-				// by the pseudo-module that gets created.
+				// In the static build, grab any @import as CSS dependencies for HMR.
 				if (config.buildOptions.experimentalStaticBuild) {
-					value = value.replace(/(?:@import)\s(?:url\()?\s?["\'](.*?)["\']\s?\)?(?:[^;]*);?/gi, (match) => {
-						prefix += match;
-						// Replace with an empty string of the same length, to preserve source maps.
-						return new Array(match.length).fill(' ').join('');
+					value.replace(/(?:@import)\s(?:url\()?\s?["\'](.*?)["\']\s?\)?(?:[^;]*);?/gi, (match, spec) => {
+						rawCSSDeps.add(spec);
+						return match;
 					});
 				}
+
 				const result = await transformWithVite({
 					value,
 					lang,
-					id: filename,
+					id: normalizedID,
 					transformHook: viteTransform,
 					ssr: isSSR(opts),
 				});
@@ -79,7 +81,7 @@ async function compile(config: AstroConfig, filename: string, source: string, vi
 						map = result.map.toString();
 					}
 				}
-				const code = (prefix += result.code);
+				const code = result.code;
 				return { code, map };
 			} catch (err) {
 				// save error to throw in plugin context
@@ -92,7 +94,17 @@ async function compile(config: AstroConfig, filename: string, source: string, vi
 	// throw CSS transform errors here if encountered
 	if (cssTransformError) throw cssTransformError;
 
-	return transformResult;
+	const compileResult: CompileResult =  Object.create(transformResult, {
+		rawCSSDeps: {
+			value: rawCSSDeps
+		}
+	});
+
+	return compileResult;
+}
+
+export function isCached(config: AstroConfig, filename: string) {
+	return configCache.has(config) && (configCache.get(config)!).has(filename);
 }
 
 export function invalidateCompilation(config: AstroConfig, filename: string) {
@@ -102,7 +114,7 @@ export function invalidateCompilation(config: AstroConfig, filename: string) {
 	}
 }
 
-export async function cachedCompilation(config: AstroConfig, filename: string, source: string | null, viteTransform: TransformHook, opts: boolean | undefined) {
+export async function cachedCompilation(config: AstroConfig, filename: string, source: string | null, viteTransform: TransformHook, opts: boolean | undefined): Promise<CompileResult> {
 	let cache: CompilationCache;
 	if (!configCache.has(config)) {
 		cache = new Map();
@@ -118,7 +130,7 @@ export async function cachedCompilation(config: AstroConfig, filename: string, s
 		const fileUrl = new URL(`file://${filename}`);
 		source = await fs.promises.readFile(fileUrl, 'utf-8');
 	}
-	const transformResult = await compile(config, filename, source, viteTransform, opts);
-	cache.set(filename, transformResult);
-	return transformResult;
+	const compileResult = await compile(config, filename, source, viteTransform, opts);
+	cache.set(filename, compileResult);
+	return compileResult;
 }
diff --git a/packages/astro/src/vite-plugin-astro/hmr.ts b/packages/astro/src/vite-plugin-astro/hmr.ts
new file mode 100644
index 000000000..f39d3fecd
--- /dev/null
+++ b/packages/astro/src/vite-plugin-astro/hmr.ts
@@ -0,0 +1,68 @@
+import type { AstroConfig } from '../@types/astro';
+import type { ViteDevServer, ModuleNode, HmrContext } from '../core/vite';
+import type { PluginContext as RollupPluginContext, ResolvedId } from 'rollup';
+import { cachedCompilation, invalidateCompilation, isCached } from './compile.js';
+
+interface TrackCSSDependenciesOptions {
+	viteDevServer: ViteDevServer | null;
+	filename: string;
+	id: string;
+	deps: Set<string>;
+}
+
+export async function trackCSSDependencies(this: RollupPluginContext, opts: TrackCSSDependenciesOptions): Promise<void> {
+	const { viteDevServer, filename, deps, id } = opts;
+	// Dev, register CSS dependencies for HMR.
+	if(viteDevServer) {
+		const mod = viteDevServer.moduleGraph.getModuleById(id);
+		if(mod) {
+			const cssDeps = (await Promise.all(Array.from(deps).map((spec) => {
+				return this.resolve(spec, id);
+			}))).filter(Boolean).map(dep => (dep as ResolvedId).id);
+
+			const { moduleGraph } = viteDevServer;
+			// record deps in the module graph so edits to @import css can trigger
+			// main import to hot update
+			const depModules = new Set(mod.importedModules);
+			for (const dep of cssDeps) {
+				depModules.add(moduleGraph.createFileOnlyEntry(dep))
+			}
+
+			// Update the module graph, telling it about our CSS deps.
+			moduleGraph.updateModuleInfo(mod, depModules, new Set(), true);
+			for (const dep of cssDeps) {
+				this.addWatchFile(dep);					
+			}
+		}
+	}
+}
+
+export function handleHotUpdate(ctx: HmrContext, config: AstroConfig) {
+	// Invalidate the compilation cache so it recompiles
+	invalidateCompilation(config, ctx.file);
+
+	// go through each of these modules importers and invalidate any .astro compilation
+	// that needs to be rerun.
+	const filtered = new Set<ModuleNode>();
+	const files = new Set<string>();
+	for(const mod of ctx.modules) {
+		if(mod.file && isCached(config, mod.file)) {
+			filtered.add(mod);
+			files.add(mod.file);
+		}
+		for(const imp of mod.importers) {
+			if(imp.file && isCached(config, imp.file)) {
+				filtered.add(imp);
+				files.add(imp.file);
+			}
+		}
+	}
+
+	// Invalidate happens as a separate step because a single .astro file
+	// produces multiple CSS modules and we want to return all of those.
+	for(const file of files) {
+		invalidateCompilation(config, file);
+	}
+
+	return Array.from(filtered);
+}
diff --git a/packages/astro/src/vite-plugin-astro/index.ts b/packages/astro/src/vite-plugin-astro/index.ts
index c19f79842..198680388 100644
--- a/packages/astro/src/vite-plugin-astro/index.ts
+++ b/packages/astro/src/vite-plugin-astro/index.ts
@@ -9,6 +9,7 @@ import { getViteTransform, TransformHook } from './styles.js';
 import { parseAstroRequest } from './query.js';
 import { cachedCompilation, invalidateCompilation } from './compile.js';
 import ancestor from 'common-ancestor-path';
+import { trackCSSDependencies, handleHotUpdate } from './hmr.js';
 
 const FRONTMATTER_PARSE_REGEXP = /^\-\-\-(.*)^\-\-\-/ms;
 interface AstroPluginOptions {
@@ -28,6 +29,7 @@ export default function astro({ config, logging }: AstroPluginOptions): vite.Plu
 	}
 
 	let viteTransform: TransformHook;
+	let viteDevServer: vite.ViteDevServer | null = null;
 
 	// Variables for determing if an id starts with /src...
 	const srcRootWeb = config.src.pathname.slice(config.projectRoot.pathname.length - 1);
@@ -39,6 +41,9 @@ export default function astro({ config, logging }: AstroPluginOptions): vite.Plu
 		configResolved(resolvedConfig) {
 			viteTransform = getViteTransform(resolvedConfig);
 		},
+		configureServer(server) {
+			viteDevServer = server;
+		},
 		// note: don’t claim .astro files with resolveId() — it prevents Vite from transpiling the final JS (import.meta.globEager, etc.)
 		async resolveId(id) {
 			// serve sub-part requests (*?astro) as virtual modules
@@ -64,6 +69,10 @@ export default function astro({ config, logging }: AstroPluginOptions): vite.Plu
 					}
 
 					const transformResult = await cachedCompilation(config, normalizeFilename(filename), null, viteTransform, opts);
+
+					// Track any CSS dependencies so that HMR is triggered when they change.
+					await trackCSSDependencies.call(this, { viteDevServer, id, filename, deps: transformResult.rawCSSDeps });
+
 					const csses = transformResult.css;
 					const code = csses[query.index];
 
@@ -166,8 +175,7 @@ export default function astro({ config, logging }: AstroPluginOptions): vite.Plu
 			}
 		},
 		async handleHotUpdate(context) {
-			// Invalidate the compilation cache so it recompiles
-			invalidateCompilation(config, context.file);
+			return handleHotUpdate(context, config);
 		},
 	};
 }
