refactor: code quality

src/__tests__/batch.spec.luau

@@ -1,6 +1,6 @@
 local atom = require(script.Parent.Parent.atom)
-local batch = require(script.Parent.Parent.batch)
 local effect = require(script.Parent.Parent.effect)
+local store = require(script.Parent.Parent.store)
 local subscribe = require(script.Parent.Parent.subscribe)
 
 return function()
@@ -10,7 +10,7 @@ return function()
 		subscribe(source, function()
 			value += 1
 		end)
-		batch(function()
+		store.batch(function()
 			source(2)
 			source(3)
 			source(4)
@@ -28,7 +28,7 @@ return function()
 			b()
 		end)
 		value = 0 -- effect() will run once initially
-		batch(function()
+		store.batch(function()
 			a(2)
 			b(2)
 			a(3)

src/__tests__/capture.spec.luau

@@ -1,10 +1,10 @@
 local atom = require(script.Parent.Parent.atom)
-local capture = require(script.Parent.Parent.capture)
+local store = require(script.Parent.Parent.store)
 
 return function()
 	it("captures atoms from a function", function()
 		local a, b, c = atom(1), atom(1), atom(2)
-		local captured, result = capture(function()
+		local captured, result = store.capture(function()
 			return a() + b()
 		end)
 		expect(captured[a]).to.be.ok()
@@ -15,7 +15,7 @@ return function()
 
 	it("captures an atom by itself", function()
 		local a = atom(1)
-		local captured, result = capture(a)
+		local captured, result = store.capture(a)
 		expect(captured[a]).to.be.ok()
 		expect(result).to.equal(1)
 	end)

src/__tests__/derive.spec.luau

@@ -42,7 +42,7 @@ return function()
 
 		local function listenersDecreased()
 			local count = 0
-			for _ in pairs(store.listeners[source]) do
+			for _ in next, store.listeners[source] do
 				count += 1
 			end
 			max = math.max(max, count)

src/__tests__/sync/validate.spec.luau

@@ -2,7 +2,7 @@ local validate = require(script.Parent.Parent.Parent.sync.validate)
 
 return function()
 	it("validates inputs", function()
-		local cases = {
+		local cases: { [string]: { value: any, expected: any } } = {
 			array = { value = { 1, 2, 3 }, expected = true },
 			array_sparse = { value = { 1, nil :: any, 3 }, expected = true },
 			array_sparse_invalid = { value = { [1] = 0, [3] = 0 }, expected = false },
@@ -15,7 +15,7 @@ return function()
 			number_nan = { value = 0 / 0, expected = false },
 		}
 
-		for key, case in pairs(cases) do
+		for key, case in cases do
 			local success, result = pcall(function()
 				validate(case.value, key)
 			end)

src/__tests__/unwrap.spec.luau

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

src/atom.luau

@@ -21,14 +21,7 @@ local function atom<T>(state: T, options: AtomOptions<T>?): Atom<T>
 
 		if state ~= nextState and not (equals and equals(state, nextState)) then
 			state = nextState
-
+			store.notify(atom)
-			if store.isBatching() then
-				store.batched[atom] = true
-			else
-				for listener in pairs(table.clone(store.listeners[atom])) do
-					listener()
-				end
-			end
 		end
 
 		return state

src/batch.luau

@@ -1,26 +0,0 @@
-local store = require(script.Parent.store)
-
-local function batch(callback: () -> ())
-	if store.isBatching() then
-		return callback()
-	end
-
-	local visited: { [() -> ()]: true } = {}
-
-	store.setBatching(true)
-	callback()
-	store.setBatching(false)
-
-	for atom in store.batched do
-		for listener in pairs(store.listeners[atom]) do
-			if not visited[listener] then
-				visited[listener] = true
-				listener()
-			end
-		end
-	end
-
-	table.clear(store.batched)
-end
-
-return batch

src/capture.luau

@@ -1,20 +0,0 @@
-local store = require(script.Parent.store)
-local types = require(script.Parent.types)
-type Atom<T> = types.Atom<T>
-
-local function capture<T>(callback: () -> T): ({ [Atom<any>]: true }, T)
-	if store.listeners[callback] then
-		return { [callback] = true }, callback()
-	end
-
-	store.setCapturing(true)
-	local result = callback()
-	store.setCapturing(false)
-
-	local dependencies = table.clone(store.captured)
-	table.clear(store.captured)
-
-	return dependencies, result
-end
-
-return capture

src/derive.luau

@@ -1,5 +1,4 @@
 local atom = require(script.Parent.atom)
-local capture = require(script.Parent.capture)
 local store = require(script.Parent.store)
 local types = require(script.Parent.types)
 type Atom<T> = types.Atom<T>
@@ -7,7 +6,7 @@ type AtomOptions<T> = types.AtomOptions<T>
 type Molecule<T> = types.Molecule<T>
 
 local function derive<T>(molecule: Molecule<T>, options: AtomOptions<T>?): Molecule<T>
-	local atoms, state = capture(molecule)
+	local atoms, state = store.capture(molecule)
 	local derived = atom(state, options)
 	local derivedRef = setmetatable({ current = derived }, { __mode = "v" })
 

src/effect.luau

@@ -1,8 +1,7 @@
-local capture = require(script.Parent.capture)
 local store = require(script.Parent.store)
 
 local function effect(molecule: () -> nil): () -> ()
-	local atoms = capture(molecule)
+	local atoms = store.capture(molecule)
 
 	local function listener()
 		molecule()

src/index.d.ts

@@ -22,8 +22,6 @@ declare namespace Charm {
 
 	function atom<State>(state: State, options?: AtomOptions<State>): Atom<State>;
 
-	function isAtom(value: unknown): value is Atom<any>;
-
 	function derive<State>(molecule: Molecule<State>, options?: AtomOptions<State>): Molecule<State>;
 
 	function subscribe<State>(molecule: Molecule<State>, callback: (state: State, prev: State) => void): Cleanup;
@@ -32,10 +30,14 @@ declare namespace Charm {
 
 	function unwrap<State, Args extends unknown[]>(molecule: State | ((...args: Args) => State), ...args: Args): State;
 
-	function capture<State>(molecule: Molecule<State>): LuaTuple<[captured: Set<Atom<unknown>>, state: State]>;
+	function isAtom(value: unknown): value is Atom<any>;
 
 	function batch(callback: () => void): void;
 
+	function capture<State>(molecule: Molecule<State>): LuaTuple<[captured: Set<Atom<unknown>>, state: State]>;
+
+	function notify(atom: Atom<unknown>): void;
+
 	function observe<Item>(
 		molecule: Molecule<readonly Item[]>,
 		factory: (item: Item, index: number) => Cleanup | void,

src/init.luau

@@ -1,16 +1,14 @@
 _G[script.modules.Promise] = _G[script]
 
 local atom = require(script.atom)
-local batch = require(script.batch)
-local capture = require(script.capture)
 local derive = require(script.derive)
 local effect = require(script.effect)
-local isAtom = require(script.isAtom)
 local map = require(script.map)
 local observe = require(script.observe)
 local useAtom = require(script.react.useAtom)
 local useAtomState = require(script.react.useAtomState)
 local useSetAtom = require(script.react.useSetAtom)
+local store = require(script.store)
 local subscribe = require(script.subscribe)
 local sync = require(script.sync)
 local types = require(script.types)
@@ -22,11 +20,12 @@ export type SyncPayload = types.SyncPayload
 
 return {
 	atom = atom,
-	batch = batch,
+	batch = store.batch,
-	capture = capture,
+	capture = store.capture,
+	notify = store.notify,
+	isAtom = store.isAtom,
 	derive = derive,
 	effect = effect,
-	isAtom = isAtom,
 	map = map,
 	observe = observe,
 	useAtom = useAtom,

src/isAtom.luau

@@ -1,7 +0,0 @@
-local store = require(script.Parent.store)
-
-local function isAtom(value: any)
-	return value ~= nil and store.listeners[value] :: any ~= nil
-end
-
-return isAtom

src/store.luau

@@ -3,11 +3,61 @@ type Atom<T> = types.Atom<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 batched: { [Atom<any>]: true } = {}
 local capturing = false
+
+local batched: { [() -> ()]: true } = {}
 local batching = false
 
+local function isAtom(value: any)
+	return not not (value and listeners[value])
+end
+
+local function notify(atom: Atom<any>)
+	if batching then
+		for listener in next, listeners[atom] do
+			batched[listener] = true
+		end
+		return
+	end
+
+	for listener in next, table.clone(listeners[atom]) do
+		listener()
+	end
+end
+
+local function capture<T>(callback: () -> T): ({ [Atom<any>]: true }, T)
+	if listeners[callback] then
+		return { [callback] = true }, callback()
+	end
+
+	capturing = true
+	local result = callback()
+	capturing = false
+
+	local dependencies = table.clone(captured)
+	table.clear(captured)
+
+	return dependencies, result
+end
+
+local function batch(callback: () -> ())
+	if batching then
+		return callback()
+	end
+
+	batching = true
+	callback()
+	batching = false
+
+	for listener in batched do
+		listener()
+	end
+
+	table.clear(batched)
+end
+
 local function setCapturing(value: boolean)
 	capturing = value
 end
@@ -16,20 +66,13 @@ local function isCapturing()
 	return capturing
 end
 
-local function setBatching(value: boolean)
-	batching = value
-end
-
-local function isBatching()
-	return batching
-end
-
 return {
 	listeners = listeners,
 	captured = captured,
-	batched = batched,
+	isAtom = isAtom,
+	notify = notify,
+	capture = capture,
+	batch = batch,
 	setCapturing = setCapturing,
 	isCapturing = isCapturing,
-	setBatching = setBatching,
-	isBatching = isBatching,
 }

src/subscribe.luau

@@ -1,10 +1,9 @@
-local capture = require(script.Parent.capture)
 local store = require(script.Parent.store)
 local types = require(script.Parent.types)
 type Molecule<T> = types.Molecule<T>
 
 local function subscribe<T>(molecule: Molecule<T>, callback: (state: T, prev: T) -> ()): () -> ()
-	local atoms, state = capture(molecule)
+	local atoms, state = store.capture(molecule)
 
 	local function listener()
 		local prevState = state

src/sync/validate.luau

@@ -1,4 +1,4 @@
-local isAtom = require(script.Parent.Parent.isAtom)
+local store = require(script.Parent.Parent.store)
 
 local SAFE_KEYS = { string = true, number = true }
 local SAFE_VALUES = { string = true, number = true, boolean = true, table = true }
@@ -40,7 +40,7 @@ local function validate(value: any, key: any)
 	elseif typeOfValue == "table" then
 		if getmetatable(value) ~= nil then
 			error(`Cannot sync tables with metatables! Got {value} at key '{key}'`)
-		elseif isAtom(value) then
+		elseif store.isAtom(value) then
 			error(`Cannot sync nested atoms! Got atom '{value.debugLabel}' at key '{key}'`)
 		elseif isUnsafeTable(value) then
 			error(