feat: allow conditional dependencies

src/__tests__/benchmarks/atom.bench.luau

@@ -7,34 +7,23 @@ local atom = require(Package.atom)
 local computed = require(Package.computed)
 
 local function parameterGenerator()
-	return atom({ count = 0 })
+	return atom(0)
 end
 
 local benchmarks = {
 	atom = function(_, source: any)
 		for _ = 1, 500 do
-			source({ count = math.random(1, 2) })
+			source(math.random(1, 2))
 		end
 	end,
 
 	computed = function(_, source: any)
 		computed(function()
-			return { count = 2 * source().count }
+			return 2 * source()
 		end)
-		for _ = 1, 500 do
-			source({ count = math.random(1, 2) })
-		end
-	end,
 
-	computed_nested = function(_, source: any)
-		local double = computed(function()
-			return { count = 2 * source().count }
-		end)
-		computed(function()
-			return { count = 2 * double().count }
-		end)
 		for _ = 1, 500 do
-			source({ count = math.random(1, 2) })
+			source(math.random(1, 2))
 		end
 	end,
 }

src/__tests__/capture.spec.luau

@@ -2,21 +2,30 @@ local atom = require(script.Parent.Parent.atom)
 local store = require(script.Parent.Parent.store)
 
 return function()
+	it("gracefully handles yielding", function()
+		expect(function()
+			store.capture(function()
+				task.wait(10)
+				return 1
+			end)
+		end).to.throw()
+	end)
+
 	it("captures atoms from a function", function()
 		local a, b, c = atom(1), atom(1), atom(2)
-		local captured, result = store.capture(function()
+		local deps, result = store.capture(function()
 			return a() + b()
 		end)
-		expect(captured[a]).to.be.ok()
-		expect(captured[b]).to.be.ok()
-		expect(captured[c]).to.never.be.ok()
+		expect(deps[a]).to.be.ok()
+		expect(deps[b]).to.be.ok()
+		expect(deps[c]).to.never.be.ok()
 		expect(result).to.equal(2)
 	end)
 
 	it("captures an atom by itself", function()
 		local a = atom(1)
-		local captured, result = store.capture(a)
-		expect(captured[a]).to.be.ok()
+		local deps, result = store.capture(a)
+		expect(deps[a]).to.be.ok()
 		expect(result).to.equal(1)
 	end)
 end

src/__tests__/computed.spec.luau

@@ -68,4 +68,35 @@ return function()
 
 		assert(count(store.listeners[source]) == 0, "nested computed atom did not gc")
 	end)
+
+	it("tracks conditional dependencies", function()
+		local condition = atom(false)
+		local a = atom(0)
+		local b = atom(1)
+		local reruns = 0
+		local current
+
+		-- assign variable to avoid gc
+		local _ = computed(function()
+			reruns += 1
+			current = if condition() then a() else b()
+			return current
+		end)
+
+		expect(reruns).to.equal(1)
+		expect(current).to.equal(1)
+
+		condition(true)
+		expect(reruns).to.equal(2)
+		expect(current).to.equal(0)
+
+		a(1)
+		b(2)
+		expect(reruns).to.equal(3)
+		expect(current).to.equal(1)
+
+		condition(false)
+		expect(reruns).to.equal(4)
+		expect(current).to.equal(2)
+	end)
 end

src/__tests__/effect.spec.luau

@@ -43,4 +43,33 @@ return function()
 		unsubscribe()
 		expect(cleanups).to.equal(2)
 	end)
+
+	it("tracks conditional dependencies", function()
+		local condition = atom(false)
+		local a = atom(0)
+		local b = atom(1)
+		local reruns = 0
+		local current
+
+		effect(function()
+			reruns += 1
+			current = if condition() then a() else b()
+		end)
+
+		expect(reruns).to.equal(1)
+		expect(current).to.equal(1)
+
+		condition(true)
+		expect(reruns).to.equal(2)
+		expect(current).to.equal(0)
+
+		a(1)
+		b(2)
+		expect(reruns).to.equal(3)
+		expect(current).to.equal(1)
+
+		condition(false)
+		expect(reruns).to.equal(4)
+		expect(current).to.equal(2)
+	end)
 end

src/__tests__/peek.spec.luau

@@ -4,12 +4,12 @@ local store = require(script.Parent.Parent.store)
 return function()
 	it("calls an atom without capturing it", function()
 		local a, b, c = atom(1), atom(1), atom(1)
-		local captured, result = store.capture(function()
+		local deps, result = store.capture(function()
 			return a() + store.peek(b) + c()
 		end)
-		expect(captured[a]).to.be.ok()
-		expect(captured[b]).to.never.be.ok()
-		expect(captured[c]).to.be.ok()
+		expect(deps[a]).to.be.ok()
+		expect(deps[b]).to.never.be.ok()
+		expect(deps[c]).to.be.ok()
 		expect(result).to.equal(3)
 	end)
 

src/__tests__/subscribe.spec.luau

@@ -67,4 +67,32 @@ return function()
 
 		assert(count(store.listeners[source]) == 0, "computed atom did not collect")
 	end)
+
+	it("tracks conditional dependencies", function()
+		local condition = atom(false)
+		local a = atom(0)
+		local b = atom(1)
+		local reruns = 0
+		local current
+
+		subscribe(function()
+			return if condition() then a() else b()
+		end, function(state)
+			reruns += 1
+			current = state
+		end)
+
+		condition(true)
+		expect(reruns).to.equal(1)
+		expect(current).to.equal(0)
+
+		a(1)
+		b(2)
+		expect(reruns).to.equal(2)
+		expect(current).to.equal(1)
+
+		condition(false)
+		expect(reruns).to.equal(3)
+		expect(current).to.equal(2)
+	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.isCapturing() then
-				store.captured[atom] = true
+			if store.capturing.current then
+				store.capturing.current[atom] = true
 			end
 
 			return state

src/computed.luau

@@ -15,19 +15,18 @@ type Molecule<T> = types.Molecule<T>
 	@return A new read-only atom.
 ]=]
 local function computed<T>(molecule: Molecule<T>, options: AtomOptions<T>?): Molecule<T>
-	local captured, state = store.capture(molecule)
+	local dependencies, state = store.capture(molecule)
 	local computedAtom = atom(state, options)
 	local computedRef = setmetatable({ current = computedAtom }, { __mode = "v" })
 
 	local function listener()
 		if computedRef.current then
-			computedRef.current(molecule())
+			dependencies, state = store.reconnect(molecule, dependencies, listener, computedRef.current)
+			computedRef.current(state)
 		end
 	end
 
-	for atom in captured do
-		store.listeners[atom][listener] = computedAtom
-	end
+	store.connect(dependencies, listener, computedAtom)
 
 	return computedAtom
 end

src/effect.luau

@@ -13,24 +13,20 @@ type Cleanup = () -> ()
 	@return A function that unsubscribes the callback.
 ]=]
 local function effect(callback: () -> any): () -> ()
-	local captured, cleanup = store.capture(callback)
+	local dependencies, cleanup = store.capture(callback)
 
 	local function listener()
 		if cleanup then
 			cleanup()
 		end
 
-		cleanup = callback()
+		dependencies, cleanup = store.reconnect(callback, dependencies, listener)
 	end
 
-	for atom in captured do
-		store.listeners[atom][listener] = true
-	end
+	store.connect(dependencies, listener)
 
 	return function()
-		for atom in captured do
-			store.listeners[atom][listener] = nil
-		end
+		store.disconnect(dependencies, listener)
 
 		if cleanup then
 			cleanup()

src/index.d.ts

@@ -123,7 +123,7 @@ declare namespace Charm {
 	 * @param molecule The function to run.
 	 * @returns A tuple containing the captured atoms and the result of the function.
 	 */
-	function capture<State>(molecule: Molecule<State>): LuaTuple<[captured: Set<Atom<unknown>>, state: State]>;
+	function capture<State>(molecule: Molecule<State>): LuaTuple<[dependencies: Set<Atom<unknown>>, state: State]>;
 
 	/**
 	 * Notifies all subscribers of the given atom that the state has changed.

src/store.luau

@@ -4,19 +4,55 @@ type WeakMap<K, V> = types.WeakMap<K, V>
 
 local listeners: WeakMap<Atom<any>, WeakMap<() -> (), unknown>> = setmetatable({}, { __mode = "k" })
 
-local captured: { [Atom<any>]: true } = {}
-local capturing = false
-
+local capturing: { current: { [Atom<any>]: true }? } = {}
 local batched: { [() -> ()]: true } = {}
 local batching = false
 
+--[=[
+	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
+	the outcome, the `finally` function is called to clean up any resources.
+	
+	@param callback The function to run.
+	@param finally Cleanup logic to run before error handling.
+	@param ... Arguments to pass to the callback.
+	@return The result of the callback.
+]=]
+local function try<T, U...>(callback: (U...) -> T | any, finally: (() -> ())?, ...: U...): T
+	if _G.__DEV__ then
+		local thread = coroutine.create(callback)
+		local success, result = coroutine.resume(thread, ...)
+
+		if finally then
+			finally()
+		end
+
+		if coroutine.status(thread) == "suspended" then
+			coroutine.close(thread)
+			error("Yielding is not allowed in atom functions")
+		elseif not success then
+			error("An error occurred while running an atom function")
+		end
+
+		return result
+	end
+
+	local result = callback(...)
+
+	if finally then
+		finally()
+	end
+
+	return result
+end
+
 --[=[
 	Returns whether the given value is an atom.
 
 	@param value The value to check.
 	@return `true` if the value is an atom, otherwise `false`.
 ]=]
-local function isAtom(value: any)
+local function isAtom(value: any): boolean
 	return not not (value and listeners[value])
 end
 
@@ -34,7 +70,7 @@ local function notify(atom: Atom<any>)
 	end
 
 	for listener in next, table.clone(listeners[atom]) do
-		listener()
+		try(listener)
 	end
 end
 
@@ -45,17 +81,19 @@ end
 	@param molecule The function to run.
 	@return A tuple containing the captured atoms and the result of the function.
 ]=]
-local function capture<T>(callback: () -> T): ({ [Atom<any>]: true }, T)
-	if listeners[callback] then
-		return { [callback] = true }, callback()
+local function capture<T>(molecule: () -> T): ({ [Atom<any>]: true }, T)
+	if listeners[molecule] then
+		return { [molecule] = true }, molecule()
 	end
 
-	capturing = true
-	local result = callback()
-	capturing = false
+	local dependencies = {}
+	local previous = capturing.current
 
-	local dependencies = table.clone(captured)
-	table.clear(captured)
+	capturing.current = dependencies
+
+	local result = try(molecule, function()
+		capturing.current = previous
+	end)
 
 	return dependencies, result
 end
@@ -68,16 +106,17 @@ end
 ]=]
 local function batch(callback: () -> ())
 	if batching then
-		callback()
-		return
+		return callback()
 	end
 
 	batching = true
-	callback()
+
+	try(callback, function()
 		batching = false
+	end)
 
 	for listener in batched do
-		listener()
+		try(listener)
 	end
 
 	table.clear(batched)
@@ -96,44 +135,77 @@ local function peek<T, U...>(callback: ((U...) -> T) | T, ...: U...): T
 		return callback
 	end
 
-	if not capturing then
+	if not capturing.current then
 		return callback(...)
 	end
 
-	capturing = false
-	local result = callback(...)
-	capturing = true
+	local previous = capturing.current
+
+	capturing.current = nil
+
+	local result = try(callback, function()
+		capturing.current = previous
+	end, ...)
 
 	return result
 end
 
 --[=[
-	Sets whether the store is capturing atoms. If `true`, atoms that are read
-	will be tracked and returned by `capture`.
+	Subscribes the listener to the changes of the given atoms.
 	
-	@param value Whether to capture atoms.
+	@param dependencies 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 setCapturing(value: boolean)
-	capturing = value
+local function connect(dependencies: { [Atom<any>]: true }, listener: () -> (), ref: unknown?)
+	for atom in dependencies do
+		listeners[atom][listener] = ref or true
+	end
 end
 
 --[=[
-	Returns whether the store is currently capturing atoms.
+	Unsubscribes the listener from the changes of the given atoms.
 	
-	@return `true` if the store is capturing atoms, otherwise `false`.
+	@param dependencies The atoms to stop listening to.
+	@param listener The function to stop calling when the atoms change.
 ]=]
-local function isCapturing()
-	return capturing
+local function disconnect(dependencies: { [Atom<any>]: true }, listener: () -> ())
+	for atom in dependencies 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: { [Atom<any>]: true },
+	listener: () -> (),
+	ref: unknown?
+): ({ [Atom<any>]: true }, T)
+	disconnect(dependencies, listener)
+	local nextDependencies, state = capture(molecule)
+	connect(nextDependencies, listener, ref)
+	return nextDependencies, state
 end
 
 return {
 	listeners = listeners,
-	captured = captured,
+	capturing = capturing,
 	isAtom = isAtom,
 	notify = notify,
 	capture = capture,
 	batch = batch,
 	peek = peek,
-	setCapturing = setCapturing,
-	isCapturing = isCapturing,
+	connect = connect,
+	disconnect = disconnect,
+	reconnect = reconnect,
 }

src/subscribe.luau

@@ -16,20 +16,18 @@ local function subscribe<T>(molecule: Molecule<T>, callback: (state: T, prev: T)
 
 	local function listener()
 		local prevState = state
-		state = molecule()
+
+		captured, state = store.reconnect(molecule, captured, listener)
+
 		if state ~= prevState then
 			callback(state, prevState)
 		end
 	end
 
-	for atom in captured do
-		store.listeners[atom][listener] = true
-	end
+	store.connect(captured, listener)
 
 	return function()
-		for atom in captured do
-			store.listeners[atom][listener] = nil
-		end
+		store.disconnect(captured, listener)
 	end
 end