feat: add cleanup function to effects

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

src/effect.luau

@@ -1,10 +1,26 @@
 local store = require(script.Parent.store)
 
-local function effect(molecule: () -> nil): () -> ()
+type Cleanup = () -> ()
-	local captured = store.capture(molecule)
+
+--[=[
+	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