fix: cancel reconnection when unsubscribing

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

src/computed.luau

@@ -20,9 +20,13 @@ local function computed<T>(molecule: Molecule<T>, options: AtomOptions<T>?): 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
 

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

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<T, U...>(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<Atom<any>>, listener: () -> (), ref: unknown?)
-	for atom in dependencies do
+local function connect(atoms: Set<Atom<any>>, 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<Atom<any>>, listener: () -> ())
-	for atom in dependencies do
+local function disconnect(atoms: Set<Atom<any>>, 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<T>(
-	molecule: Atom<T>,
-	dependencies: Set<Atom<any>>,
-	listener: () -> (),
-	ref: unknown?
-): (Set<Atom<any>>, 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,
 }

src/subscribe.luau

@@ -12,22 +12,31 @@ type Molecule<T> = types.Molecule<T>
 	@return A function that unsubscribes the callback.
 ]=]
 local function subscribe<T>(molecule: Molecule<T>, 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