feat: pass cleanup function to effect (#40)

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 function effect(callback: (cleanup: Cleanup) -> Cleanup?): Cleanup
-	local dependencies, cleanup = store.capture(callback)
+	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)
+	function disconnect()
-
-	return function()
 		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

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)
+local function capture<T, U...>(callback: (U...) -> T, ...: U...): (Set<Atom<any>>, T)
-	if listeners[callback] then
+	-- If the callback is an atom, return it immediately
-		return { [callback] = true }, peek(callback)
+	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

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