refactor: move peek function to store module

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/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/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,
 }