diff --git a/src/__tests__/effect.spec.luau b/src/__tests__/effect.spec.luau index fd077c9..9e50c0b 100644 --- a/src/__tests__/effect.spec.luau +++ b/src/__tests__/effect.spec.luau @@ -72,4 +72,22 @@ return function() expect(reruns).to.equal(4) expect(current).to.equal(2) end) + + it("allows self-cleanup", function() + local a = atom(1) + local cleanedUp = false + local cleanup + + cleanup = effect(function() + if cleanup then + assert(not cleanedUp, "cleanup ran twice") + cleanup() + cleanedUp = true + end + a() + end) + + a(2) + a(3) + end) end diff --git a/src/computed.luau b/src/computed.luau index cdfa072..7fed8f7 100644 --- a/src/computed.luau +++ b/src/computed.luau @@ -20,9 +20,13 @@ local function computed(molecule: Molecule, options: AtomOptions?): Mol local computedRef = setmetatable({ current = computedAtom }, { __mode = "v" }) local function listener() - if computedRef.current then - dependencies, state = store.reconnect(molecule, dependencies, listener, computedRef.current) - computedRef.current(state) + local computedAtom = computedRef.current + + if computedAtom then + store.disconnect(dependencies, listener) + dependencies, state = store.capture(molecule) + store.connect(dependencies, listener, computedAtom) + computedAtom(state) end end diff --git a/src/effect.luau b/src/effect.luau index 7da63af..ae36188 100644 --- a/src/effect.luau +++ b/src/effect.luau @@ -9,20 +9,31 @@ type Cleanup = () -> () @param callback The function to run. @return A function that unsubscribes the callback. ]=] -local function effect(callback: () -> any): () -> () +local function effect(callback: () -> (() -> ())?): () -> () local dependencies, cleanup = store.capture(callback) + local disconnected = false local function listener() if cleanup then cleanup() end - dependencies, cleanup = store.reconnect(callback, dependencies, listener) + store.disconnect(dependencies, listener) + dependencies, cleanup = store.capture(callback) + + if not disconnected then + store.connect(dependencies, listener) + end end store.connect(dependencies, listener) return function() + if disconnected then + return + end + + disconnected = true store.disconnect(dependencies, listener) if cleanup then diff --git a/src/store.luau b/src/store.luau index 229d0dd..e197048 100644 --- a/src/store.luau +++ b/src/store.luau @@ -126,7 +126,7 @@ local function batch(callback: () -> ()) batching = false end) - for listener in batched do + for listener in next, batched do try(listener) end @@ -155,7 +155,7 @@ local function peek(callback: ((U...) -> T) | T, ...: U...): T table.clear(capturing) local result = try(callback, function() - for set in snapshot do + for set in next, snapshot do capturing[set] = true end end, ...) @@ -166,50 +166,28 @@ end --[=[ Subscribes the listener to the changes of the given atoms. - @param dependencies The atoms to listen to. + @param atoms The atoms to listen to. @param listener The function to call when the atoms change. @param ref Optionally bind the lifetime of the listener to a value. ]=] -local function connect(dependencies: Set>, listener: () -> (), ref: unknown?) - for atom in dependencies do +local function connect(atoms: Set>, listener: () -> (), ref: unknown?) + for atom in next, atoms do listeners[atom][listener] = ref or true end end --[=[ - Unsubscribes the listener from the changes of the given atoms. + Unsubscribes the listener from every atom it was connected to. - @param dependencies The atoms to stop listening to. + @param atoms The atoms to stop listening to. @param listener The function to stop calling when the atoms change. ]=] -local function disconnect(dependencies: Set>, listener: () -> ()) - for atom in dependencies do +local function disconnect(atoms: Set>, listener: () -> ()) + for atom in next, atoms do listeners[atom][listener] = nil end end ---[=[ - Disconnects the listener and evaluates the molecule again to capture any - new dependencies. Useful for tracking conditional dependencies. - - @param molecule The atom or molecule to evaluate. - @param dependencies The current dependencies. - @param listener The listener to disconnect. - @param ref Optionally bind the lifetime of the listener to a value. - @return The new dependencies and the state of the molecule. -]=] -local function reconnect( - molecule: Atom, - dependencies: Set>, - listener: () -> (), - ref: unknown? -): (Set>, T) - disconnect(dependencies, listener) - local nextDependencies, state = capture(molecule) - connect(nextDependencies, listener, ref) - return nextDependencies, state -end - return { listeners = listeners, capturing = capturing, @@ -220,5 +198,4 @@ return { peek = peek, connect = connect, disconnect = disconnect, - reconnect = reconnect, } diff --git a/src/subscribe.luau b/src/subscribe.luau index 1083d99..8f13ddd 100644 --- a/src/subscribe.luau +++ b/src/subscribe.luau @@ -12,22 +12,31 @@ type Molecule = types.Molecule @return A function that unsubscribes the callback. ]=] local function subscribe(molecule: Molecule, callback: (state: T, prev: T) -> ()): () -> () - local captured, state = store.capture(molecule) + local dependencies, state = store.capture(molecule) + local disconnected = false local function listener() local prevState = state - captured, state = store.reconnect(molecule, captured, listener) + store.disconnect(dependencies, listener) + dependencies, state = store.capture(molecule) + + if not disconnected then + store.connect(dependencies, listener) + end if state ~= prevState then callback(state, prevState) end end - store.connect(captured, listener) + store.connect(dependencies, listener) return function() - store.disconnect(captured, listener) + if not disconnected then + disconnected = true + store.disconnect(dependencies, listener) + end end end