From 89f1dfeaa5d3956c7b79a7e3c1b8d653ab727bc3 Mon Sep 17 00:00:00 2001
From: richard <56808540+littensy@users.noreply.github.com>
Date: Sat, 25 Jan 2025 12:14:21 -0800
Subject: [PATCH] feat: pass cleanup function to effect (#40)

---
 src/charm/src/effect.luau    | 22 +++++++++++++++-------
 src/charm/src/store.luau     |  9 +++++----
 tests/charm/effect.spec.luau | 15 +++++++++++++++
 3 files changed, 35 insertions(+), 11 deletions(-)

diff --git a/src/charm/src/effect.luau b/src/charm/src/effect.luau
index 1b822b7..23f3294 100644
--- a/src/charm/src/effect.luau
+++ b/src/charm/src/effect.luau
@@ -9,9 +9,11 @@ type Cleanup = () -> ()
 	@param callback The function to run.
 	@return A function that unsubscribes the callback.
 ]=]
-local function effect(callback: () -> Cleanup?): Cleanup
-	local dependencies, cleanup = store.capture(callback)
+local function effect(callback: (cleanup: Cleanup) -> Cleanup?): Cleanup
+	local dependencies = {}
+	local cleanup: Cleanup?
 	local disconnected = false
+	local disconnect
 
 	local function listener()
 		if cleanup then
@@ -19,16 +21,14 @@ local function effect(callback: () -> Cleanup?): Cleanup
 		end
 
 		store.disconnect(dependencies, listener)
-		dependencies, cleanup = store.capture(callback)
+		dependencies, cleanup = store.capture(callback, disconnect)
 
 		if not disconnected then
 			store.connect(dependencies, listener)
 		end
 	end
 
-	store.connect(dependencies, listener)
-
-	return function()
+	function disconnect()
 		if disconnected then
 			return
 		end
@@ -40,6 +40,14 @@ local function effect(callback: () -> Cleanup?): Cleanup
 			cleanup()
 		end
 	end
+
+	dependencies, cleanup = store.capture(callback, disconnect)
+
+	if not disconnected then
+		store.connect(dependencies, listener)
+	end
+
+	return disconnect
 end
 
-return effect :: (callback: (() -> ()) | (() -> Cleanup)) -> Cleanup
+return effect :: (callback: ((cleanup: Cleanup) -> ()) | ((cleanup: Cleanup) -> Cleanup?)) -> Cleanup
diff --git a/src/charm/src/store.luau b/src/charm/src/store.luau
index 957b5a0..fbffa90 100644
--- a/src/charm/src/store.luau
+++ b/src/charm/src/store.luau
@@ -133,9 +133,10 @@ end
 	@param callback The function to run.
 	@return A tuple containing the captured atoms and the result of the function.
 ]=]
-local function capture<T>(callback: () -> T): (Set<Atom<any>>, T)
-	if listeners[callback] then
-		return { [callback] = true }, peek(callback)
+local function capture<T, U...>(callback: (U...) -> T, ...: U...): (Set<Atom<any>>, T)
+	-- If the callback is an atom, return it immediately
+	if listeners[callback :: any] then
+		return { [callback :: any] = true }, peek(callback)
 	end
 
 	local dependencies: Set<Atom<any>> = {}
@@ -146,7 +147,7 @@ local function capture<T>(callback: () -> T): (Set<Atom<any>>, T)
 	local result = try(callback, function()
 		capturing.stack[capturing.index] = nil
 		capturing.index -= 1
-	end)
+	end, ...)
 
 	return dependencies, result
 end
diff --git a/tests/charm/effect.spec.luau b/tests/charm/effect.spec.luau
index faa95b3..1e834d0 100644
--- a/tests/charm/effect.spec.luau
+++ b/tests/charm/effect.spec.luau
@@ -93,4 +93,19 @@ return function()
 		a(2)
 		a(3)
 	end)
+
+	it("allows immediate self-cleanup", function()
+		local a = atom(1)
+		local cleanedUp = false
+
+		effect(function(cleanup)
+			assert(not cleanedUp, "cleanup ran twice")
+			cleanup()
+			cleanedUp = true
+			a()
+		end)
+
+		a(2)
+		a(3)
+	end)
 end