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