From a90d131fed51f90707d773acb9ed0dbe7a469407 Mon Sep 17 00:00:00 2001
From: richard <56808540+littensy@users.noreply.github.com>
Date: Tue, 11 Jun 2024 17:13:43 -0700
Subject: [PATCH] feat: allow conditional dependencies (#4)

* feat: allow conditional dependencies

* docs: remove conditional deps warnings

* fix: allow nested captures

* docs: add dev flag

* fix: `try()` uses pcall if passed `finally()`

* docs: interval is optional

* docs: update readme

* docs: add client/server tags
---
 README.md                                |  25 +++-
 src/__tests__/benchmarks/atom.bench.luau |  19 +--
 src/__tests__/capture.spec.luau          |  38 +++++-
 src/__tests__/computed.spec.luau         |  31 +++++
 src/__tests__/effect.spec.luau           |  29 ++++
 src/__tests__/peek.spec.luau             |   8 +-
 src/__tests__/subscribe.spec.luau        |  28 ++++
 src/atom.luau                            |   4 +-
 src/computed.luau                        |   9 +-
 src/effect.luau                          |  15 +--
 src/index.d.ts                           |   7 +-
 src/store.luau                           | 161 +++++++++++++++++------
 src/subscribe.luau                       |  12 +-
 src/types.luau                           |   2 +
 14 files changed, 290 insertions(+), 98 deletions(-)

diff --git a/README.md b/README.md
index 2f6c1ac..d92304d 100644
--- a/README.md
+++ b/README.md
@@ -24,7 +24,7 @@
 
 - ⚛️ **Manage state with _atoms_.** Decompose state into small, distinct containers called _atoms_, as opposed to combining them into a single store.
 
-- 💪 **Minimal, yet powerful.** No concept of actions or middleware — write simple functions to read from and write to state.
+- 💪 **Minimal, yet powerful.** Less boilerplate — write simple functions to read from and write to state.
 
 - 🔬 **Immediate updates.** Listeners run asynchronously by default, avoiding the cascading effects of deferred updates and improving responsiveness.
 
@@ -51,6 +51,23 @@ Charm = "littensy/charm@VERSION"
 
 ---
 
+## 🐛 Debugging
+
+Charm provides a debug mode to help you identify potential bugs in your project. To enable debug mode, set the global `_G.__DEV__` flag to `true` at the entry point of your project.
+
+Enabling `__DEV__` adds a few helpful features:
+
+- Better error handling for molecules, listeners, and batched functions:
+
+  - Errors provide the function's name and line number.
+  - Yielding in these functions will throw an error.
+
+- Server state is validated for [remote event limitations](https://create.roblox.com/docs/scripting/events/remote#argument-limitations) before being passed to the client.
+
+Enabling debug mode in unit tests, storybooks, and other development environments can help you catch potential issues early. However, remember to turn off debug mode in production to avoid the performance overhead.
+
+---
+
 ## 📚 Reference
 
 ### `atom(state, options?)`
@@ -144,10 +161,6 @@ effect(() => {
 
 `effect` returns a cleanup function.
 
-#### Caveats
-
-- Dependencies are captured when the effect is created, but not when it re-runs. This means that atoms referenced conditionally might be "skipped" in the dependency list.
-
 ---
 
 ### `computed(molecule, options?)`
@@ -379,7 +392,7 @@ remotes.requestState.connect((player) => {
 
   - `atoms`: An object containing the atoms to sync. The keys should match the keys on the client.
 
-  - `interval`: The interval at which to send state updates to clients. Defaults to `0`, meaning updates are sent on the next frame.
+  - **optional** `interval`: The interval at which to send state updates to clients. Defaults to `0`, meaning updates are sent on the next frame.
 
 #### Returns
 
diff --git a/src/__tests__/benchmarks/atom.bench.luau b/src/__tests__/benchmarks/atom.bench.luau
index e7819f9..6663caa 100644
--- a/src/__tests__/benchmarks/atom.bench.luau
+++ b/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,
 }
diff --git a/src/__tests__/capture.spec.luau b/src/__tests__/capture.spec.luau
index 59caaa4..60b6631 100644
--- a/src/__tests__/capture.spec.luau
+++ b/src/__tests__/capture.spec.luau
@@ -2,21 +2,47 @@ 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)
+
+	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/__tests__/computed.spec.luau b/src/__tests__/computed.spec.luau
index b9b3a19..87d0d6b 100644
--- a/src/__tests__/computed.spec.luau
+++ b/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
diff --git a/src/__tests__/effect.spec.luau b/src/__tests__/effect.spec.luau
index a6bb66f..fd077c9 100644
--- a/src/__tests__/effect.spec.luau
+++ b/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
diff --git a/src/__tests__/peek.spec.luau b/src/__tests__/peek.spec.luau
index bde5def..a861517 100644
--- a/src/__tests__/peek.spec.luau
+++ b/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)
 
diff --git a/src/__tests__/subscribe.spec.luau b/src/__tests__/subscribe.spec.luau
index 5097213..d5256e5 100644
--- a/src/__tests__/subscribe.spec.luau
+++ b/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
diff --git a/src/atom.luau b/src/atom.luau
index 52abf93..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.isCapturing() then
-				store.captured[atom] = true
+			for set in store.capturing do
+				set[atom] = true
 			end
 
 			return state
diff --git a/src/computed.luau b/src/computed.luau
index 0b0a910..cdfa072 100644
--- a/src/computed.luau
+++ b/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
diff --git a/src/effect.luau b/src/effect.luau
index 4e27531..7da63af 100644
--- a/src/effect.luau
+++ b/src/effect.luau
@@ -6,31 +6,24 @@ type Cleanup = () -> ()
 	Runs the given callback immediately and whenever any atom it depends on
 	changes. Returns a cleanup function that unsubscribes the callback.
 	
-	Note that dependencies are only evaluated once when the effect is created,
-	and conditional dependencies may not work as expected.
-	
 	@param callback The function to run.
 	@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()
diff --git a/src/index.d.ts b/src/index.d.ts
index edcbd90..4c4cbbf 100644
--- a/src/index.d.ts
+++ b/src/index.d.ts
@@ -82,9 +82,6 @@ declare namespace Charm {
 	 * Runs the given callback immediately and whenever any atom it depends on
 	 * changes. Returns a cleanup function that unsubscribes the callback.
 	 *
-	 * Note that dependencies are only evaluated once when the effect is created,
-	 * so conditional dependencies may not work as expected.
-	 *
 	 * @param callback The function to run.
 	 * @returns A function that unsubscribes the callback.
 	 */
@@ -123,7 +120,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.
@@ -199,6 +196,7 @@ declare namespace Charm {
 		 * Creates a `ClientSyncer` object that receives patches from the server and
 		 * applies them to the local state.
 		 *
+		 * @client
 		 * @param options The atoms to synchronize with the server.
 		 * @returns A `ClientSyncer` object.
 		 */
@@ -207,6 +205,7 @@ declare namespace Charm {
 		 * Creates a `ServerSyncer` object that sends patches to the client and
 		 * hydrates the client's state.
 		 *
+		 * @server
 		 * @param options The atoms to synchronize with the client.
 		 * @returns A `ServerSyncer` object.
 		 */
diff --git a/src/store.luau b/src/store.luau
index 1847611..229d0dd 100644
--- a/src/store.luau
+++ b/src/store.luau
@@ -1,22 +1,71 @@
 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 captured: { [Atom<any>]: true } = {}
-local capturing = false
-
-local batched: { [() -> ()]: true } = {}
+local capturing: Set<Set<Atom<any>>> = {}
+local batched: Set<() -> ()> = {}
 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
+			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}`
+					.. if name == "" then "" else ` function {name}`
+			)
+		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}`
+					.. (if name == "" then "" else ` function {name}`)
+					.. `\nError: {result}`
+			)
+		end
+
+		return result
+	end
+
+	if not finally then
+		return callback(...)
+	end
+
+	local success, result = pcall(callback, ...)
+	finally()
+	assert(success, result)
+
+	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 +83,7 @@ local function notify(atom: Atom<any>)
 	end
 
 	for listener in next, table.clone(listeners[atom]) do
-		listener()
+		try(listener)
 	end
 end
 
@@ -45,17 +94,17 @@ 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): (Set<Atom<any>>, T)
+	if listeners[molecule] then
+		return { [molecule] = true }, molecule()
 	end
 
-	capturing = true
-	local result = callback()
-	capturing = false
+	local dependencies: Set<Atom<any>> = {}
+	capturing[dependencies] = true
 
-	local dependencies = table.clone(captured)
-	table.clear(captured)
+	local result = try(molecule, function()
+		capturing[dependencies] = nil
+	end)
 
 	return dependencies, result
 end
@@ -68,16 +117,17 @@ end
 ]=]
 local function batch(callback: () -> ())
 	if batching then
-		callback()
-		return
+		return callback()
 	end
 
 	batching = true
-	callback()
-	batching = false
+
+	try(callback, function()
+		batching = false
+	end)
 
 	for listener in batched do
-		listener()
+		try(listener)
 	end
 
 	table.clear(batched)
@@ -96,44 +146,79 @@ local function peek<T, U...>(callback: ((U...) -> T) | T, ...: U...): T
 		return callback
 	end
 
-	if not capturing then
+	if not next(capturing) then
 		return callback(...)
 	end
 
-	capturing = false
-	local result = callback(...)
-	capturing = true
+	local snapshot = table.clone(capturing)
+
+	table.clear(capturing)
+
+	local result = try(callback, function()
+		for set in snapshot do
+			capturing[set] = true
+		end
+	end, ...)
 
 	return result
 end
 
 --[=[
-	Sets whether the store is capturing atoms. If `true`, atoms that are read
-	will be tracked and returned by `capture`.
-
-	@param value Whether to capture atoms.
+	Subscribes the listener to the changes of the given 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: Set<Atom<any>>, listener: () -> (), ref: unknown?)
+	for atom in dependencies do
+		listeners[atom][listener] = ref or true
+	end
 end
 
 --[=[
-	Returns whether the store is currently capturing atoms.
-
-	@return `true` if the store is capturing atoms, otherwise `false`.
+	Unsubscribes the listener from the changes of the given atoms.
+	
+	@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: Set<Atom<any>>, 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: Set<Atom<any>>,
+	listener: () -> (),
+	ref: unknown?
+): (Set<Atom<any>>, 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,
 }
diff --git a/src/subscribe.luau b/src/subscribe.luau
index f285c90..1083d99 100644
--- a/src/subscribe.luau
+++ b/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
 
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" }))
 
 --[=[