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

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(peek(function()
+		expect(store.peek(1)).to.equal(1)
+		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): () -> ()
-	local captured = store.capture(molecule)
+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 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