fix: skip nested captures (#22)

* fix: skip nested captures

* test: update atom bench

src/__tests__/benchmarks/atom.bench.luau

@@ -1,7 +1,7 @@
 -- This file is for use by Benchmarker
 -- https://boatbomber.itch.io/benchmarker
 
-local Package = game:GetService("ReplicatedStorage").Packages.charm
+local Package = game:GetService("ReplicatedStorage").DevPackages.Charm
 
 local atom = require(Package.atom)
 local computed = require(Package.computed)
@@ -11,19 +11,19 @@ local function parameterGenerator()
 end
 
 local benchmarks = {
-	atom = function(_, source: any)
+	atom = function(_, state: any)
 		for _ = 1, 500 do
-			source(math.random(1, 2))
+			state(math.random(1, 2))
 		end
 	end,
 
-	computed = function(_, source: any)
+	computed = function(_, state: any)
 		computed(function()
-			return 2 * source()
+			return 2 * state()
 		end)
 
 		for _ = 1, 500 do
-			source(math.random(1, 2))
+			state(math.random(1, 2))
 		end
 	end,
 }

src/__tests__/capture.spec.luau

@@ -1,5 +1,6 @@
 local atom = require(script.Parent.Parent.atom)
 local store = require(script.Parent.Parent.store)
+local subscribe = require(script.Parent.Parent.subscribe)
 
 return function()
 	it("gracefully handles yielding", function()
@@ -29,7 +30,7 @@ return function()
 		expect(result).to.equal(1)
 	end)
 
-	it("allows nested captures", function()
+	it("skips atoms in inner capture calls", function()
 		local a = atom(1)
 		local b = atom(1)
 		local inner
@@ -41,8 +42,20 @@ return function()
 		end)
 
 		assert(outer[a], "outer did not capture a")
-		assert(outer[b], "outer did not capture b")
+		assert(outer[b] == nil, "outer captured b")
 		assert(inner[a] == nil, "inner captured a")
 		assert(inner[b], "inner did not capture b")
 	end)
+
+	it("should not capture notified dependents", function()
+		local state = atom(1)
+
+		subscribe(state, function() end)
+
+		local deps = store.capture(function()
+			return state(2)
+		end)
+
+		expect(deps[state]).to.never.be.ok()
+	end)
 end

src/__tests__/init.spec.luau

@@ -13,5 +13,7 @@ return function()
 
 	afterEach(function()
 		table.clear(store.listeners :: any)
+		table.clear(store.capturing.stack)
+		store.capturing.index = 0
 	end)
 end

src/atom.luau

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

src/store.luau

@@ -4,10 +4,15 @@ 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 batched: Set<() -> ()> = {}
 local batching = false
 
+local capturing = {
+	stack = {} :: { Set<Atom<any>> },
+	index = 0,
+}
+
 --[=[
 	Calls the given function and returns the result. If the function yields or
 	throws an error, the thread is closed and an error is thrown. Regardless of
@@ -29,7 +34,9 @@ local function try<T, U...>(callback: (U...) -> T | any, finally: (() -> ())?, .
 
 		if coroutine.status(thread) == "suspended" then
 			local source, line, name = debug.info(callback, "sln")
+
 			coroutine.close(thread)
+
 			error(
 				"Yielding is not allowed in atom functions. Consider wrapping this code in a Promise or task.defer instead."
 					.. `\nFunction defined at: {source}:{line}`
@@ -37,6 +44,7 @@ local function try<T, U...>(callback: (U...) -> T | any, finally: (() -> ())?, .
 			)
 		elseif not success then
 			local source, line, name = debug.info(callback, "sln")
+
 			error(
 				"An error occurred while running an atom function"
 					.. `\nFunction defined at: {source}:{line}`
@@ -53,6 +61,7 @@ local function try<T, U...>(callback: (U...) -> T | any, finally: (() -> ())?, .
 	end
 
 	local success, result = pcall(callback, ...)
+
 	finally()
 	assert(success, result)
 
@@ -87,6 +96,34 @@ local function notify(atom: Atom<any>)
 	end
 end
 
+--[=[
+	Returns the result of the function without subscribing to changes. If a
+	non-function value is provided, it is returned as is.
+	
+	@param molecule The atom or molecule to get the state of.
+	@param args Arguments to pass to the molecule.
+	@return The current state.
+]=]
+local function peek<T, U...>(callback: ((U...) -> T) | T, ...: U...): T
+	if type(callback) ~= "function" then
+		return callback
+	end
+
+	if capturing.index == 0 then
+		return callback(...)
+	end
+
+	capturing.index += 1
+	capturing.stack[capturing.index] = {}
+
+	local result = try(callback, function()
+		capturing.stack[capturing.index] = nil
+		capturing.index -= 1
+	end, ...)
+
+	return result
+end
+
 --[=[
 	Captures all atoms that are read during the function call and returns them
 	along with the result of the function. Useful for tracking dependencies.
@@ -96,14 +133,17 @@ end
 ]=]
 local function capture<T>(molecule: () -> T): (Set<Atom<any>>, T)
 	if listeners[molecule] then
-		return { [molecule] = true }, molecule()
+		return { [molecule] = true }, peek(molecule)
 	end
 
 	local dependencies: Set<Atom<any>> = {}
-	capturing[dependencies] = true
+
+	capturing.index += 1
+	capturing.stack[capturing.index] = dependencies
 
 	local result = try(molecule, function()
-		capturing[dependencies] = nil
+		capturing.stack[capturing.index] = nil
+		capturing.index -= 1
 	end)
 
 	return dependencies, result
@@ -133,36 +173,6 @@ local function batch(callback: () -> ())
 	table.clear(batched)
 end
 
---[=[
-	Returns the result of the function without subscribing to changes. If a
-	non-function value is provided, it is returned as is.
-	
-	@param molecule The atom or molecule to get the state of.
-	@param args Arguments to pass to the molecule.
-	@return The current state.
-]=]
-local function peek<T, U...>(callback: ((U...) -> T) | T, ...: U...): T
-	if type(callback) ~= "function" then
-		return callback
-	end
-
-	if not next(capturing) then
-		return callback(...)
-	end
-
-	local snapshot = table.clone(capturing)
-
-	table.clear(capturing)
-
-	local result = try(callback, function()
-		for set in next, snapshot do
-			capturing[set] = true
-		end
-	end, ...)
-
-	return result
-end
-
 --[=[
 	Subscribes the listener to the changes of the given atoms.