Refactor style-only HMR cache and performance

.changeset/eight-turtles-itch.md

@@ -0,0 +1,5 @@
+---
+"astro": patch
+---
+
+Improves Astro file style-only HMR caching and performance

packages/astro/src/core/compile/index.ts

@@ -1,8 +1,3 @@
-export {
-	cachedCompilation,
-	getCachedCompileResult,
-	invalidateCompilation,
-	isCached,
-} from './cache.js';
+export { compile } from './compile.js';
 export type { CompileProps, CompileResult } from './compile.js';
 export type { TransformStyle } from './types.js';

packages/astro/src/vite-plugin-astro/compile.ts

@@ -1,15 +1,17 @@
 import { transformWithEsbuild, type ESBuildTransformResult } from 'vite';
 import type { AstroConfig } from '../@types/astro.js';
-import { cachedCompilation, type CompileProps, type CompileResult } from '../core/compile/index.js';
+import { compile, type CompileProps, type CompileResult } from '../core/compile/index.js';
 import type { Logger } from '../core/logger/core.js';
 import { getFileInfo } from '../vite-plugin-utils/index.js';
+import type { CompileMetadata } from './types.js';
 
-interface CachedFullCompilation {
+interface CompileAstroOption {
 	compileProps: CompileProps;
+	astroFileToCompileMetadata: Map<string, CompileMetadata>;
 	logger: Logger;
 }
 
-interface FullCompileResult extends Omit<CompileResult, 'map'> {
+export interface CompileAstroResult extends Omit<CompileResult, 'map'> {
 	map: ESBuildTransformResult['map'];
 }
 
@@ -23,15 +25,16 @@ interface EnhanceCompilerErrorOptions {
 
 const FRONTMATTER_PARSE_REGEXP = /^\-\-\-(.*)^\-\-\-/ms;
 
-export async function cachedFullCompilation({
+export async function compileAstro({
 	compileProps,
+	astroFileToCompileMetadata,
 	logger,
-}: CachedFullCompilation): Promise<FullCompileResult> {
+}: CompileAstroOption): Promise<CompileAstroResult> {
 	let transformResult: CompileResult;
 	let esbuildResult: ESBuildTransformResult;
 
 	try {
-		transformResult = await cachedCompilation(compileProps);
+		transformResult = await compile(compileProps);
 		// Compile all TypeScript to JavaScript.
 		// Also, catches invalid JS/TS in the compiled output before returning.
 		esbuildResult = await transformWithEsbuild(transformResult.code, compileProps.filename, {
@@ -76,6 +79,13 @@ export async function cachedFullCompilation({
 		}
 	}
 
+	// Attach compile metadata to map for use by virtual modules
+	astroFileToCompileMetadata.set(compileProps.filename, {
+		originalCode: compileProps.source,
+		css: transformResult.css,
+		scripts: transformResult.scripts,
+	});
+
 	return {
 		...transformResult,
 		code: esbuildResult.code + SUFFIX,

packages/astro/src/vite-plugin-astro/hmr.ts

@@ -1,42 +1,32 @@
 import path from 'node:path';
 import { appendForwardSlash } from '@astrojs/internal-helpers/path';
 import type { HmrContext } from 'vite';
-import type { AstroConfig } from '../@types/astro.js';
-import type { cachedCompilation } from '../core/compile/index.js';
-import { invalidateCompilation, isCached, type CompileResult } from '../core/compile/index.js';
 import type { Logger } from '../core/logger/core.js';
+import type { CompileAstroResult } from './compile.js';
+import type { CompileMetadata } from './types.js';
 
 export interface HandleHotUpdateOptions {
-	config: AstroConfig;
 	logger: Logger;
+	compile: (code: string, filename: string) => Promise<CompileAstroResult>;
 	astroFileToCssAstroDeps: Map<string, Set<string>>;
-
-	compile: () => ReturnType<typeof cachedCompilation>;
-	source: string;
+	astroFileToCompileMetadata: Map<string, CompileMetadata>;
 }
 
 export async function handleHotUpdate(
 	ctx: HmrContext,
-	{ config, logger, astroFileToCssAstroDeps, compile, source }: HandleHotUpdateOptions
+	{ logger, compile, astroFileToCssAstroDeps, astroFileToCompileMetadata }: HandleHotUpdateOptions
 ) {
-	let isStyleOnlyChange = false;
-	if (ctx.file.endsWith('.astro') && isCached(config, ctx.file)) {
-		// Get the compiled result from the cache
-		const oldResult = await compile();
-		// Skip HMR if source isn't changed
-		if (oldResult.source === source) return [];
-		// Invalidate to get fresh, uncached result to compare it to
-		invalidateCompilation(config, ctx.file);
-		const newResult = await compile();
-		if (isStyleOnlyChanged(oldResult, newResult)) {
-			isStyleOnlyChange = true;
-		}
-	} else {
-		invalidateCompilation(config, ctx.file);
-	}
-
-	if (isStyleOnlyChange) {
+	const oldCode = astroFileToCompileMetadata.get(ctx.file)?.originalCode;
+	const newCode = await ctx.read();
+	// If only the style code has changed, e.g. editing the `color`, then we can directly invalidate
+	// the Astro CSS virtual modules only. The main Astro module's JS result will be the same and doesn't
+	// need to be invalidated.
+	if (oldCode && isStyleOnlyChanged(oldCode, newCode)) {
 		logger.debug('watch', 'style-only change');
+		// Re-compile the main Astro component (even though we know its JS result will be the same)
+		// so that `astroFileToCompileMetadata` gets a fresh set of compile metadata to be used
+		// by the virtual modules later in the `load()` hook.
+		await compile(newCode, ctx.file);
 		// Only return the Astro styles that have changed!
 		return ctx.modules.filter((mod) => mod.id?.includes('astro&type=style'));
 	}
@@ -68,25 +58,39 @@ export async function handleHotUpdate(
 	}
 }
 
-function isStyleOnlyChanged(oldResult: CompileResult, newResult: CompileResult) {
-	return (
-		normalizeCode(oldResult.code) === normalizeCode(newResult.code) &&
-		// If style tags are added/removed, we need to regenerate the main Astro file
-		// so that its CSS imports are also added/removed
-		oldResult.css.length === newResult.css.length &&
-		!isArrayEqual(oldResult.css, newResult.css)
-	);
-}
+const frontmatterRE = /^\-\-\-.*?^\-\-\-/ms;
+const scriptRE = /<script(?:\s.*?)?>.*?<\/script>/gs;
+const styleRE = /<style(?:\s.*?)?>.*?<\/style>/gs;
+
+function isStyleOnlyChanged(oldCode: string, newCode: string) {
+	if (oldCode === newCode) return false;
+
+	// Before we can regex-capture style tags, we remove the frontmatter and scripts
+	// first as they could contain false-positive style tag matches. At the same time,
+	// we can also compare if they have changed and early out.
+
+	// Strip off and compare frontmatter
+	let oldFrontmatter = '';
+	let newFrontmatter = '';
+	oldCode = oldCode.replace(frontmatterRE, (m) => ((oldFrontmatter = m), ''));
+	newCode = newCode.replace(frontmatterRE, (m) => ((newFrontmatter = m), ''));
+	if (oldFrontmatter !== newFrontmatter) return false;
 
-const astroStyleImportRE = /import\s*"[^"]+astro&type=style[^"]+";/g;
-const sourceMappingUrlRE = /\/\/# sourceMappingURL=[^ ]+$/gm;
+	// Strip off and compare scripts
+	const oldScripts: string[] = [];
+	const newScripts: string[] = [];
+	oldCode = oldCode.replace(scriptRE, (m) => (oldScripts.push(m), ''));
+	newCode = newCode.replace(scriptRE, (m) => (newScripts.push(m), ''));
+	if (!isArrayEqual(oldScripts, newScripts)) return false;
 
-/**
- * Remove style-related code and sourcemap from the final astro output so they
- * can be compared between non-style code
- */
-function normalizeCode(code: string) {
-	return code.replace(astroStyleImportRE, '').replace(sourceMappingUrlRE, '').trim();
+	// Finally, we can compare styles
+	const oldStyles: string[] = [];
+	const newStyles: string[] = [];
+	oldCode.match(styleRE)?.forEach((m) => oldStyles.push(m));
+	newCode.match(styleRE)?.forEach((m) => newStyles.push(m));
+	// The length must also be the same for style only change.  If style tags are added/removed,
+	// we need to regenerate the main Astro file so that its CSS imports are also added/removed
+	return oldStyles.length === newStyles.length && !isArrayEqual(oldStyles, newStyles);
 }
 
 function isArrayEqual(a: any[], b: any[]) {