feat: add effect cleanups, add docs (#2)

* docs: update readme

* feat: add cleanup function to effects

* refactor: move peek function to store module

* test: add nested atom check

* docs: add doc comments

* fix: `effect()` type

* refactor: remove divider comment

README.md

@@ -20,9 +20,6 @@
 
 **Charm** is an atomic state management library inspired by [Jotai](https://jotai.org) and [Nanostores](https://github.com/nanostores/nanostores). Designed to be a more composable alternative to Reflex, Charm aims to address common criticisms of Rodux-like state containers to better address certain use cases.
 
-> [!NOTE]
-> The documentation is a work-in-progress! Please refer to the [examples](#🚀-examples) section for more information on how to use Charm.
-
 ## 🍀 Features
 
 - ⚛️ **Manage state with _atoms_.** Decompose state into small, distinct containers called _atoms_, as opposed to combining them into a single store.
@@ -54,39 +51,347 @@ Charm = "littensy/charm@VERSION"
 
 ---
 
-## 📚 API Reference
+## 📚 Reference
 
-### Charm
+### `atom(state, options?)`
 
-#### `atom(state, options?)`
+Call `atom` to create a state container with the value `state`.
 
-#### `effect(callback)`
+```ts
+const nameAtom = atom("John");
+const todosAtom = atom<string[]>([]);
+```
 
-#### `subscribe(atom, listener)`
+#### Parameters
 
-#### `computed(atom, options?)`
+- `state`: The value to assign to the atom initially.
 
-#### `observe(atom, observer)`
+- **optional** `options`: An object that configures the behavior of this atom.
 
-#### `mapped(atom, mapper)`
+  - **optional** `equals`: An equality function to determine whether the state has changed. By default, strict equality (`===`) is used.
 
-#### `peek(atom)`
+#### Returns
 
-### React
+The `atom` constructor returns an atom function with two possible operations:
 
-#### `useAtom(atom, dependencies?)`
+1. **Read the state.** Call the atom without arguments to get the current state.
+2. **Set the state.** Pass a new value or an updater function to change the state.
 
-### Sync
+```ts
+function newTodo() {
+	nameAtom("Jane");
+	nameAtom(); // "Jane"
+	todosAtom((todos) => [...todos, "New todo"]);
+}
+```
 
-#### `sync.client(options)`
+---
 
-#### `Client.sync(payload)`
+### `subscribe(atom, listener)`
 
-#### `sync.server(options)`
+Call `subscribe` to listen for changes to an atom. Changes to the atom will immediately notify all subscribers, passing the new state and the previous state as arguments.
 
-#### `Server.connect(callback)`
+```ts
+const nameAtom = atom("John");
 
-#### `Server.hydrate(player)`
+subscribe(nameAtom, (name, prevName) => {
+	print(name);
+});
+```
+
+You may also pass a _molecule_, or a function that derives a value from one or more atoms. The molecule will be memoized and only re-run when its dependencies change.
+
+```ts
+const getUppercase = () => nameAtom().upper();
+
+subscribe(getUppercase, (name) => {
+	print(name);
+});
+```
+
+#### Parameters
+
+- `atom`: An atom or molecule that you want to subscribe to. This can be an atom, or a function that reads from one or more atoms.
+
+- `listener`: A function that will be called whenever the atom changes. The listener will receive the new state and the previous state as arguments.
+
+#### Returns
+
+`subscribe` returns a cleanup function.
+
+---
+
+### `effect(callback)`
+
+Call `effect` to declare a side effect that runs when any atom that it depends on changes. The effect will run immediately and whenever its dependencies change.
+
+```ts
+const nameAtom = atom("John");
+
+effect(() => {
+	print(nameAtom());
+	return () => {
+		print("Changing name!");
+	};
+});
+```
+
+#### Parameters
+
+- `callback`: The function that runs your effect. The function is called once to retrieve its dependencies, and then again whenever they change. Your callback may return a cleanup function to run when the effect is removed or about to re-run.
+
+#### Returns
+
+`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?)`
+
+Call `computed` when you want to derive an expensive value from one or more atoms. The derived value will be memoized and only re-run when its dependencies change.
+
+```ts
+const todosAtom = atom<string[]>([]);
+const mapToUppercase = computed(() => {
+	return todosAtom().map((todo) => todo.upper());
+});
+```
+
+`computed` is useful when you have multiple subscribers that depend on the same derived value. By memoizing the value, you can avoid re-calculating it for each subscriber.
+
+#### Parameters
+
+- `molecule`: A function that reads from one or more atoms and returns a derived value.
+
+- **optional** [`options`](#parameters): An object that configures the behavior of this atom.
+
+#### Returns
+
+`computed` returns a read-only atom.
+
+---
+
+### `observe(atom, factory)`
+
+Call `observe` to run the factory when a key is added to the atom's state. Your factory can return a cleanup function to run when the key is removed or the observer is disposed.
+
+```ts
+const todosAtom = atom<{ [Id in string]?: Todo }>({});
+
+observe(todosAtom, (todo, key) => {
+	print(todo);
+	return () => {
+		print("Removing todo!");
+	};
+});
+```
+
+#### Parameters
+
+- `atom`: An atom or molecule that you want to observe. This can be a primitive atom, or a function that reads from one or more atoms. The atom should return a dictionary or an array of objects.
+
+- `factory`: A function called for each key in the atom's state. The factory will receive the entry as an argument and should return a cleanup function.
+
+#### Returns
+
+`observe` returns a cleanup function.
+
+#### Caveats
+
+- The factory will only run when a key is added or removed, not when the value at that key changes. If your data is not keyed by a unique and stable identifier, consider using `mapped` to transform it into a keyed object before passing it to `observe`.
+
+---
+
+### `mapped(atom, mapper)`
+
+Call `mapped` to transform the key-value pairs of an atom's state. The mapper function will be called for each key-value pair in the atom's state, and the result will be stored in a new atom.
+
+```ts
+const todosAtom = atom<Todo[]>([]);
+const todosById = mapped(todosAtom, (todo, index) => {
+	return $tuple(todo, todo.id);
+});
+```
+
+#### Parameters
+
+- `atom`: An atom or molecule that you want to map. This can be a primitive atom, or a function that reads from one or more atoms. The atom should return a dictionary or an array of objects.
+
+- `mapper`: A function called for each key-value pair in the atom's state. The mapper will receive the entry as an argument and should return a new key-value pair. The value you return determines the type of the mapped atom:
+
+  - If you return a tuple, the mapped atom returns a dictionary with the first element as the value and the second element as the key.
+
+  - If you return a value without a key, the mapped atom returns an array of the given values in the order they were mapped.
+
+  - If the first element is `undefined`, the entry will be omitted from the mapped atom.
+
+#### Returns
+
+`mapped` returns a read-only atom.
+
+---
+
+### `peek(value)`
+
+Call `peek` to get the current state of an atom without tracking it as a dependency in functions like `effect` and `subscribe`.
+
+```ts
+const nameAtom = atom("John");
+const ageAtom = atom(25);
+
+effect(() => {
+	const name = nameAtom();
+	const age = peek(ageAtom);
+});
+```
+
+#### Parameters
+
+- `value`: Any value. If the value is an atom, `peek` will return the current state of the atom without tracking it. Otherwise, it will return the value as-is.
+
+- **optional** `...args`: Additional arguments to pass to the value if it is a function.
+
+#### Returns
+
+`peek` returns the current state of the atom. If the value is not a function, it will return the value as-is.
+
+---
+
+### `batch(callback)`
+
+Call `batch` to group multiple state changes into a single update. The callback will run immediately and listeners will only be notified once all changes have been applied.
+
+```ts
+const nameAtom = atom("John");
+const ageAtom = atom(25);
+
+batch(() => {
+	nameAtom("Jane");
+	ageAtom(26);
+});
+```
+
+#### Parameters
+
+- `callback`: A function that makes multiple state changes. The changes will be batched together and listeners will only be notified once all changes have been applied.
+
+#### Returns
+
+`batch` does not return anything.
+
+---
+
+## 📘 React
+
+### `useAtom(atom, dependencies?)`
+
+Call `useAtom` at the top-level of a React component to read from an atom.
+
+```tsx
+import { useAtom } from "@rbxts/charm";
+import { todosAtom } from "./todos-atom";
+
+function TodosApp() {
+	const todos = useAtom(todosAtom);
+	// ...
+}
+```
+
+By default, the atom is subscribed to once when the component initially mounts. Optionally, you may pass an array of dependencies to `useAtom` if your atom should be memoized based on other values.
+
+```tsx
+const todos = useAtom(searchTodos(filter), [filter]);
+```
+
+#### Parameters
+
+- `atom`: An atom or molecule that you want to read from. This can be an atom, or a function that reads from one or more atoms.
+
+- **optional** `dependencies`: An array of values that the atom depends on. If the dependencies change, the atom will be re-subscribed to.
+
+#### Returns
+
+`useAtom` returns the current state of the atom.
+
+---
+
+## 📗 Charm Sync
+
+### `sync.client(options)`
+
+Call `sync.client` to create a client sync object. The object will sync the client's copy of the state with the server's state.
+
+```ts
+import { sync } from "@rbxts/charm";
+
+const client = sync.client({ atoms: atomsToSync });
+
+remotes.syncState.connect((payload) => {
+	client.sync(payload);
+});
+
+remotes.requestState.fire();
+```
+
+#### Parameters
+
+- `options`: An object to configure the client syncer.
+
+  - `atoms`: An object containing the atoms to sync. The keys should match the keys on the server.
+
+#### Returns
+
+`sync.client` returns a client sync object. The sync object has the following methods:
+
+- `client.sync(payload)` applies a state update from the server.
+
+#### Caveats
+
+- The client sync object does not handle network communication. You must implement your own network layer to send and receive state updates. This includes requesting the initial state, which is implemented via `requestState` in the example above.
+
+---
+
+### `sync.server(options)`
+
+Call `sync.server` to create a server sync object. The object handles sending state updates to clients at a specified interval, and hydrating clients with the initial state.
+
+```ts
+import { sync } from "@rbxts/charm";
+
+const server = sync.server({ atoms: atomsToSync });
+
+server.connect((player, payload) => {
+	remotes.syncState.fire(player, payload);
+});
+
+remotes.requestState.connect((player) => {
+	server.hydrate(player);
+});
+```
+
+#### Parameters
+
+- `options`: An object to configure the server syncer.
+
+  - `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.
+
+#### Returns
+
+`sync.server` returns a server sync object. The sync object has the following methods:
+
+- `server.connect(callback)` registers a callback to send state updates to clients. The callback will receive the player and the payload to send, and should send the payload to the client. The payload should not be mutated, so changes should be applied to a copy of the payload.
+
+- `server.hydrate(player)` sends the initial state to a player, calling the callback passed to `connect` with a payload containing the initial state.
+
+#### Caveats
+
+- The server sync object does not handle network communication. You must implement your own network layer to send and receive state updates. This includes sending the initial state, which is implemented via `requestState` in the example above.
 
 ---
 
@@ -98,18 +403,17 @@ Charm = "littensy/charm@VERSION"
 import { atom, subscribe } from "@rbxts/charm";
 
 const counterAtom = atom(0);
+
+// Create a derived atom that returns double the counter value
 const doubleCounterAtom = () => counterAtom() * 2;
 
-function incrementCounter() {
+// Runs after counterAtom is updated and prints double the new value
-	counterAtom((count) => count + 1);
-}
-
 subscribe(doubleCounterAtom, (value) => {
 	print(value);
 });
 
-counterAtom(1);
+counterAtom(1); // 2
-incrementCounter();
+counterAtom((count) => count + 1); // 4
 ```
 
 ### Counter component
@@ -154,11 +458,11 @@ const server = sync.server({ atoms });
 
 // Broadcast a state update to a specific player
 server.connect((player, payload) => {
-	remotes.sync.fire(player, payload);
+	remotes.syncState.fire(player, payload);
 });
 
 // Send initial state to a player upon request
-remotes.init.connect((player) => {
+remotes.requestState.connect((player) => {
 	server.hydrate(player);
 });
 ```
@@ -173,12 +477,12 @@ import * as atoms from "./atoms";
 const client = sync.client({ atoms });
 
 // Listen for incoming state changes from the server
-remotes.sync.connect((payload) => {
+remotes.syncState.connect((payload) => {
 	client.sync(payload);
 });
 
 // Request initial state from the server
-remotes.init.fire();
+remotes.requestState.fire();
 ```
 
 ---

src/__tests__/effect.spec.luau

@@ -24,7 +24,23 @@ return function()
 
 		unsubscribe()
 		a(4)
-
 		expect(current).to.equal(6)
 	end)
+
+	it("accepts a cleanup function", function()
+		local a = atom(0)
+		local cleanups = 0
+		local unsubscribe = effect(function()
+			a()
+			return function()
+				cleanups += 1
+			end
+		end)
+
+		expect(cleanups).to.equal(0)
+		a(1)
+		expect(cleanups).to.equal(1)
+		unsubscribe()
+		expect(cleanups).to.equal(2)
+	end)
 end

src/__tests__/peek.spec.luau

@@ -1,12 +1,11 @@
 local atom = require(script.Parent.Parent.atom)
-local peek = require(script.Parent.Parent.peek)
 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()
-			return a() + peek(b) + c()
+			return a() + store.peek(b) + c()
 		end)
 		expect(captured[a]).to.be.ok()
 		expect(captured[b]).to.never.be.ok()
@@ -15,8 +14,8 @@ return function()
 	end)
 
 	it("passes through arguments", function()
-		expect(peek(1)).to.equal(1)
+		expect(store.peek(1)).to.equal(1)
-		expect(peek(function()
+		expect(store.peek(function()
 			return 1
 		end)).to.equal(1)
 	end)

src/__tests__/sync/validate.spec.luau

@@ -1,3 +1,4 @@
+local atom = require(script.Parent.Parent.Parent.atom)
 local validate = require(script.Parent.Parent.Parent.sync.validate)
 
 return function()
@@ -13,6 +14,8 @@ return function()
 			number = { value = 1, expected = true },
 			number_infinity = { value = math.huge, expected = false },
 			number_nan = { value = 0 / 0, expected = false },
+
+			nested_atom = { value = atom(0), expected = false },
 		}
 
 		for key, case in cases do

src/atom.luau

@@ -1,9 +1,15 @@
-local peek = require(script.Parent.peek)
 local store = require(script.Parent.store)
 local types = require(script.Parent.types)
 type Atom<T> = types.Atom<T>
 type AtomOptions<T> = types.AtomOptions<T>
 
+--[=[
+	Creates a new atom with the given state.
+	
+	@param state The initial state.
+	@param options Optional configuration.
+	@return A new atom.
+]=]
 local function atom<T>(state: T, options: AtomOptions<T>?): Atom<T>
 	local equals = options and options.equals
 
@@ -16,7 +22,7 @@ local function atom<T>(state: T, options: AtomOptions<T>?): Atom<T>
 			return state
 		end
 
-		local nextState = peek(..., state)
+		local nextState = store.peek(..., state)
 
 		if state ~= nextState and not (equals and equals(state, nextState)) then
 			state = nextState

src/computed.luau

@@ -5,6 +5,15 @@ type Atom<T> = types.Atom<T>
 type AtomOptions<T> = types.AtomOptions<T>
 type Molecule<T> = types.Molecule<T>
 
+--[=[
+	Creates a read-only atom that derives its state from one or more atoms.
+	Used to avoid unnecessary recomputations if multiple listeners depend on
+	the same molecule.
+	
+	@param molecule The function that produces the state.
+	@param options Optional configuration.
+	@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 computedAtom = atom(state, options)

src/effect.luau

@@ -1,10 +1,26 @@
 local store = require(script.Parent.store)
 
-local function effect(molecule: () -> nil): () -> ()
+type Cleanup = () -> ()
-	local captured = store.capture(molecule)
+
+--[=[
+	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 function listener()
-		molecule()
+		if cleanup then
+			cleanup()
+		end
+
+		cleanup = callback()
 	end
 
 	for atom in captured do
@@ -15,7 +31,11 @@ local function effect(molecule: () -> nil): () -> ()
 		for atom in captured do
 			store.listeners[atom][listener] = nil
 		end
+
+		if cleanup then
+			cleanup()
+		end
 	end
 end
 
-return effect
+return effect :: (callback: (() -> ()) | (() -> Cleanup)) -> Cleanup

src/index.d.ts

@@ -9,35 +9,138 @@ type AnyMap<K, V> =
 	| (K extends string | number | symbol ? { readonly [Key in K]: V } : never);
 
 declare namespace Charm {
+	/**
+	 * A primitive state container that can be read from and written to. When the
+	 * state changes, all subscribers are notified.
+	 *
+	 * @template State The type of the state.
+	 * @param state The next state or a function that produces the next state.
+	 * @returns The current state, if no arguments are provided.
+	 */
 	interface Atom<State> extends Molecule<State> {
+		/**
+		 * @deprecated This property is not meant to be accessed directly.
+		 */
 		readonly __nominal: unique symbol;
+		/**
+		 * @param state The next state or a function that produces the next state.
+		 * @returns The current state, if no arguments are provided.
+		 */
 		(state: State | ((prev: State) => State)): void;
 	}
 
+	/**
+	 * A function that depends on one or more atoms and produces a state. Can be
+	 * used to derive state from atoms.
+	 *
+	 * @template State The type of the state.
+	 * @returns The current state.
+	 */
 	type Molecule<State> = () => State;
 
 	interface AtomOptions<State> {
+		/**
+		 * A function that determines whether the state has changed. By default,
+		 * a strict equality check (`===`) is used.
+		 */
 		equals?: (prev: State, next: State) => boolean;
 	}
 
+	/**
+	 * Creates a new atom with the given state.
+	 *
+	 * @template State The type of the state.
+	 * @param state The initial state.
+	 * @param options Optional configuration.
+	 * @returns A new atom.
+	 */
 	function atom<State>(state: State, options?: AtomOptions<State>): Atom<State>;
 
+	/**
+	 * Creates a read-only atom that derives its state from one or more atoms.
+	 * Used to avoid unnecessary recomputations if multiple listeners depend on
+	 * the same molecule.
+	 *
+	 * @param molecule The function that produces the state.
+	 * @param options Optional configuration.
+	 * @returns A new read-only atom.
+	 */
 	function computed<State>(molecule: Molecule<State>, options?: AtomOptions<State>): Molecule<State>;
 
+	/**
+	 * Subscribes to changes in the given atom or molecule. The callback is
+	 * called with the current state and the previous state immediately after a
+	 * change occurs.
+	 *
+	 * @param molecule The atom or molecule to subscribe to.
+	 * @param callback The function to call when the state changes.
+	 * @returns A function that unsubscribes the callback.
+	 */
 	function subscribe<State>(molecule: Molecule<State>, callback: (state: State, prev: State) => void): Cleanup;
 
-	function effect(callback: () => void): 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,
+	 * so conditional dependencies may not work as expected.
+	 *
+	 * @param callback The function to run.
+	 * @returns A function that unsubscribes the callback.
+	 */
+	function effect(callback: () => Cleanup | void): Cleanup;
 
+	/**
+	 * 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.
+	 * @returns The current state.
+	 */
 	function peek<State, Args extends unknown[]>(molecule: State | ((...args: Args) => State), ...args: Args): State;
 
+	/**
+	 * Returns whether the given value is an atom.
+	 *
+	 * @param value The value to check.
+	 * @returns `true` if the value is an atom, otherwise `false`.
+	 */
 	function isAtom(value: unknown): value is Atom<any>;
 
+	/**
+	 * Runs the given function and schedules listeners to be notified only once
+	 * after the function has completed. Useful for batching multiple changes.
+	 *
+	 * @param callback The function to run.
+	 */
 	function batch(callback: () => void): void;
 
+	/**
+	 * Captures all atoms that are read during the function call and returns them
+	 * along with the result of the function. Useful for tracking dependencies.
+	 *
+	 * @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]>;
 
+	/**
+	 * Notifies all subscribers of the given atom that the state has changed.
+	 *
+	 * @param atom The atom to notify.
+	 */
 	function notify(atom: Atom<unknown>): void;
 
+	/**
+	 * Creates an instance of `factory` for each item in the atom's state, and
+	 * cleans up the instance when the item is removed. Returns a cleanup function
+	 * that unsubscribes all instances.
+	 *
+	 * @param molecule The atom or molecule to observe.
+	 * @param factory The function that tracks the lifecycle of each item.
+	 * @returns A function that unsubscribes all instances.
+	 */
 	function observe<Item>(
 		molecule: Molecule<readonly Item[]>,
 		factory: (item: Item, index: number) => Cleanup | void,
@@ -48,6 +151,16 @@ declare namespace Charm {
 		factory: (item: Item, key: Key) => Cleanup | void,
 	): Cleanup;
 
+	/**
+	 * Maps each entry in the atom's state to a new key-value pair. If the `mapper`
+	 * function returns `undefined`, the entry is omitted from the resulting map.
+	 * When the atom changes, the `mapper` is called for each entry in the state
+	 * to compute the new state.
+	 *
+	 * @param molecule The atom or molecule to map.
+	 * @param mapper The function that maps each entry.
+	 * @returns A new atom with the mapped state.
+	 */
 	function mapped<V0, K1, V1>(
 		molecule: Molecule<readonly V0[]>,
 		mapper: (value: V0, index: number) => LuaTuple<[value: V1 | undefined, key: K1]>,
@@ -63,43 +176,110 @@ declare namespace Charm {
 		mapper: (value: V0, key: K0) => LuaTuple<[value: V1 | undefined, key: K1]> | V1,
 	): Molecule<ReadonlyMap<K1, V1>>;
 
-	// React
+	/**
-
+	 * A hook that subscribes to changes in the given atom or molecule. The
+	 * component is re-rendered whenever the state changes.
+	 *
+	 * If the `dependencies` array is provided, the subscription to the atom or
+	 * molecule is re-created whenever the dependencies change. Otherwise, the
+	 * subscription is created once when the component is mounted.
+	 *
+	 * @param molecule The atom or molecule to subscribe to.
+	 * @param dependencies An array of values that the subscription depends on.
+	 * @returns The current state.
+	 */
 	function useAtom<State>(molecule: Molecule<State>, dependencies?: unknown[]): State;
 
-	// Sync
+	/**
-
+	 * Synchronizes state between the client and server. The server sends patches
+	 * to the client, which applies them to its local state.
+	 */
 	const sync: {
+		/**
+		 * Creates a `ClientSyncer` object that receives patches from the server and
+		 * applies them to the local state.
+		 *
+		 * @param options The atoms to synchronize with the server.
+		 * @returns A `ClientSyncer` object.
+		 */
 		client: <T extends Record<string, Atom<any>>>(options: ClientOptions<T>) => ClientSyncer<T>;
+		/**
+		 * Creates a `ServerSyncer` object that sends patches to the client and
+		 * hydrates the client's state.
+		 *
+		 * @param options The atoms to synchronize with the client.
+		 * @returns A `ServerSyncer` object.
+		 */
 		server: <T extends Record<string, Atom<any>>>(options: ServerOptions<T>) => ServerSyncer<T>;
-		collect: (root: Instance) => Record<string, Atom<any>>;
 	};
 
+	/**
+	 * A special value that denotes the absence of a value. Used to represent
+	 * undefined values in patches.
+	 */
 	interface None {
-		__none: "__none";
+		readonly __none: "__none";
 	}
 
 	type SyncPatch<T> = {
 		readonly [P in keyof T]?: (T[P] extends object ? SyncPatch<T[P]> : T[P]) | (T[P] extends undefined ? None : never);
 	};
 
+	/**
+	 * A payload that can be sent from the server to the client to synchronize
+	 * state between the two.
+	 */
 	type SyncPayload<T> = { type: "init"; data: T } | { type: "patch"; data: SyncPatch<T> };
 
 	interface ClientOptions<T extends Record<string, Atom<any>>> {
+		/**
+		 * The atoms to synchronize with the server.
+		 */
 		atoms: T;
 	}
 
 	interface ServerOptions<T extends Record<string, Atom<any>>> {
+		/**
+		 * The atoms to synchronize with the client.
+		 */
 		atoms: T;
+		/**
+		 * The interval at which to send patches to the client, in seconds.
+		 * Defaults to `0` (patches are sent up to once per frame). Set to a
+		 * negative value to disable automatic syncing.
+		 */
 		interval?: number;
 	}
 
 	interface ClientSyncer<T extends Record<string, Atom<any>>> {
+		/**
+		 * Applies a patch or initializes the state of the atoms with the given
+		 * payload from the server.
+		 *
+		 * @param payload The patch or hydration payload to apply.
+		 */
 		sync(payload: SyncPayload<T>): void;
 	}
 
 	interface ServerSyncer<T extends Record<string, Atom<any>>> {
+		/**
+		 * Sets up a subscription to each atom that schedules a patch to be sent to
+		 * the client whenever the state changes. When a change occurs, the `callback`
+		 * is called with the player and the payload to send.
+		 *
+		 * Note that the `payload` object should not be mutated. If you need to
+		 * modify the payload, apply the changes to a copy of the object.
+		 *
+		 * @param callback The function to call when the state changes.
+		 * @returns A cleanup function that unsubscribes all listeners.
+		 */
 		connect(callback: (player: Player, payload: SyncPayload<T>) => void): Cleanup;
+		/**
+		 * Hydrates the client's state with the server's state. This should be
+		 * called when a player joins the game and requires the server's state.
+		 *
+		 * @param player The player to hydrate.
+		 */
 		hydrate(player: Player): void;
 	}
 }

src/init.luau

@@ -5,7 +5,6 @@ local computed = require(script.computed)
 local effect = require(script.effect)
 local mapped = require(script.mapped)
 local observe = require(script.observe)
-local peek = require(script.peek)
 local useAtom = require(script.react.useAtom)
 local store = require(script.store)
 local subscribe = require(script.subscribe)
@@ -22,12 +21,12 @@ return {
 	effect = effect,
 	mapped = mapped,
 	observe = observe,
-	peek = peek,
 	useAtom = useAtom,
 	batch = store.batch,
 	capture = store.capture,
-	notify = store.notify,
 	isAtom = store.isAtom,
+	notify = store.notify,
+	peek = store.peek,
 	subscribe = subscribe,
 	sync = sync,
 }

src/mapped.luau

@@ -1,6 +1,6 @@
 local atom = require(script.Parent.atom)
 local effect = require(script.Parent.effect)
-local peek = require(script.Parent.peek)
+local store = require(script.Parent.store)
 local types = require(script.Parent.types)
 type Atom<T> = types.Atom<T>
 type Molecule<T> = types.Molecule<T>
@@ -10,6 +10,16 @@ type Map =
 	& (<K0, V0, V1>(fn: Molecule<{ [K0]: V0 }>, mapper: (V0, K0) -> V1?) -> Molecule<{ [K0]: V1 }>)
 	& (<K0, V0, K1, V1>(fn: Molecule<{ [K0]: V0 }>, mapper: (V0, K0) -> (V1?, K1?)) -> Molecule<{ [K1]: V1 }>)
 
+--[=[
+	Maps each entry in the atom's state to a new key-value pair. If the `mapper`
+	function returns `undefined`, the entry is omitted from the resulting map.
+	When the atom changes, the `mapper` is called for each entry in the state
+	to compute the new state.
+	
+	@param molecule The atom or molecule to map.
+	@param mapper The function that maps each entry.
+	@return A new atom with the mapped state.
+]=]
 local function mapped<K0, V0, K1, V1>(molecule: Molecule<{ [K0]: V0 }>, mapper: (V0, K0) -> (V1?, K1?)): Molecule<{ [K1]: V1 }>
 	local mappedAtom = atom({})
 	local mappedAtomRef = setmetatable({ current = mappedAtom }, { __mode = "v" })
@@ -22,9 +32,10 @@ local function mapped<K0, V0, K1, V1>(molecule: Molecule<{ [K0]: V0 }>, mapper:
 		end
 
 		local items = molecule()
-		local mappedItems = table.clone(peek(mappedAtomRef.current))
+		local mappedItems = table.clone(store.peek(mappedAtomRef.current))
 		local mappedKeys = {}
 
+		-- TODO: Only call mapper if the item has changed.
 		for key, item in items do
 			local newItem, newKey = mapper(item, key)
 			if newKey == nil then

src/observe.luau

@@ -2,6 +2,15 @@ local effect = require(script.Parent.effect)
 local types = require(script.Parent.types)
 type Molecule<T> = types.Molecule<T>
 
+--[=[
+	Creates an instance of `factory` for each item in the atom's state, and
+	cleans up the instance when the item is removed. Returns a cleanup function
+	that unsubscribes all instances.
+	
+	@param molecule The atom or molecule to observe.
+	@param factory The function that tracks the lifecycle of each item.
+	@return A function that unsubscribes all instances.
+]=]
 local function observe<K, V>(molecule: Molecule<{ [K]: V }>, factory: (value: V, key: K) -> () -> ()): () -> ()
 	local connections: { [K]: () -> () } = {}
 

src/peek.luau

@@ -1,19 +0,0 @@
-local store = require(script.Parent.store)
-
-local function peek<T, U...>(callback: ((U...) -> T) | T, ...: U...): T
-	if type(callback) ~= "function" then
-		return callback
-	end
-
-	if not store.isCapturing() then
-		return callback(...)
-	end
-
-	store.setCapturing(false)
-	local result = callback(...)
-	store.setCapturing(true)
-
-	return result
-end
-
-return peek

src/react/useAtom.luau

@@ -3,6 +3,18 @@ local subscribe = require(script.Parent.Parent.subscribe)
 local types = require(script.Parent.Parent.types)
 type Molecule<T> = types.Molecule<T>
 
+--[=[
+	A hook that subscribes to changes in the given atom or molecule. The
+	component is re-rendered whenever the state changes.
+	
+	If the `dependencies` array is provided, the subscription to the atom or
+	molecule is re-created whenever the dependencies change. Otherwise, the
+	subscription is created once when the component is mounted.
+	
+	@param molecule The atom or molecule to subscribe to.
+	@param dependencies An array of values that the subscription depends on.
+	@return The current state.
+]=]
 local function useAtom<State>(molecule: Molecule<State>, dependencies: { any }?): State
 	local React = requireReact()
 

src/store.luau

@@ -10,10 +10,21 @@ local capturing = false
 local batched: { [() -> ()]: true } = {}
 local batching = false
 
+--[=[
+	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)
 	return not not (value and listeners[value])
 end
 
+--[=[
+	Notifies all subscribers of the given atom that the state has changed.
+	
+	@param atom The atom to notify.
+]=]
 local function notify(atom: Atom<any>)
 	if batching then
 		for listener in next, listeners[atom] do
@@ -27,6 +38,13 @@ local function notify(atom: Atom<any>)
 	end
 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.
+	
+	@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()
@@ -42,9 +60,16 @@ local function capture<T>(callback: () -> T): ({ [Atom<any>]: true }, T)
 	return dependencies, result
 end
 
+--[=[
+	Runs the given function and schedules listeners to be notified only once
+	after the function has completed. Useful for batching multiple changes.
+	
+	@param callback The function to run.
+]=]
 local function batch(callback: () -> ())
 	if batching then
-		return callback()
+		callback()
+		return
 	end
 
 	batching = true
@@ -58,10 +83,45 @@ 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 capturing then
+		return callback(...)
+	end
+
+	capturing = false
+	local result = callback(...)
+	capturing = true
+
+	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.
+]=]
 local function setCapturing(value: boolean)
 	capturing = value
 end
 
+--[=[
+	Returns whether the store is currently capturing atoms.
+
+	@return `true` if the store is capturing atoms, otherwise `false`.
+]=]
 local function isCapturing()
 	return capturing
 end
@@ -73,6 +133,7 @@ return {
 	notify = notify,
 	capture = capture,
 	batch = batch,
+	peek = peek,
 	setCapturing = setCapturing,
 	isCapturing = isCapturing,
 }

src/subscribe.luau

@@ -2,13 +2,24 @@ local store = require(script.Parent.store)
 local types = require(script.Parent.types)
 type Molecule<T> = types.Molecule<T>
 
+--[=[
+	Subscribes to changes in the given atom or molecule. The callback is
+	called with the current state and the previous state immediately after a
+	change occurs.
+	
+	@param molecule The atom or molecule to subscribe to.
+	@param callback The function to call when the state changes.
+	@return A function that unsubscribes the callback.
+]=]
 local function subscribe<T>(molecule: Molecule<T>, callback: (state: T, prev: T) -> ()): () -> ()
 	local captured, state = store.capture(molecule)
 
 	local function listener()
 		local prevState = state
 		state = molecule()
-		callback(state, prevState)
+		if state ~= prevState then
+			callback(state, prevState)
+		end
 	end
 
 	for atom in captured do

src/sync/client.luau

@@ -4,13 +4,29 @@ type SyncPayload = types.SyncPayload
 local patch = require(script.Parent.patch)
 
 type ClientOptions = {
+	--[=[
+		The atoms to synchronize with the server.
+	]=]
 	atoms: { [string]: Atom<any> },
 }
 
 type ClientSyncer = {
+	--[=[
+		Applies a patch or initializes the state of the atoms with the given
+		payload from the server.
+		
+		@param payload The patch or hydration payload to apply.
+	]=]
 	sync: (self: ClientSyncer, payload: SyncPayload) -> (),
 }
 
+--[=[
+	Creates a `ClientSyncer` object that receives patches from the server and
+	applies them to the local state.
+	
+	@param options The atoms to synchronize with the server.
+	@return A `ClientSyncer` object.
+]=]
 local function client(options: ClientOptions): ClientSyncer
 	local atoms = options.atoms
 	local self = {} :: ClientSyncer

src/sync/init.luau

@@ -1,6 +1,10 @@
 local client = require(script.client)
 local server = require(script.server)
 
+--[=[
+	Synchronizes state between the client and server. The server sends patches
+	to the client, which applies them to its local state.
+]=]
 return {
 	client = client,
 	server = server,

src/sync/patch.luau

@@ -1,5 +1,9 @@
 local validate = require(script.Parent.validate)
 
+--[=[
+	A special value that denotes the absence of a value. Used to represent
+	undefined values in patches.
+]=]
 local NONE = { __none = "__none" }
 
 local function diff(prevState: { [any]: any }, nextState: { [any]: any })

src/sync/server.luau

@@ -8,16 +8,51 @@ local setInterval = require(script.Parent.Parent.utils.setInterval)
 local patch = require(script.Parent.patch)
 
 type ServerOptions = {
+	--[=[
+		The atoms to synchronize with the client.
+	]=]
 	atoms: { [string]: Atom<any> },
+	--[=[
+		The interval at which to send patches to the client, in seconds.
+		Defaults to `0` (patches are sent up to once per frame). Set to a
+		negative value to disable automatic syncing.
+	]=]
 	interval: number?,
 }
 
 type ServerSyncer = {
+	--[=[
+		Sets up a subscription to each atom that schedules a patch to be sent to
+		the client whenever the state changes. When a change occurs, the `callback`
+		is called with the player and the payload to send.
+
+		Note that the `payload` object should not be mutated. If you need to
+		modify the payload, apply the changes to a copy of the object.
+
+		@param callback The function to call when the state changes.
+		@return A cleanup function that unsubscribes all listeners.
+	]=]
 	connect: (self: ServerSyncer, callback: (player: Player, payload: SyncPayload) -> ()) -> () -> (),
+	--[=[
+		Hydrates the client's state with the server's state. This should be
+		called when a player joins the game and requires the server's state.
+
+		@param player The player to hydrate.
+	]=]
 	hydrate: (self: ServerSyncer, player: Player) -> (),
+	--[=[
+		@deprecated For internal use only.
+	]=]
 	_sendPatch: (self: ServerSyncer, player: Player) -> (),
 }
 
+--[=[
+	Creates a `ServerSyncer` object that sends patches to the client and
+	hydrates the client's state.
+	
+	@param options The atoms to synchronize with the client.
+	@return A `ServerSyncer` object.
+]=]
 local function server(options: ServerOptions): ServerSyncer
 	local atoms = options.atoms
 	local interval = options.interval or 0
@@ -42,7 +77,7 @@ local function server(options: ServerOptions): ServerSyncer
 			end)
 
 			state[key] = atom()
-			stateSnapshot[key] = atom()
+			stateSnapshot[key] = state[key]
 		end
 
 		local disconnect = setInterval(function()
@@ -51,7 +86,7 @@ local function server(options: ServerOptions): ServerSyncer
 			end
 
 			local diffs = patch.diff(stateSnapshot, state)
-			stateSnapshot = table.clone(state)
+			stateSnapshot = state
 			changed = false
 
 			for _, player in Players:GetPlayers() do

src/sync/validate.luau

@@ -29,6 +29,13 @@ local function isUnsafeTable(object: { [any]: any })
 	return false
 end
 
+--[=[
+	Validates a value to ensure it can be synced over a remote event.
+
+	@param value The value to validate.
+	@param key The key of the value in the table.
+	@error Throws an error if the value cannot be synced.
+]=]
 local function validate(value: any, key: any)
 	local typeOfKey = type(key)
 	local typeOfValue = type(value)

src/types.luau

@@ -1,13 +1,38 @@
 export type WeakMap<K, V> = typeof(setmetatable({} :: { [K]: V }, { __mode = "k" }))
 
+--[=[
+	A primitive state container that can be read from and written to. When the
+	state changes, all subscribers are notified.
+
+	@param state The next state or a function that produces the next state.
+	@return The current state, if no arguments are provided.
+]=]
 export type Atom<T> = (() -> T) & (state: T | (T) -> T) -> T
 
+--[=[
+	A function that depends on one or more atoms and produces a state. Can be
+	used to derive state from atoms.
+
+	@template State The type of the state.
+	@return The current state.
+]=]
 export type Molecule<T> = () -> T
 
+--[=[
+	Optional configuration for creating an atom.
+]=]
 export type AtomOptions<T> = {
+	--[=[
+		A function that determines whether the state has changed. By default,
+		a strict equality check (`===`) is used.
+	]=]
 	equals: (prev: T, next: T) -> boolean,
 }
 
+--[=[
+	A payload that can be sent from the server to the client to synchronize
+	state between the two.
+]=]
 export type SyncPayload = {
 	type: "patch" | "init",
 	data: { [any]: any },

src/utils/collect.luau

@@ -1,3 +1,7 @@
+--[=[
+	Trigger a garbage collection cycle by allocating a large amount of junk memory
+	and then waiting for the next frame.
+]=]
 local function collect()
 	for _ = 1, 1e4 do
 		local _ = table.create(1e3)

src/utils/count.luau

@@ -1,3 +1,9 @@
+--[=[
+	Counts the number of elements in a table.
+
+	@param object The table to count.
+	@return The number of elements in the table.
+]=]
 local function count(object: any): number
 	local count = 0
 	for _ in next, object do

src/utils/setInterval.luau

@@ -1,5 +1,12 @@
 local RunService = game:GetService("RunService")
 
+--[=[
+	Schedules a function to be called at a fixed interval of `seconds`.
+	
+	@param callback The function to call.
+	@param seconds The interval in seconds.
+	@return A function that cancels the interval.
+]=]
 local function setInterval(callback: () -> (), seconds: number): () -> ()
 	if seconds < 0 then
 		return function() end