From ce5ae96c2bc8459334ece0a43b09dabf5311ba28 Mon Sep 17 00:00:00 2001
From: Richard <richerd.rdc@gmail.com>
Date: Wed, 10 Jul 2024 15:08:50 -0700
Subject: [PATCH] feat: allow factory to return nil

---
 src/observe.luau | 18 ++++++++++++------
 1 file changed, 12 insertions(+), 6 deletions(-)

diff --git a/src/observe.luau b/src/observe.luau
index f289c03..b71b0fe 100644
--- a/src/observe.luau
+++ b/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) -> ()
+) -> () -> ()