From 98ec6a8e46c4be4652d613fb8e54ab71a854de9e Mon Sep 17 00:00:00 2001
From: Rich Harris <richard.a.harris@gmail.com>
Date: Thu, 9 May 2019 21:18:13 -0400
Subject: [PATCH] allow derivers to return cleanup functions - fixes #2553

---
 store.mjs           |  4 ++++
 test/store/index.js | 31 +++++++++++++++++++++++++++++++
 2 files changed, 35 insertions(+)

diff --git a/store.mjs b/store.mjs
index bc9bc46ed1..624ede3dde 100644
--- a/store.mjs
+++ b/store.mjs
@@ -51,11 +51,14 @@ export function derived(stores, fn, initial_value) {
 		const values = [];
 
 		let pending = 0;
+		let cleanup = noop;
 
 		const sync = () => {
 			if (pending) return;
+			cleanup();
 			const result = fn(single ? values[0] : values, set);
 			if (auto) set(result);
+			else cleanup = result || noop;
 		};
 
 		const unsubscribers = stores.map((store, i) => store.subscribe(
@@ -74,6 +77,7 @@ export function derived(stores, fn, initial_value) {
 
 		return function stop() {
 			run_all(unsubscribers);
+			cleanup();
 		};
 	});
 }
diff --git a/test/store/index.js b/test/store/index.js
index 44617d6d00..5f9176cffb 100644
--- a/test/store/index.js
+++ b/test/store/index.js
@@ -211,6 +211,37 @@ describe('store', () => {
 
 			unsubscribe();
 		});
+
+		it('calls a cleanup function', () => {
+			const num = writable(1);
+
+			const values = [];
+			const cleaned_up = [];
+
+			const d = derived(num, ($num, set) => {
+				set($num * 2);
+
+				return function cleanup() {
+					cleaned_up.push($num);
+				};
+			});
+
+			num.set(2);
+
+			const unsubscribe = d.subscribe(value => {
+				values.push(value);
+			});
+
+			num.set(3);
+			num.set(4);
+
+			assert.deepEqual(values, [4, 6, 8]);
+			assert.deepEqual(cleaned_up, [2, 3]);
+
+			unsubscribe();
+
+			assert.deepEqual(cleaned_up, [2, 3, 4]);
+		});
 	});
 
 	describe('get', () => {