From 45d94fae012310d740f162e09c9ddf07c58f6f51 Mon Sep 17 00:00:00 2001
From: Simon Holthausen <simon.holthausen@vercel.com>
Date: Mon, 22 Jul 2024 15:57:00 +0200
Subject: [PATCH] fix: store a persistent HMR wrapper per component

The previous HMR logic created a new wrapper and source per HMR update, which meant things would either get wrapped more and more (prior to #12454) or would not get updated after the first update because the reference to the original would get lost (after #12454).

This fixes that by creating a registry within the HMR wrappers by filename are stored, retrieved and its signals updated. That way nothing gets lost, and nothing gets wrapped more than needed.

Fixes #12506
---
 .changeset/hip-ducks-roll.md                  |  5 +++
 .../3-transform/client/transform-client.js    | 20 ++--------
 .../svelte/src/internal/client/dev/hmr.js     | 39 ++++++++++++++++---
 3 files changed, 42 insertions(+), 22 deletions(-)
 create mode 100644 .changeset/hip-ducks-roll.md

diff --git a/.changeset/hip-ducks-roll.md b/.changeset/hip-ducks-roll.md
new file mode 100644
index 0000000000..bdce28fbc4
--- /dev/null
+++ b/.changeset/hip-ducks-roll.md
@@ -0,0 +1,5 @@
+---
+'svelte': patch
+---
+
+fix: store a persistent HMR wrapper per component
diff --git a/packages/svelte/src/compiler/phases/3-transform/client/transform-client.js b/packages/svelte/src/compiler/phases/3-transform/client/transform-client.js
index 0ef05eeac0..6d2dc84d59 100644
--- a/packages/svelte/src/compiler/phases/3-transform/client/transform-client.js
+++ b/packages/svelte/src/compiler/phases/3-transform/client/transform-client.js
@@ -417,9 +417,7 @@ export function client_component(source, analysis, options) {
 	);
 
 	if (options.hmr) {
-		const accept_fn_body = [
-			b.stmt(b.call('$.set', b.id('s'), b.member(b.id('module.default'), b.id('$.ORIGINAL'), true)))
-		];
+		const accept_fn_body = [b.stmt(b.call('$$hmr.update', b.id('module.default')))];
 
 		if (analysis.css.hash) {
 			// remove existing `<style>` element, in case CSS changed
@@ -438,20 +436,8 @@ export function client_component(source, analysis, options) {
 		}
 
 		const hmr = b.block([
-			b.const(b.id('s'), b.call('$.source', b.id(analysis.name))),
-			b.const(b.id('filename'), b.member(b.id(analysis.name), b.id('filename'))),
-			b.const(b.id('$$original'), b.id(analysis.name)),
-			b.stmt(b.assignment('=', b.id(analysis.name), b.call('$.hmr', b.id('s')))),
-			b.stmt(b.assignment('=', b.member(b.id(analysis.name), b.id('filename')), b.id('filename'))),
-			// Assign the original component to the wrapper so we can use it on hot reload patching,
-			// else we would call the HMR function two times
-			b.stmt(
-				b.assignment(
-					'=',
-					b.member(b.id(analysis.name), b.id('$.ORIGINAL'), true),
-					b.id('$$original')
-				)
-			),
+			b.const(b.id('$$hmr'), b.call('$.hmr', b.id(analysis.name))),
+			b.stmt(b.assignment('=', b.id(analysis.name), b.id('$$hmr.wrapper'))),
 			b.stmt(b.call('import.meta.hot.accept', b.arrow([b.id('module')], b.block(accept_fn_body))))
 		]);
 
diff --git a/packages/svelte/src/internal/client/dev/hmr.js b/packages/svelte/src/internal/client/dev/hmr.js
index 050901cd7c..44b9b72289 100644
--- a/packages/svelte/src/internal/client/dev/hmr.js
+++ b/packages/svelte/src/internal/client/dev/hmr.js
@@ -1,19 +1,37 @@
-/** @import { Source, Effect } from '#client' */
+/** @import { Effect } from '#client' */
+import { FILENAME, ORIGINAL } from '../../../constants.js';
 import { EFFECT_TRANSPARENT } from '../constants.js';
 import { block, branch, destroy_effect } from '../reactivity/effects.js';
+import { set, source } from '../reactivity/sources.js';
 import { set_should_intro } from '../render.js';
 import { get } from '../runtime.js';
 
+/**
+ * For each original component, store a persistent reference to the HMR wrapper
+ * @type {Map<string, {wrapper: any, update: (update: any) => void}>}
+ */
+const registry = new Map();
+
 /**
  * @template {(anchor: Comment, props: any) => any} Component
- * @param {Source<Component>} source
+ * @param {Component} original
  */
-export function hmr(source) {
+export function hmr(original) {
+	let result = registry.get(/** @type {any} */ (original)[FILENAME]);
+
+	if (result) {
+		// update the reference to the original component as it's now updated
+		result.wrapper[ORIGINAL] = original;
+		return result;
+	}
+
+	const component_source = source(original);
+
 	/**
 	 * @param {Comment} anchor
 	 * @param {any} props
 	 */
-	return function (anchor, props) {
+	const wrapper = function (anchor, props) {
 		let instance = {};
 
 		/** @type {Effect} */
@@ -22,7 +40,7 @@ export function hmr(source) {
 		let ran = false;
 
 		block(() => {
-			const component = get(source);
+			const component = get(component_source);
 
 			if (effect) {
 				// @ts-ignore
@@ -51,4 +69,15 @@ export function hmr(source) {
 
 		return instance;
 	};
+
+	// stash a reference to the original component to avoid adding more and more wrappers per HMR update
+	wrapper[ORIGINAL] = original;
+
+	result = {
+		wrapper,
+		update: (update) => set(component_source, update[ORIGINAL])
+	};
+	registry.set(/** @type {any} */ (original)[FILENAME], result);
+
+	return result;
 }