From 5fd064e70a2f798a97b5239cec594c4b26668d56 Mon Sep 17 00:00:00 2001
From: Richard <richerd.rdc@gmail.com>
Date: Sat, 1 Jun 2024 18:44:23 -0700
Subject: [PATCH] feat: add cleanup function to effects

---
 src/__tests__/effect.spec.luau | 18 +++++++++++++++++-
 src/effect.luau                | 28 ++++++++++++++++++++++++----
 2 files changed, 41 insertions(+), 5 deletions(-)

diff --git a/src/__tests__/effect.spec.luau b/src/__tests__/effect.spec.luau
index d1b0337..a6bb66f 100644
--- a/src/__tests__/effect.spec.luau
+++ b/src/__tests__/effect.spec.luau
@@ -24,7 +24,23 @@ return function()
 
 		unsubscribe()
 		a(4)
-
 		expect(current).to.equal(6)
 	end)
+
+	it("accepts a cleanup function", function()
+		local a = atom(0)
+		local cleanups = 0
+		local unsubscribe = effect(function()
+			a()
+			return function()
+				cleanups += 1
+			end
+		end)
+
+		expect(cleanups).to.equal(0)
+		a(1)
+		expect(cleanups).to.equal(1)
+		unsubscribe()
+		expect(cleanups).to.equal(2)
+	end)
 end
diff --git a/src/effect.luau b/src/effect.luau
index 7dfa960..4e27531 100644
--- a/src/effect.luau
+++ b/src/effect.luau
@@ -1,10 +1,26 @@
 local store = require(script.Parent.store)
 
-local function effect(molecule: () -> nil): () -> ()
-	local captured = store.capture(molecule)
+type Cleanup = () -> ()
+
+--[=[
+	Runs the given callback immediately and whenever any atom it depends on
+	changes. Returns a cleanup function that unsubscribes the callback.
+	
+	Note that dependencies are only evaluated once when the effect is created,
+	and conditional dependencies may not work as expected.
+	
+	@param callback The function to run.
+	@return A function that unsubscribes the callback.
+]=]
+local function effect(callback: () -> any): () -> ()
+	local captured, cleanup = store.capture(callback)
 
 	local function listener()
-		molecule()
+		if cleanup then
+			cleanup()
+		end
+
+		cleanup = callback()
 	end
 
 	for atom in captured do
@@ -15,7 +31,11 @@ local function effect(molecule: () -> nil): () -> ()
 		for atom in captured do
 			store.listeners[atom][listener] = nil
 		end
+
+		if cleanup then
+			cleanup()
+		end
 	end
 end
 
-return effect
+return effect :: (callback: (() -> ()) | (() -> Cleanup)) -> Cleanup