From 82fcdfc2fa0321cd1014fe9e5c9d79b8454f2f15 Mon Sep 17 00:00:00 2001
From: Tan Li Hau <tanhauhau@users.noreply.github.com>
Date: Tue, 5 Jan 2021 03:41:11 +0800
Subject: [PATCH] rework SSR store autosubscription (#5419)

---
 CHANGELOG.md                                  |   4 +
 src/compiler/compile/render_dom/Renderer.ts   |  53 +--------
 src/compiler/compile/render_dom/invalidate.ts |  65 +++++++++-
 src/compiler/compile/render_ssr/index.ts      | 112 ++++++++++++++----
 .../_config.js                                |   1 -
 .../_config.js                                |   1 -
 .../_config.js                                |   1 -
 .../_config.js                                |   3 +
 .../main.svelte                               |  31 +++++
 9 files changed, 193 insertions(+), 78 deletions(-)
 create mode 100644 test/runtime/samples/store-auto-resubscribe-immediate/_config.js
 create mode 100644 test/runtime/samples/store-auto-resubscribe-immediate/main.svelte

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 28121604d948..e77f65dea83c 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,9 @@
 # Svelte changelog
 
+## Unreleased
+
+* Rework SSR store handling to subscribe and unsubscribe as in DOM mode ([#3375](https://github.com/sveltejs/svelte/issues/3375), [#3582](https://github.com/sveltejs/svelte/issues/3582), [#3636](https://github.com/sveltejs/svelte/issues/3636))
+
 ## 3.31.1
 
 * Fix scrolling of element with resize listener by making the `<iframe>` have `z-index: -1` ([#5448](https://github.com/sveltejs/svelte/issues/5448))
diff --git a/src/compiler/compile/render_dom/Renderer.ts b/src/compiler/compile/render_dom/Renderer.ts
index be3feed1431c..2e76ac591d3e 100644
--- a/src/compiler/compile/render_dom/Renderer.ts
+++ b/src/compiler/compile/render_dom/Renderer.ts
@@ -6,6 +6,7 @@ import { x } from 'code-red';
 import { Node, Identifier, MemberExpression, Literal, Expression, BinaryExpression } from 'estree';
 import flatten_reference from '../utils/flatten_reference';
 import { reserved_keywords } from '../utils/reserved_keywords';
+import { renderer_invalidate } from './invalidate';
 
 interface ContextMember {
 	name: string;
@@ -168,57 +169,7 @@ export default class Renderer {
 	}
 
 	invalidate(name: string, value?, main_execution_context: boolean = false) {
-		const variable = this.component.var_lookup.get(name);
-		const member = this.context_lookup.get(name);
-
-		if (variable && (variable.subscribable && (variable.reassigned || variable.export_name))) {
-			return main_execution_context
-			  ? x`${`$$subscribe_${name}`}(${value || name})`
-			  : x`${`$$subscribe_${name}`}($$invalidate(${member.index}, ${value || name}))`;
-		}
-
-		if (name[0] === '$' && name[1] !== '$') {
-			return x`${name.slice(1)}.set(${value || name})`;
-		}
-
-		if (
-			variable && (
-				variable.module || (
-					!variable.referenced &&
-					!variable.is_reactive_dependency &&
-					!variable.export_name &&
-					!name.startsWith('$$')
-				)
-			)
-		) {
-			return value || name;
-		}
-
-		if (value) {
-			return x`$$invalidate(${member.index}, ${value})`;
-		}
-
-		// if this is a reactive declaration, invalidate dependencies recursively
-		const deps = new Set([name]);
-
-		deps.forEach(name => {
-			const reactive_declarations = this.component.reactive_declarations.filter(x =>
-				x.assignees.has(name)
-			);
-			reactive_declarations.forEach(declaration => {
-				declaration.dependencies.forEach(name => {
-					deps.add(name);
-				});
-			});
-		});
-
-		// TODO ideally globals etc wouldn't be here in the first place
-		const filtered = Array.from(deps).filter(n => this.context_lookup.has(n));
-		if (!filtered.length) return null;
-
-		return filtered
-			.map(n => x`$$invalidate(${this.context_lookup.get(n).index}, ${n})`)
-			.reduce((lhs, rhs) => x`${lhs}, ${rhs}`);
+		return renderer_invalidate(this, name, value, main_execution_context);
 	}
 
 	dirty(names: string[], is_reactive_declaration = false): Expression {
diff --git a/src/compiler/compile/render_dom/invalidate.ts b/src/compiler/compile/render_dom/invalidate.ts
index f88c566d96b5..46ffe48284e0 100644
--- a/src/compiler/compile/render_dom/invalidate.ts
+++ b/src/compiler/compile/render_dom/invalidate.ts
@@ -33,7 +33,7 @@ export function invalidate(renderer: Renderer, scope: Scope, node: Node, names:
 		if (main_execution_context && !variable.subscribable && variable.name[0] !== '$') {
 			return node;
 		}
-		return renderer.invalidate(variable.name, undefined, main_execution_context);
+		return renderer_invalidate(renderer, variable.name, undefined, main_execution_context);
 	}
 
 	if (!head) {
@@ -79,3 +79,66 @@ export function invalidate(renderer: Renderer, scope: Scope, node: Node, names:
 
 	return invalidate;
 }
+
+export function renderer_invalidate(renderer: Renderer, name: string, value?, main_execution_context: boolean = false) {
+	const variable = renderer.component.var_lookup.get(name);
+
+	if (variable && (variable.subscribable && (variable.reassigned || variable.export_name))) {
+		if (main_execution_context) {
+			return x`${`$$subscribe_${name}`}(${value || name})`;
+		} else {
+			const member = renderer.context_lookup.get(name);
+			return x`${`$$subscribe_${name}`}($$invalidate(${member.index}, ${value || name}))`;
+		}
+	}
+
+	if (name[0] === '$' && name[1] !== '$') {
+		return x`${name.slice(1)}.set(${value || name})`;
+	}
+
+	if (
+		variable && (
+			variable.module || (
+				!variable.referenced &&
+				!variable.is_reactive_dependency &&
+				!variable.export_name &&
+				!name.startsWith('$$')
+			)
+		)
+	) {
+		return value || name;
+	}
+
+	if (value) {
+		if (main_execution_context) {
+			return x`${value}`;
+		} else {
+			const member = renderer.context_lookup.get(name);
+			return x`$$invalidate(${member.index}, ${value})`;
+		}
+	}
+
+	if (main_execution_context) return;
+
+	// if this is a reactive declaration, invalidate dependencies recursively
+	const deps = new Set([name]);
+
+	deps.forEach(name => {
+		const reactive_declarations = renderer.component.reactive_declarations.filter(x =>
+			x.assignees.has(name)
+		);
+		reactive_declarations.forEach(declaration => {
+			declaration.dependencies.forEach(name => {
+				deps.add(name);
+			});
+		});
+	});
+
+	// TODO ideally globals etc wouldn't be here in the first place
+	const filtered = Array.from(deps).filter(n => renderer.context_lookup.has(n));
+	if (!filtered.length) return null;
+
+	return filtered
+		.map(n => x`$$invalidate(${renderer.context_lookup.get(n).index}, ${n})`)
+		.reduce((lhs, rhs) => x`${lhs}, ${rhs}`);
+}
diff --git a/src/compiler/compile/render_ssr/index.ts b/src/compiler/compile/render_ssr/index.ts
index a156a4fe6a4b..e00e698d6f95 100644
--- a/src/compiler/compile/render_ssr/index.ts
+++ b/src/compiler/compile/render_ssr/index.ts
@@ -6,6 +6,9 @@ import Renderer from './Renderer';
 import { INode as TemplateNode } from '../nodes/interfaces'; // TODO
 import Text from '../nodes/Text';
 import { LabeledStatement, Statement, Node } from 'estree';
+import { walk } from 'estree-walker';
+import { extract_names } from 'periscopic';
+import { invalidate } from '../render_dom/invalidate';
 
 export default function ssr(
 	component: Component,
@@ -38,24 +41,94 @@ export default function ssr(
 	const slots = uses_slots ? b`let $$slots = @compute_slots(#slots);` : null;
 
 	const reactive_stores = component.vars.filter(variable => variable.name[0] === '$' && variable.name[1] !== '$');
-	const reactive_store_values = reactive_stores
+	const reactive_store_subscriptions = reactive_stores
+		.filter(store => {
+			const variable = component.var_lookup.get(store.name.slice(1));
+			return !variable || variable.hoistable;
+		})
+		.map(({ name }) => {
+			const store_name = name.slice(1);
+			return b`
+				${component.compile_options.dev && b`@validate_store(${store_name}, '${store_name}');`}
+				${`$$unsubscribe_${store_name}`} = @subscribe(${store_name}, #value => ${name} = #value)
+				${store_name}.subscribe($$value => ${name} = $$value);
+			`;
+		});
+	const reactive_store_unsubscriptions = reactive_stores.map(
+		({ name }) => b`${`$$unsubscribe_${name.slice(1)}`}()`
+	);
+
+	const reactive_store_declarations = reactive_stores
 		.map(({ name }) => {
 			const store_name = name.slice(1);
 			const store = component.var_lookup.get(store_name);
-			if (store && store.hoistable) return null;
 
-			const assignment = b`${name} = @get_store_value(${store_name});`;
+			if (store && store.reassigned) {
+				const unsubscribe = `$$unsubscribe_${store_name}`;
+				const subscribe = `$$subscribe_${store_name}`;
 
-			return component.compile_options.dev
-				? b`@validate_store(${store_name}, '${store_name}'); ${assignment}`
-				: assignment;
-		})
-		.filter(Boolean);
+				return b`let ${name}, ${unsubscribe} = @noop, ${subscribe} = () => (${unsubscribe}(), ${unsubscribe} = @subscribe(${store_name}, $$value => ${name} = $$value), ${store_name})`;
+			}
+			return b`let ${name}, ${`$$unsubscribe_${store_name}`};`;
+		});
 
-	component.rewrite_props(({ name }) => {
+	// instrument get/set store value
+	if (component.ast.instance) {
+		let scope = component.instance_scope;
+		const map = component.instance_scope_map;
+
+		walk(component.ast.instance.content, {
+			enter(node: Node) {
+				if (map.has(node)) {
+					scope = map.get(node);
+				}
+			},
+			leave(node: Node) {
+				if (map.has(node)) {
+					scope = scope.parent;
+				}
+
+				if (node.type === 'AssignmentExpression' || node.type === 'UpdateExpression') {
+					const assignee = node.type === 'AssignmentExpression' ? node.left : node.argument;
+					const names = new Set(extract_names(assignee));
+					const to_invalidate = new Set<string>();
+
+					for (const name of names) {
+						const variable = component.var_lookup.get(name);
+						if (variable &&
+							!variable.hoistable &&
+							!variable.global &&
+							!variable.module &&
+							(
+								variable.subscribable || variable.name[0] === '$'
+							)) {
+								to_invalidate.add(variable.name);
+							}
+					}
+
+					if (to_invalidate.size) {
+						this.replace(
+							invalidate(
+								{ component } as any,
+								scope,
+								node,
+								to_invalidate,
+								true
+							)
+						);
+					}
+				}
+			}
+		});
+	}
+
+	component.rewrite_props(({ name, reassigned }) => {
 		const value = `$${name}`;
 
-		let insert = b`${value} = @get_store_value(${name})`;
+		let insert = reassigned
+			? b`${`$$subscribe_${name}`}()`
+			: b`${`$$unsubscribe_${name}`} = @subscribe(${name}, #value => $${value} = #value)`;
+
 		if (component.compile_options.dev) {
 			insert = b`@validate_store(${name}, '${name}'); ${insert}`;
 		}
@@ -99,35 +172,28 @@ export default function ssr(
 			do {
 				$$settled = true;
 
-				${reactive_store_values}
-
 				${reactive_declarations}
 
 				$$rendered = ${literal};
 			} while (!$$settled);
 
+			${reactive_store_unsubscriptions}
+
 			return $$rendered;
 		`
 		: b`
-			${reactive_store_values}
-
 			${reactive_declarations}
 
+			${reactive_store_unsubscriptions}
+
 			return ${literal};`;
 
 	const blocks = [
 		...injected.map(name => b`let ${name};`),
 		rest,
 		slots,
-		...reactive_stores.map(({ name }) => {
-			const store_name = name.slice(1);
-			const store = component.var_lookup.get(store_name);
-			if (store && store.hoistable) {
-				return b`let ${name} = @get_store_value(${store_name});`;
-			}
-			return b`let ${name};`;
-		}),
-
+		...reactive_store_declarations,
+		...reactive_store_subscriptions,
 		instance_javascript,
 		...parent_bindings,
 		css.code && b`$$result.css.add(#css);`,
diff --git a/test/runtime/samples/reactive-assignment-in-complex-declaration-with-store-2/_config.js b/test/runtime/samples/reactive-assignment-in-complex-declaration-with-store-2/_config.js
index a22aa878ffae..62f47a01c3ef 100644
--- a/test/runtime/samples/reactive-assignment-in-complex-declaration-with-store-2/_config.js
+++ b/test/runtime/samples/reactive-assignment-in-complex-declaration-with-store-2/_config.js
@@ -1,6 +1,5 @@
 // destructure to store value
 export default {
-	skip_if_ssr: true, // pending https://github.com/sveltejs/svelte/issues/3582
 	html: '<h1>2 2 xxx 5 6 9 10 2</h1>',
 	async test({ assert, target, component }) {
 		await component.update();
diff --git a/test/runtime/samples/reactive-assignment-in-complex-declaration-with-store/_config.js b/test/runtime/samples/reactive-assignment-in-complex-declaration-with-store/_config.js
index 6a613f73c4aa..d235d2a1a066 100644
--- a/test/runtime/samples/reactive-assignment-in-complex-declaration-with-store/_config.js
+++ b/test/runtime/samples/reactive-assignment-in-complex-declaration-with-store/_config.js
@@ -1,7 +1,6 @@
 // destructure to store
 export default {
 	html: '<h1>2 2 xxx 5 6 9 10 2</h1>',
-	skip_if_ssr: true,
 	async test({ assert, target, component }) {
 		await component.update();
 		assert.htmlEqual(target.innerHTML, '<h1>11 11 yyy 12 13 14 15 11</h1>');
diff --git a/test/runtime/samples/store-assignment-updates-property/_config.js b/test/runtime/samples/store-assignment-updates-property/_config.js
index 168f7d27a01d..a1c41b988861 100644
--- a/test/runtime/samples/store-assignment-updates-property/_config.js
+++ b/test/runtime/samples/store-assignment-updates-property/_config.js
@@ -5,7 +5,6 @@ export default {
 		<button></button>
 		<button></button>
 	`,
-	skip_if_ssr: true,
 
 	async test({ assert, component, target, window }) {
 		const [btn1, btn2] = target.querySelectorAll('button');
diff --git a/test/runtime/samples/store-auto-resubscribe-immediate/_config.js b/test/runtime/samples/store-auto-resubscribe-immediate/_config.js
new file mode 100644
index 000000000000..47a901d793e1
--- /dev/null
+++ b/test/runtime/samples/store-auto-resubscribe-immediate/_config.js
@@ -0,0 +1,3 @@
+export default {
+	html: '{"answer":4}'
+};
diff --git a/test/runtime/samples/store-auto-resubscribe-immediate/main.svelte b/test/runtime/samples/store-auto-resubscribe-immediate/main.svelte
new file mode 100644
index 000000000000..8dd40ae7aea1
--- /dev/null
+++ b/test/runtime/samples/store-auto-resubscribe-immediate/main.svelte
@@ -0,0 +1,31 @@
+<script>
+	import { writable } from '../../../../store';
+
+	let value = writable({ foo: 1, bar: 2 });
+	$value.foo = $value.foo + $value.bar; // 3
+	$value.bar = $value.foo * $value.bar; // 6
+
+	// should resubscribe immediately
+	value = writable({ foo: $value.foo + 2, bar: $value.bar - 2 }); // { foo: 5, bar: 4 }
+
+	// should mutate the store value
+	$value.baz = $value.foo + $value.bar; // { foo: 5, bar: 4, baz: 9 }
+
+	// should resubscribe immediately
+	value = writable({ qux: $value.baz - $value.foo }); // { qux: 4 }
+
+	// making sure instrumentation returns the expression value
+	$value = {
+		one: writable(
+			$value = {
+				two: ({ $value } = { $value: { fred: $value.qux } }) // { fred: 4 }
+			} // { two: { $value: { fred: 4 } } }
+		) // { one: { two: { $value: { fred: 4 } } } }
+	};
+
+	const one = $value.one;
+
+	value.update(val => ({ answer: $one.two.$value.fred })); // { answer: 4 }
+</script>
+
+{JSON.stringify($value)}