From 37a7a8347ce49cb773ed907260ffe169bd3aa15f Mon Sep 17 00:00:00 2001
From: Nate Moore <natemoo-re@users.noreply.github.com>
Date: Wed, 20 Apr 2022 16:46:40 -0500
Subject: [PATCH] Improved HMR (#3138)

* WIP: improved HMR

* fix(hmr): improve hmr filtering to avoid full reloads

* chore: add changeset
---
 .changeset/cyan-dots-admire.md                |  5 ++++
 packages/astro/src/runtime/client/hmr.ts      | 30 +++++++++++++++----
 packages/astro/src/vite-plugin-astro/hmr.ts   | 15 +++++-----
 packages/astro/src/vite-plugin-astro/index.ts | 17 +++++++++--
 4 files changed, 51 insertions(+), 16 deletions(-)
 create mode 100644 .changeset/cyan-dots-admire.md

diff --git a/.changeset/cyan-dots-admire.md b/.changeset/cyan-dots-admire.md
new file mode 100644
index 000000000..97c3991e1
--- /dev/null
+++ b/.changeset/cyan-dots-admire.md
@@ -0,0 +1,5 @@
+---
+'astro': patch
+---
+
+General HMR Improvements, including new HMR support for framework components that are only server-side rendered (do not have a `client:*` directive)
diff --git a/packages/astro/src/runtime/client/hmr.ts b/packages/astro/src/runtime/client/hmr.ts
index 180b91fbe..119dd421f 100644
--- a/packages/astro/src/runtime/client/hmr.ts
+++ b/packages/astro/src/runtime/client/hmr.ts
@@ -1,11 +1,8 @@
 if (import.meta.hot) {
-	// signal to Vite that we accept HMR
-	import.meta.hot.accept((mod) => mod);
+	import.meta.hot.accept(mod => mod);
 	const parser = new DOMParser();
-	import.meta.hot.on('astro:update', async ({ file }) => {
+	async function updatePage() {
 		const { default: diff } = await import('micromorph');
-		// eslint-disable-next-line no-console
-		console.log(`[vite] hot updated: ${file}`);
 		const html = await fetch(`${window.location}`).then((res) => res.text());
 		const doc = parser.parseFromString(html, 'text/html');
 
@@ -17,6 +14,27 @@ if (import.meta.hot) {
 				root.innerHTML = current?.innerHTML;
 			}
 		}
-		diff(document, doc);
+		return diff(document, doc);
+	}
+	async function updateAll(files: any[]) {
+		let hasAstroUpdate = false;
+		for (const file of files) {
+			if (file.acceptedPath.endsWith('.astro')) {
+				hasAstroUpdate = true;
+				continue;
+			}
+			if (file.acceptedPath.includes('vue&type=style')) {
+				const link = document.querySelector(`link[href="${file.acceptedPath}"]`);
+				if (link) {
+					link.replaceWith(link.cloneNode(true));
+				}
+			}
+		}
+		if (hasAstroUpdate) {
+			return updatePage()
+		}
+	}
+	import.meta.hot.on('vite:beforeUpdate', async (event) => {
+		await updateAll(event.updates);
 	});
 }
diff --git a/packages/astro/src/vite-plugin-astro/hmr.ts b/packages/astro/src/vite-plugin-astro/hmr.ts
index 433b15901..f63a63969 100644
--- a/packages/astro/src/vite-plugin-astro/hmr.ts
+++ b/packages/astro/src/vite-plugin-astro/hmr.ts
@@ -82,18 +82,17 @@ export async function handleHotUpdate(ctx: HmrContext, config: AstroConfig, logg
 		invalidateCompilation(config, file);
 	}
 
-	const mod = ctx.modules.find((m) => m.file === ctx.file);
+	// Bugfix: sometimes style URLs get normalized and end with `lang.css=`
+	// These will cause full reloads, so filter them out here
+	const mods = ctx.modules.filter(m => !m.url.endsWith('='));
+	const isSelfAccepting = mods.every(m => m.isSelfAccepting || m.url.endsWith('.svelte'));
 
-	// Note: this intentionally ONLY applies to Astro components
-	// HMR is handled for other file types by their respective plugins
 	const file = ctx.file.replace(config.root.pathname, '/');
-	if (ctx.file.endsWith('.astro')) {
-		ctx.server.ws.send({ type: 'custom', event: 'astro:update', data: { file } });
-	}
-	if (mod?.isSelfAccepting) {
+	if (isSelfAccepting) {
 		info(logging, 'astro', msg.hmr({ file }));
 	} else {
 		info(logging, 'astro', msg.reload({ file }));
 	}
-	return Array.from(filtered);
+
+	return mods
 }
diff --git a/packages/astro/src/vite-plugin-astro/index.ts b/packages/astro/src/vite-plugin-astro/index.ts
index d82bd2f33..428b70d55 100644
--- a/packages/astro/src/vite-plugin-astro/index.ts
+++ b/packages/astro/src/vite-plugin-astro/index.ts
@@ -174,7 +174,20 @@ export default function astro({ config, logging }: AstroPluginOptions): vite.Plu
 				let SUFFIX = '';
 				// Add HMR handling in dev mode.
 				if (!resolvedConfig.isProduction) {
-					SUFFIX += `\nif (import.meta.hot) import.meta.hot.accept((mod) => mod);`;
+					// HACK: extract dependencies from metadata until compiler static extraction handles them
+					const metadata = transformResult.code.split('$$createMetadata(')[1].split('});\n')[0]
+					const pattern = /specifier:\s*'([^']*)'/g;
+					const deps = new Set();
+					let match;
+					while (match = pattern.exec(metadata)?.[1]) {
+						deps.add(match);
+					}
+					// // import.meta.hot.accept(["${id}", "${Array.from(deps).join('","')}"], (...mods) => mods);
+					// We need to be self-accepting, AND
+					// we need an explicity array of deps to track changes for SSR-only components
+					SUFFIX += `\nif (import.meta.hot) {
+						import.meta.hot.accept(mod => mod);
+					}`;
 				}
 				// Add handling to inject scripts into each page JS bundle, if needed.
 				if (isPage) {
@@ -242,7 +255,7 @@ export default function astro({ config, logging }: AstroPluginOptions): vite.Plu
 					}
 				}
 
-				throw err;
+			throw err;
 			}
 		},
 		async handleHotUpdate(context) {