fix: avoid reconnecting to dependencies after cleanup (#14)

* fix: cancel reconnection when unsubscribing

* test: remove print

* feat: allow factory to return nil

* perf: use `next` iteration

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/__tests__/sync/server.spec.luau

@@ -80,7 +80,6 @@ return function()
 		b(3)
 		server:_sendPatch(target)
 
-		print(payloads)
 		expect(#payloads).to.equal(2)
 		expect(payloads[1].type).to.equal("patch")
 		expect(payloads[1].data.a).to.equal(3)

src/__tests__/sync/validate.spec.luau

@@ -18,7 +18,7 @@ return function()
 			nested_atom = { value = atom(0), expected = false },
 		}
 
-		for key, case in cases do
+		for key, case in next, cases do
 			local success, result = pcall(function()
 				validate(case.value, key)
 			end)

src/atom.luau

@@ -15,7 +15,7 @@ local function atom<T>(state: T, options: AtomOptions<T>?): Atom<T>
 
 	local function atom(...)
 		if select("#", ...) == 0 then
-			for set in store.capturing do
+			for set in next, store.capturing do
 				set[atom] = true
 			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/mapped.luau

@@ -36,7 +36,7 @@ local function mapped<K0, V0, K1, V1>(molecule: Molecule<{ [K0]: V0 }>, mapper:
 		local mappedKeys = {}
 
 		-- TODO: Only call mapper if the item has changed.
-		for key, item in items do
+		for key, item in next, items do
 			local newItem, newKey = mapper(item, key)
 			if newKey == nil then
 				newKey = key :: any
@@ -48,7 +48,7 @@ local function mapped<K0, V0, K1, V1>(molecule: Molecule<{ [K0]: V0 }>, mapper:
 			end
 		end
 
-		for key in prevMappedItems do
+		for key in next, prevMappedItems do
 			if mappedKeys[key] == nil and mappedItems[key] == prevMappedItems[key] then
 				mappedItems[key] = nil
 			end

src/observe.luau

@@ -2,6 +2,8 @@ local effect = require(script.Parent.effect)
 local types = require(script.Parent.types)
 type Molecule<T> = types.Molecule<T>
 
+local function noop() end
+
 --[=[
 	Creates an instance of `factory` for each item in the atom's state, and
 	cleans up the instance when the item is removed. Returns a cleanup function
@@ -11,34 +13,38 @@ type Molecule<T> = types.Molecule<T>
 	@param factory The function that tracks the lifecycle of each item.
 	@return A function that unsubscribes all instances.
 ]=]
-local function observe<K, V>(molecule: Molecule<{ [K]: V }>, factory: (value: V, key: K) -> () -> ()): () -> ()
+local function observe<K, V>(molecule: Molecule<{ [K]: V }>, factory: (value: V, key: K) -> (() -> ())?): () -> ()
 	local connections: { [K]: () -> () } = {}
 
 	local unsubscribe = effect(function()
 		local state = molecule()
 
-		for key, disconnect in connections do
+		for key, disconnect in next, connections do
 			if state[key] == nil then
 				connections[key] = nil
 				disconnect()
 			end
 		end
 
-		for key, value in state do
+		for key, value in next, state do
 			if not connections[key] then
-				connections[key] = factory(value, key)
+				connections[key] = factory(value, key) or noop
 			end
 		end
 	end)
 
 	local function cleanup()
 		unsubscribe()
-		for _, disconnect in connections do
+		for _, disconnect in next, connections do
 			disconnect()
 		end
+		table.clear(connections)
 	end
 
 	return cleanup
 end
 
-return observe
+return observe :: <K, V>(
+	molecule: Molecule<{ [K]: V }>,
+	factory: ((value: V, key: K) -> () -> ()) | (value: V, key: K) -> ()
+) -> () -> ()

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
 

src/sync/client.luau

@@ -36,7 +36,7 @@ local function client(options: ClientOptions): ClientSyncer
 		for index = 1, select("#", ...) do
 			local payload = select(index, ...)
 
-			for key, state in payload.data do
+			for key, state in next, payload.data do
 				if payload.type == "patch" then
 					atoms[key](patch.apply(atoms[key](), state))
 				else

src/sync/patch.luau

@@ -13,7 +13,7 @@ end
 local function diff(prevState: { [any]: any }, nextState: { [any]: any })
 	local patches = table.clone(nextState)
 
-	for key, previous in prevState do
+	for key, previous in next, prevState do
 		local next = nextState[key]
 
 		if previous == next then
@@ -26,7 +26,7 @@ local function diff(prevState: { [any]: any }, nextState: { [any]: any })
 	end
 
 	if _G.__DEV__ then
-		for key, value in prevState do
+		for key, value in next, prevState do
 			validate(value, key)
 		end
 
@@ -50,7 +50,7 @@ local function apply(state: any, patches: any): any
 	local nextState = table.clone(state)
 	local stateIsArray = state[1] ~= nil
 
-	for key, patch in patches do
+	for key, patch in next, patches do
 		-- Diff-checking an array produces a sparse array, which will not be
 		-- preserved when converted to JSON. To prevent this, we turn string
 		-- keys back into numeric keys.

src/sync/server.luau

@@ -72,7 +72,7 @@ local function server(options: ServerOptions): ServerSyncer
 
 		local function getSnapshot()
 			local snapshot: { [string]: any } = {}
-			for key, atom in options.atoms do
+			for key, atom in next, options.atoms do
 				snapshot[key] = atom()
 			end
 			return snapshot
@@ -105,7 +105,7 @@ local function server(options: ServerOptions): ServerSyncer
 
 		-- Populate the initial state and snapshot for each atom.
 		-- Subscribe to each atom and update the state when it changes.
-		for key, atom in options.atoms do
+		for key, atom in next, options.atoms do
 			cleanups[key] = subscribe(atom, function(current, previous)
 				pushSnapshot(key, current, previous)
 				changed = true
@@ -122,7 +122,7 @@ local function server(options: ServerOptions): ServerSyncer
 			local payloads: { SyncPayload } = {}
 			local lastSnapshot
 
-			for index, snapshot in snapshots do
+			for index, snapshot in next, snapshots do
 				lastSnapshot = snapshot
 
 				if index == 1 then
@@ -138,14 +138,14 @@ local function server(options: ServerOptions): ServerSyncer
 			snapshots = { lastSnapshot }
 			changed = false
 
-			for _, player in Players:GetPlayers() do
+			for _, player in next, Players:GetPlayers() do
 				callback(player, unpack(payloads))
 			end
 		end, options.interval or 0)
 
 		return function()
 			disconnect()
-			for _, cleanup in cleanups do
+			for _, cleanup in next, cleanups do
 				cleanup()
 			end
 		end
@@ -170,7 +170,7 @@ local function server(options: ServerOptions): ServerSyncer
 
 		local payloads: { SyncPayload } = {}
 
-		for index, snapshot in snapshots do
+		for index, snapshot in next, snapshots do
 			if index == 1 then
 				continue
 			end

src/sync/validate.luau

@@ -8,7 +8,7 @@ local function isUnsafeTable(object: { [any]: any })
 	local objectSize = 0
 
 	-- All keys must have the same type
-	for key in object do
+	for key in next, object do
 		local currentType = type(key)
 
 		if not keyType and SAFE_KEYS[currentType] then