refactor: improve logic

pkgs/charm/src/atom.luau

@@ -10,6 +10,14 @@ local signal = signals.signal
 
 type Equals<T> = (current: T, incoming: T) -> boolean
 
+local function createAtom<T>(getter: Getter<T>, setter: Setter<T>): Atom<T>
+	local function atom(...: Update<T>): T
+		return if select("#", ...) == 0 then getter() else setter((...))
+	end
+
+	return atom
+end
+
 local function createAtomWithEqualityCheck<T>(getter: Getter<T>, setter: Setter<T>, equals: Equals<T>): Atom<T>
 	local function atom(...: Update<T>): T
 		if select("#", ...) == 0 then
@@ -27,18 +35,10 @@ local function createAtomWithEqualityCheck<T>(getter: Getter<T>, setter: Setter<
 	return atom
 end
 
-local function createAtom<T>(getter: Getter<T>, setter: Setter<T>): Atom<T>
-	local function atom(...: Update<T>): T
-		return if select("#", ...) == 0 then getter() else setter((...))
-	end
-
-	return atom
-end
-
 local function atom<T>(state: T, equals: Equals<T>?): Atom<T>
 	local getter, setter = signal(state)
 
-	if equals == nil then
+	if not equals then
 		return createAtom(getter, setter)
 	else
 		return createAtomWithEqualityCheck(getter, setter, equals)

pkgs/charm/src/batched.luau

@@ -1,16 +1,22 @@
 local signals = require("./signals")
 
+type ReactiveNode = signals.ReactiveNode
+
 local startBatch = signals.startBatch
 local endBatch = signals.endBatch
 
-local function batched(callback: () -> ())
-	startBatch()
-	local success, result: any = pcall(callback)
+local function handleBatched(ok: boolean, ...)
 	endBatch()
-
-	if not success then
-		error(result)
+	if not ok then
+		return ...
+	else
+		error(...)
 	end
 end
 
+local function batched<Args..., Results...>(callback: (Args...) -> Results..., ...: Args...): Results...
+	startBatch()
+	return handleBatched(pcall(callback, ...))
+end
+
 return batched

pkgs/charm/src/lib.luau

@@ -4,6 +4,7 @@ local effect = require("./effect")
 local listen = require("./listen")
 local mapped = require("./mapped")
 local observe = require("./observe")
+local onCleanup = require("./onCleanup")
 local signals = require("./signals")
 local subscribe = require("./subscribe")
 local types = require("./types")
@@ -17,8 +18,6 @@ export type Update<T = any> = types.Update<T>
 
 local computed = signals.computed
 local effectScope = signals.effectScope
-local onEffectCleanup = signals.onEffectCleanup
-local onScopeDispose = signals.onScopeDispose
 
 return {
 	atom = atom,
@@ -31,7 +30,6 @@ return {
 	observe = observe,
 	subscribe = subscribe,
 	untracked = untracked,
-	onEffectCleanup = onEffectCleanup,
-	onScopeDispose = onScopeDispose,
+	onCleanup = onCleanup,
 	signals = signals,
 }

pkgs/charm/src/listen.luau

@@ -2,6 +2,8 @@ local signals = require("./signals")
 
 local computed = signals.computed
 local effect = signals.effect
+local effectScope = signals.effectScope
+local onEffectCleanup = signals.onEffectCleanup
 
 local function listen<T>(selector: () -> T, listener: (state: T, previous: T?, dispose: () -> ()) -> ()): () -> ()
 	local tracked = computed(function()
@@ -18,7 +20,12 @@ local function listen<T>(selector: () -> T, listener: (state: T, previous: T?, d
 	dispose = effect(function()
 		local previous = current
 		current = tracked()
-		listener(current, previous, dispose)
+
+		local stopScope = effectScope(function()
+			listener(current, previous, dispose)
+		end)
+
+		onEffectCleanup(stopScope)
 	end)
 
 	if disposed then

pkgs/charm/src/noYield.luau

@@ -1,11 +1,6 @@
 local function handleNoYield(thread: thread, ok: boolean, ...)
 	if not ok then
-		local err = (...)
-		if typeof(err) == "string" then
-			return false, debug.traceback(thread, err)
-		else
-			return false, tostring(err)
-		end
+		return false, debug.traceback(thread, tostring((...)))
 	end
 
 	if coroutine.status(thread) ~= "dead" then
@@ -20,6 +15,6 @@ local function noYield<Args..., Results...>(callback: (Args...) -> Results..., .
 	return handleNoYield(thread, coroutine.resume(thread, ...))
 end
 
--- To eliminate the overhead of noYield returning callback(...) in debug mode,
--- replace it with pcall to avoid wrapping the callback in another function.
+-- To eliminate the overhead of noYield returning callback(...) in production,
+-- replace it with pcall to simulate the same behavior.
 return if _G.__DEV__ then noYield else pcall :: never

pkgs/charm/src/observe.luau

@@ -1,10 +1,9 @@
+local onCleanup = require("./onCleanup")
 local signals = require("./signals")
 
 local effect = signals.effect
 local effectScope = signals.effectScope
-local onEffectCleanup = signals.onEffectCleanup
 local onScopeDispose = signals.onScopeDispose
-local getCurrentSub = signals.getCurrentSub
 
 local function observe<K, V>(selector: () -> { [K]: V }, observer: (item: V, key: K) -> ...(() -> ())?): () -> ()
 	local scopes: { [K]: () -> () } = {}
@@ -46,11 +45,7 @@ local function observe<K, V>(selector: () -> { [K]: V }, observer: (item: V, key
 		end
 	end)
 
-	if getCurrentSub() then
-		onEffectCleanup(disposeScopes)
-	else
-		onScopeDispose(disposeScopes)
-	end
+	onCleanup(disposeScopes, true)
 
 	return function()
 		disposeEffect()

pkgs/charm/src/onCleanup.luau

@@ -0,0 +1,20 @@
+local signals = require("./signals")
+
+local onEffectCleanup = signals.onEffectCleanup
+local onScopeDispose = signals.onScopeDispose
+local getCurrentScope = signals.getCurrentScope
+local getCurrentSub = signals.getCurrentSub
+
+local warn = warn or error :: never
+
+local function onCleanup(cleanup: () -> (), failSilently: boolean?)
+	if getCurrentSub() then
+		onEffectCleanup(cleanup)
+	elseif getCurrentScope() then
+		onScopeDispose(cleanup)
+	elseif _G.__DEV__ and not failSilently then
+		warn(debug.traceback("onCleanup() can only be called inside an effect or a scope.", 2))
+	end
+end
+
+return onCleanup

pkgs/charm/src/subscribe.luau

@@ -2,21 +2,31 @@ local signals = require("./signals")
 
 local computed = signals.computed
 local effect = signals.effect
+local effectScope = signals.effectScope
+local onEffectCleanup = signals.onEffectCleanup
 
 local function subscribe<T>(selector: () -> T, listener: (state: T, lastState: T) -> ()): () -> ()
 	local tracked = computed(function()
 		return selector()
 	end)
 
-	local current = tracked()
+	local current: T
+	local mounted = true
 
 	return effect(function()
 		local previous = current
 		current = tracked()
 
-		if current ~= previous then
-			listener(current, previous)
+		if mounted then
+			mounted = false
+			return
 		end
+
+		local stopScope = effectScope(function()
+			listener(current, previous)
+		end)
+
+		onEffectCleanup(stopScope)
 	end)
 end
 

pkgs/charm/src/types.luau

@@ -1,9 +1,11 @@
 export type Getter<T = any> = () -> T
+
 export type Setter<T = any> = (update: Update<T>) -> T
+
 export type Update<T = any> = T | (previous: T) -> T
 
 export type Selector<T = any> = () -> T
 
-export type Atom<T = any> = Setter<T> & Getter<T>
+export type Atom<T = any> = (...Update<T>) -> T
 
 return {}

pkgs/charm/src/untracked.luau

@@ -1,18 +1,21 @@
 local noYield = require("./noYield")
 local signals = require("./signals")
 
+type ReactiveNode = signals.ReactiveNode
+
 local setCurrentSub = signals.setCurrentSub
 
-local function untracked<T, Args...>(callback: ((Args...) -> T) | (Args...) -> (), ...: Args...): T
-	local currentSub = setCurrentSub(nil)
-	local success, result = noYield(callback, ...)
+local function handleUntracked(currentSub: ReactiveNode?, ok: boolean, ...)
 	setCurrentSub(currentSub)
-
-	if not success then
-		error(result)
+	if ok then
+		return ...
+	else
+		error(...)
 	end
+end
 
-	return result
+local function untracked<Args..., Results...>(callback: (Args...) -> Results..., ...: Args...): Results...
+	return handleUntracked(setCurrentSub(nil), noYield(callback, ...))
 end
 
 return untracked

pkgs/charm/tests/listen.spec.luau

@@ -3,6 +3,7 @@ local suite = require("@test/suite")
 
 local atom = Charm.atom
 local listen = Charm.listen
+local effect = Charm.effect
 local test = suite.test
 
 test("should listen to selector", function()
@@ -26,6 +27,47 @@ test("should listen to selector", function()
 	assert(current == 2 and previous == 1, "should not run after unsub")
 end)
 
+test("should run in untracked", function()
+	local runs = 0
+	local cleanups = 0
+	local a = atom(0)
+	local b = atom(0)
+
+	listen(a, function()
+		runs += 1
+		b()
+		effect(function()
+			return function()
+				cleanups += 1
+			end
+		end)
+	end)
+
+	assert(runs == 1, "should run once on init")
+	assert(cleanups == 0, "should not cleanup yet")
+
+	a(1)
+	assert(runs == 2, "should run on a change")
+	assert(cleanups == 1, "should cleanup first call")
+
+	b(1)
+	assert(runs == 2, "should not run on b change")
+	assert(cleanups == 1, "should not cleanup second call")
+end)
+
+test("should run if first value is nil", function()
+	local runs = 0
+	local source = atom(nil :: number?)
+
+	listen(source, function()
+		runs += 1
+	end)
+
+	assert(runs == 1, "should run once on init")
+	source(1)
+	assert(runs == 2, "should run once on change")
+end)
+
 test("should track inner subs", function()
 	local order = {}
 	local a = atom(1)

pkgs/charm/tests/signals/effect.spec.luau

@@ -326,4 +326,34 @@ test("should run inner effect cleanup", function()
 	assert(cleanups == 4, "should not cleanup after update")
 end)
 
+test("should run cleanups outside any evaluation context", function()
+	local calls = 0
+	local a, setA = signal(0)
+	local b, setB = signal(0)
+	local c = computed(function()
+		if a() == 0 then
+			effect(function()
+				onEffectCleanup(function()
+					b()
+				end)
+			end)
+		end
+		return a()
+	end)
+
+	effect(function()
+		calls += 1
+		c()
+	end)
+	assert(calls == 1, "should run once")
+
+	calls = 0
+	setA(1)
+	assert(calls == 1, "should run once after a update")
+
+	calls = 0
+	setB(1)
+	assert(calls == 0, "should not run after b update")
+end)
+
 return {}

pkgs/charm/tests/spec.luau

@@ -15,6 +15,9 @@ suite.describe("charm", function()
 	suite.describe("subscribe", function()
 		require("./subscribe.spec")
 	end)
+	suite.describe("untracked", function()
+		require("./untracked.spec")
+	end)
 end)
 
 return {}

pkgs/charm/tests/subscribe.spec.luau

@@ -7,6 +7,7 @@ local atom = Charm.atom
 local computed = Charm.computed
 local subscribe = Charm.subscribe
 local listen = Charm.listen
+local effect = Charm.effect
 local test = suite.test
 
 test("should subscribe to selector", function()
@@ -30,6 +31,38 @@ test("should subscribe to selector", function()
 	assert(current == 2 and previous == 1, "should not run after unsub")
 end)
 
+test("should run in untracked", function()
+	local runs = 0
+	local cleanups = 0
+	local a = atom(0)
+	local b = atom(0)
+
+	subscribe(a, function()
+		runs += 1
+		b()
+		effect(function()
+			return function()
+				cleanups += 1
+			end
+		end)
+	end)
+
+	assert(runs == 0, "should not run once on init")
+	assert(cleanups == 0, "should not cleanup yet")
+
+	a(1)
+	assert(runs == 1, "should run once on change")
+	assert(cleanups == 0, "should not cleanup yet")
+
+	b(1)
+	assert(runs == 1, "should not run on b change")
+	assert(cleanups == 0, "should not cleanup first call")
+
+	a(2)
+	assert(runs == 2, "should run on a change")
+	assert(cleanups == 1, "should cleanup first call")
+end)
+
 test("should ensure correct execution order", function()
 	local order = {}
 	local n1 = atom(0)

pkgs/charm/tests/untracked.spec.luau

@@ -0,0 +1,86 @@
+local Charm = require("../src/lib")
+local suite = require("@test/suite")
+
+local atom = Charm.atom
+local computed = Charm.computed
+local effect = Charm.effect
+local untracked = Charm.untracked
+local test = suite.test
+
+test("should block tracking inside effects", function()
+	local a = atom(1)
+	local b = atom(2)
+	local spyCalls = 0
+
+	effect(function()
+		untracked(function()
+			spyCalls += 1
+			return a() + b()
+		end)
+	end)
+	assert(spyCalls == 1, "should be called once")
+
+	a(10)
+	b(20)
+	assert(spyCalls == 1, "should be called once")
+end)
+
+test("should block tracking even when run inside effect run inside untracked", function()
+	local s = atom(1)
+	local spyCalls = 0
+
+	untracked(function()
+		effect(function()
+			untracked(function()
+				spyCalls += 1
+				return s()
+			end)
+		end)
+	end)
+	assert(spyCalls == 1, "should be called once")
+
+	s(2)
+	assert(spyCalls == 1, "should be called once")
+end)
+
+test("should not cause atom assignments throw", function()
+	local a = atom(1)
+	local aChangedTime = atom(0)
+
+	local dispose = effect(function()
+		a()
+		untracked(function()
+			aChangedTime(aChangedTime() + 1)
+		end)
+	end)
+
+	a(2)
+	assert(aChangedTime() == 2, "aChangedTime should be 2")
+	a(3)
+	assert(aChangedTime() == 3, "aChangedTime should be 3")
+
+	dispose()
+end)
+
+test("should block tracking inside computed atoms", function()
+	local a = atom(1)
+	local b = atom(2)
+	local spyCalls = 0
+	local c = computed(function()
+		return untracked(function()
+			spyCalls += 1
+			return a() + b()
+		end)
+	end)
+
+	assert(spyCalls == 0, "should not be called yet")
+	assert(c() == 3, "c should be 3")
+	a(10)
+	c()
+	b(20)
+	c()
+	assert(spyCalls == 1, "should be called once")
+	assert(c() == 3, "c should be 3")
+end)
+
+return {}