fix: allow nested captures

src/__tests__/capture.spec.luau

@@ -28,4 +28,21 @@ return function()
 		expect(deps[a]).to.be.ok()
 		expect(result).to.equal(1)
 	end)
+
+	it("allows nested captures", function()
+		local a = atom(1)
+		local b = atom(1)
+		local inner
+		local outer = store.capture(function()
+			inner = store.capture(function()
+				return b()
+			end)
+			return a()
+		end)
+
+		assert(outer[a], "outer did not capture a")
+		assert(outer[b], "outer did not capture b")
+		assert(inner[a] == nil, "inner captured a")
+		assert(inner[b], "inner did not capture b")
+	end)
 end

src/atom.luau

@@ -15,8 +15,8 @@ local function atom<T>(state: T, options: AtomOptions<T>?): Atom<T>
 
 	local function atom(...)
 		if select("#", ...) == 0 then
-			if store.capturing.current then
+			for set in store.capturing do
-				store.capturing.current[atom] = true
+				set[atom] = true
 			end
 
 			return state

src/store.luau

@@ -1,11 +1,11 @@
 local types = require(script.Parent.types)
 type Atom<T> = types.Atom<T>
+type Set<T> = types.Set<T>
 type WeakMap<K, V> = types.WeakMap<K, V>
 
 local listeners: WeakMap<Atom<any>, WeakMap<() -> (), unknown>> = setmetatable({}, { __mode = "k" })
-
+local capturing: Set<Set<Atom<any>>> = {}
-local capturing: { current: { [Atom<any>]: true }? } = {}
+local batched: Set<() -> ()> = {}
-local batched: { [() -> ()]: true } = {}
 local batching = false
 
 --[=[
@@ -81,18 +81,16 @@ end
 	@param molecule The function to run.
 	@return A tuple containing the captured atoms and the result of the function.
 ]=]
-local function capture<T>(molecule: () -> T): ({ [Atom<any>]: true }, T)
+local function capture<T>(molecule: () -> T): (Set<Atom<any>>, T)
 	if listeners[molecule] then
 		return { [molecule] = true }, molecule()
 	end
 
-	local dependencies = {}
+	local dependencies: Set<Atom<any>> = {}
-	local previous = capturing.current
+	capturing[dependencies] = true
-
-	capturing.current = dependencies
 
 	local result = try(molecule, function()
-		capturing.current = previous
+		capturing[dependencies] = nil
 	end)
 
 	return dependencies, result
@@ -135,16 +133,18 @@ local function peek<T, U...>(callback: ((U...) -> T) | T, ...: U...): T
 		return callback
 	end
 
-	if not capturing.current then
+	if not next(capturing) then
 		return callback(...)
 	end
 
-	local previous = capturing.current
+	local snapshot = table.clone(capturing)
 
-	capturing.current = nil
+	table.clear(capturing)
 
 	local result = try(callback, function()
-		capturing.current = previous
+		for set in snapshot do
+			capturing[set] = true
+		end
 	end, ...)
 
 	return result
@@ -157,7 +157,7 @@ end
 	@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: { [Atom<any>]: true }, listener: () -> (), ref: unknown?)
+local function connect(dependencies: Set<Atom<any>>, listener: () -> (), ref: unknown?)
 	for atom in dependencies do
 		listeners[atom][listener] = ref or true
 	end
@@ -169,7 +169,7 @@ end
 	@param dependencies The atoms to stop listening to.
 	@param listener The function to stop calling when the atoms change.
 ]=]
-local function disconnect(dependencies: { [Atom<any>]: true }, listener: () -> ())
+local function disconnect(dependencies: Set<Atom<any>>, listener: () -> ())
 	for atom in dependencies do
 		listeners[atom][listener] = nil
 	end
@@ -187,10 +187,10 @@ end
 ]=]
 local function reconnect<T>(
 	molecule: Atom<T>,
-	dependencies: { [Atom<any>]: true },
+	dependencies: Set<Atom<any>>,
 	listener: () -> (),
 	ref: unknown?
-): ({ [Atom<any>]: true }, T)
+): (Set<Atom<any>>, T)
 	disconnect(dependencies, listener)
 	local nextDependencies, state = capture(molecule)
 	connect(nextDependencies, listener, ref)

src/types.luau

@@ -1,3 +1,5 @@
+export type Set<T> = { [T]: true }
+
 export type WeakMap<K, V> = typeof(setmetatable({} :: { [K]: V }, { __mode = "k" }))
 
 --[=[