From b9d576c876e4a4dc565f42e8e25259fc2402e1ca Mon Sep 17 00:00:00 2001
From: Richard <richerd.rdc@gmail.com>
Date: Sun, 9 Jun 2024 20:39:17 -0700
Subject: [PATCH] fix: allow nested captures

---
 src/__tests__/capture.spec.luau | 17 +++++++++++++++++
 src/atom.luau                   |  4 ++--
 src/store.luau                  | 34 ++++++++++++++++-----------------
 src/types.luau                  |  2 ++
 4 files changed, 38 insertions(+), 19 deletions(-)

diff --git a/src/__tests__/capture.spec.luau b/src/__tests__/capture.spec.luau
index 37b37ff..60b6631 100644
--- a/src/__tests__/capture.spec.luau
+++ b/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
diff --git a/src/atom.luau b/src/atom.luau
index b17b9b0..fc6bf64 100644
--- a/src/atom.luau
+++ b/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
-				store.capturing.current[atom] = true
+			for set in store.capturing do
+				set[atom] = true
 			end
 
 			return state
diff --git a/src/store.luau b/src/store.luau
index 7cfe094..fb371ab 100644
--- a/src/store.luau
+++ b/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: { current: { [Atom<any>]: true }? } = {}
-local batched: { [() -> ()]: true } = {}
+local capturing: Set<Set<Atom<any>>> = {}
+local batched: Set<() -> ()> = {}
 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 previous = capturing.current
-
-	capturing.current = dependencies
+	local dependencies: Set<Atom<any>> = {}
+	capturing[dependencies] = true
 
 	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)
diff --git a/src/types.luau b/src/types.luau
index 846fd7f..16d19f2 100644
--- a/src/types.luau
+++ b/src/types.luau
@@ -1,3 +1,5 @@
+export type Set<T> = { [T]: true }
+
 export type WeakMap<K, V> = typeof(setmetatable({} :: { [K]: V }, { __mode = "k" }))
 
 --[=[