perf: optimize subscribe and mapped

pkgs/charm/benchmarks/propagate.luau

@@ -0,0 +1,74 @@
+local Charm = require("@pkgs/charm/src/lib")
+
+local atom = Charm.atom
+local computed = Charm.computed
+local effect = Charm.effect
+local listen = Charm.listen
+
+local signal = Charm.signals.signal
+local computedSignal = Charm.signals.computed
+local effectSignal = Charm.signals.effect
+
+local w = 10
+local h = 10
+
+local function bench(name: string, test: () -> ())
+	local start = os.clock()
+	local iter = 0
+	while os.clock() - start < 2 do
+		iter += 1
+		test()
+	end
+	local endTime = os.clock()
+	local elapsed = endTime - start
+	print(`{name}: {elapsed / iter * 1000 * 1000} us/iter`)
+end
+
+bench(`listen {w}*{h}`, function()
+	local src = atom(1)
+	for _ = 1, w do
+		local last = src
+		for _ = 1, h do
+			local prev = last
+			last = computed(function()
+				return prev() + 1
+			end)
+		end
+		listen(last, function() end)
+	end
+	src(src() + 1)
+end)
+
+bench(`effect {w}*{h}`, function()
+	local src = atom(1)
+	for _ = 1, w do
+		local last = src
+		for _ = 1, h do
+			local prev = last
+			last = computed(function()
+				return prev() + 1
+			end)
+		end
+		effect(function()
+			last()
+		end)
+	end
+	src(src() + 1)
+end)
+
+bench(`signals {w}*{h}`, function()
+	local src, setSrc = signal(1)
+	for _ = 1, w do
+		local last = src
+		for _ = 1, h do
+			local prev = last
+			last = computedSignal(function()
+				return prev() + 1
+			end)
+		end
+		effectSignal(function()
+			last()
+		end)
+	end
+	setSrc(src() + 1)
+end)

pkgs/charm/src/atom.luau

@@ -2,7 +2,6 @@ local signals = require("./signals")
 local types = require("./types")
 
 type Atom<T> = types.Atom<T>
-type AtomOptions<T> = types.AtomOptions<T>
 
 local signal = signals.signal
 

pkgs/charm/src/effect.luau

@@ -1,16 +1,17 @@
 local signals = require("./signals")
 
-local _effect = signals.effect
 local onEffectCleanup = signals.onEffectCleanup
 
-local function effect(callback: (dispose: () -> ()) -> ...() -> ()?): () -> ()
+type Cleanup = () -> ()
+
+local function effect(callback: (dispose: () -> ()) -> ...Cleanup?): () -> ()
 	local disposeImmediately = false
 
 	local function dispose()
 		disposeImmediately = true
 	end
 
-	dispose = _effect(function()
+	dispose = signals.effect(function()
 		local cleanup = callback(dispose)
 		if cleanup then
 			onEffectCleanup(cleanup)

pkgs/charm/src/effectScope.luau

@@ -1,10 +1,11 @@
 local signals = require("./signals")
 
-local _effectScope = signals.effectScope
 local onScopeDispose = signals.onScopeDispose
 
-local function effectScope(callback: () -> ...() -> ()?, detached: boolean?): () -> ()
-	return _effectScope(function()
+type Cleanup = () -> ()
+
+local function effectScope(callback: () -> ...Cleanup?, detached: boolean?): () -> ()
+	return signals.effectScope(function()
 		local cleanup = callback()
 		if cleanup then
 			onScopeDispose(cleanup)

pkgs/charm/src/init.luau

@@ -1,7 +1,6 @@
 local lib = require(script.lib)
 
 export type Atom<T = any> = lib.Atom<T>
-export type AtomOptions<T = any> = lib.AtomOptions<T>
 export type Selector<T = any> = lib.Selector<T>
 
 return lib

pkgs/charm/src/lib.luau

@@ -10,7 +10,6 @@ local types = require("./types")
 local untracked = require("./untracked")
 
 export type Atom<T = any> = types.Atom<T>
-export type AtomOptions<T = any> = types.AtomOptions<T>
 export type Selector<T = any> = types.Selector<T>
 
 local computed = signals.computed

pkgs/charm/src/listen.luau

@@ -1,5 +1,4 @@
 local signals = require("./signals")
-local spawn = require("./spawn")
 
 local effect = signals.effect
 local pauseTracking = signals.pauseTracking
@@ -7,22 +6,24 @@ local resumeTracking = signals.resumeTracking
 
 local function listen<T>(selector: () -> T, listener: (state: T, previous: T?, dispose: () -> ()) -> ()): () -> ()
 	local disposeImmediately = false
-	local current: T?
+	local current: T
 
 	local function dispose()
 		disposeImmediately = true
 	end
 
 	dispose = effect(function()
-		local value = selector()
-
-		if value ~= current then
-			local previous = current
-			current = value
+		local previous: T? = current
+		current = selector()
 
+		if current ~= previous then
 			pauseTracking()
-			spawn(listener, value, previous, dispose)
+			local success, result: unknown = pcall(listener, current, previous, dispose)
 			resumeTracking()
+
+			if not success then
+				error(`listen() encountered an error: {result}`)
+			end
 		end
 	end)
 

pkgs/charm/src/mapped.luau

@@ -4,128 +4,86 @@ local computed = signals.computed
 local pauseTracking = signals.pauseTracking
 local resumeTracking = signals.resumeTracking
 
-type MappedKey<VI, KO, VO> = {
-	key: KO,
-	value: VO,
-}
+type Mapped =
+	& (<KI, VI, KO, VO>(selector: () -> { [KI]: VI }, mapper: (item: VI, key: KI) -> (VO, KO)) -> () -> { [KO]: VO })
+	& (<K, VI, VO>(selector: () -> { [K]: VI }, mapper: (item: VI, key: K) -> VO) -> () -> { [K]: VO })
 
-local function mapped<KI, VI, KO, VO>(selector: () -> { [KI]: VI }, transform: (item: VI, key: KI) -> (VO, KO?)): () -> { [KO]: VO }
-	local mappedKeys: { [KI]: MappedKey<VI, KO, VO> } = {}
-	local previousItems: { [KI]: VI }?
+local function ensureCanMutateOutput<KO, VO>(output: { [KO]: VO }, previous: { [KO]: VO }?): { [KO]: VO }
+	return if output == previous then table.clone(output) else output
+end
 
-	local function getTransform(value: VI, key: KI): (VO, KO?)
+local function mapped<KI, VI, KO, VO>(selector: () -> { [KI]: VI }, mapper: (item: VI, key: KI) -> (VO, KO)): () -> { [KO]: VO }
+	local currentData: { [KI]: VI } = {}
+	local keyCache: { [KI]: KO } = {}
+
+	local function map(value: VI, key: KI): (VO, KO)
 		pauseTracking()
-		local success, valueOut, keyOut = pcall(transform, value, key)
+		local success, valueOut, keyOut = pcall(mapper, value, key)
 		resumeTracking()
 
 		if not success then
 			error(`mapped() transform function failed: {valueOut}`)
 		end
 
-		return valueOut, keyOut
+		return valueOut, if keyOut ~= nil then keyOut else key :: any
 	end
 
-	return computed(function(previousMappedItems)
-		local currentItems = selector()
+	return computed(function(previousOutput)
+		local output: { [KO]: VO } = previousOutput or {}
+		local previousData = currentData
+		currentData = selector()
 
-		-- If this is the initial run, create a map of the items and store the
-		-- updated keys and values.
-		if not previousItems or not previousMappedItems then
-			local mappedItems: { [KO]: VO } = {}
-			previousItems = currentItems
-
-			for key, value in currentItems do
-				local valueOut, keyOut = getTransform(value, key)
-				local keyOut = if keyOut ~= nil then keyOut else key :: never
-
-				if valueOut ~= nil then
-					mappedKeys[key] = { key = keyOut, value = valueOut }
-					mappedItems[keyOut] = valueOut
-				end
-			end
-
-			return mappedItems
+		-- Check equality since the selector might not be a computed value
+		if currentData == previousData then
+			return output
 		end
 
-		if previousItems == currentItems then
-			return previousMappedItems
+		-- Clear removed values from the output and key cache
+		for key in previousData do
+			if currentData[key] == nil then
+				local keyOut = keyCache[key]
+
+				if keyOut then
+					output = ensureCanMutateOutput(output, previousOutput)
+					output[keyOut] = nil
+					keyCache[key] = nil
+				end
+			end
 		end
 
-		local mappedItems = table.clone(previousMappedItems)
-
-		-- Update and remove mapped keys if the source values changed
-		for key, oldValue in previousItems do
-			local currentValue = currentItems[key]
-			local mappedKey = mappedKeys[key]
-
-			-- Remove the key if the source value changed from non-nil to nil
-			if currentValue == nil then
-				mappedKeys[key] = nil
-				if mappedKey then
-					mappedItems[mappedKey.key] = nil
-				end
+		-- Process new or changed values
+		for key, value in currentData do
+			if value == previousData[key] then
 				continue
 			end
 
-			-- Do not call transform if the value has not changed
-			if currentValue == oldValue then
-				continue
-			end
-
-			local valueOut, keyOut = getTransform(currentValue, key)
-			local keyOut = if keyOut ~= nil then keyOut else key :: never
-
-			-- If there is no previous mappedKey, then this item was filtered
-			-- out to optimize memory. We need to check if the value is nil to
-			-- determine if we should add it back to the list.
-			if not mappedKey then
-				if valueOut ~= nil then
-					mappedKeys[key] = { key = keyOut, value = valueOut }
-					mappedItems[keyOut] = valueOut
-				end
-				continue
-			end
-
-			-- If the mapped value is nil, then delete mappedKey to save memory.
-			-- This is treated the same as setting mappedKey.value to nil.
-			if valueOut == nil then
-				mappedKeys[key] = nil
-				mappedItems[mappedKey.key] = nil
-				continue
-			end
-
-			-- If the key changed since the last transform call, remove
-			-- the old key and update mappedKey.
-			if keyOut ~= mappedKey.key then
-				mappedItems[mappedKey.key] = nil
-				mappedKey.key = keyOut
-			end
-
-			mappedKey.value = valueOut
-			mappedItems[keyOut] = valueOut
-		end
-
-		-- Add new keys if the value was not present in the previous state
-		for key, currentValue in currentItems do
-			local oldValue = previousItems[key]
-
-			if oldValue ~= nil then
-				continue
-			end
-
-			local valueOut, keyOut = getTransform(currentValue, key)
-			local keyOut = if keyOut ~= nil then keyOut else key :: never
+			local valueOut, keyOut = map(value, key)
+			local cachedKeyOut = keyCache[key]
 
 			if valueOut ~= nil then
-				mappedKeys[key] = { key = keyOut, value = valueOut }
-				mappedItems[keyOut] = valueOut
+				-- Skip update if the key and value are the same
+				if keyOut == cachedKeyOut and valueOut == output[cachedKeyOut] then
+					continue
+				end
+
+				output = ensureCanMutateOutput(output, previousOutput)
+				output[keyOut] = valueOut
+				keyCache[key] = keyOut
+
+				-- Remove the old key if a new key is assigned
+				if cachedKeyOut ~= nil and keyOut ~= cachedKeyOut then
+					output[cachedKeyOut] = nil
+				end
+			elseif cachedKeyOut ~= nil then
+				-- Remove the cached key if the new value is nil
+				output = ensureCanMutateOutput(output, previousOutput)
+				output[cachedKeyOut] = nil
+				keyCache[key] = nil
 			end
 		end
 
-		previousItems = currentItems
-
-		return mappedItems
+		return output
 	end)
 end
 
-return mapped
+return mapped :: Mapped

pkgs/charm/src/observe.luau

@@ -8,11 +8,13 @@ local onScopeDispose = signals.onScopeDispose
 local pauseTracking = signals.pauseTracking
 local resumeTracking = signals.resumeTracking
 
-local function observe<K, V>(selector: () -> { [K]: V }, observer: (item: V, key: K) -> ...() -> ()?): () -> ()
+type Cleanup = () -> ()
+
+local function observe<K, V>(selector: () -> { [K]: V }, observer: (item: V, key: K) -> ...Cleanup?): () -> ()
 	local scopes: { [K]: () -> () } = {}
 	local currentItems: { [K]: V } = {}
 
-	local function update(currentItems: { [K]: V }, previousItems: { [K]: V })
+	local function updateScopes(currentItems: { [K]: V }, previousItems: { [K]: V })
 		for key in previousItems do
 			if currentItems[key] == nil then
 				local dispose = scopes[key]
@@ -41,14 +43,12 @@ local function observe<K, V>(selector: () -> { [K]: V }, observer: (item: V, key
 	end
 
 	local disposeEffect = effect(function()
-		local nextItems = selector()
-
-		if nextItems ~= currentItems then
-			local previousItems = currentItems
-			currentItems = nextItems
+		local previousItems = currentItems
+		currentItems = selector()
 
+		if currentItems ~= previousItems then
 			pauseTracking()
-			local success, result: unknown = pcall(update, currentItems, previousItems)
+			local success, result: unknown = pcall(updateScopes, currentItems, previousItems)
 			resumeTracking()
 
 			if not success then

pkgs/charm/src/signals.luau

@@ -78,7 +78,7 @@ local batchDepth = 0
 local notifyIndex = 0
 local queuedEffectsLength = 0
 local activeSub: ReactiveNode?
-local activeScope: EffectScope?
+local activeScope: Effect | EffectScope?
 
 local function getCurrentSub(): ReactiveNode?
 	return activeSub
@@ -90,26 +90,28 @@ local function setCurrentSub(sub: ReactiveNode?)
 	return prevSub
 end
 
-local function getCurrentScope(): EffectScope?
+local function getCurrentScope(): Effect | EffectScope?
 	return activeScope
 end
 
-local function setCurrentScope(scope: EffectScope?)
+local function setCurrentScope(scope: Effect | EffectScope?)
 	local prevScope = activeScope
 	activeScope = scope
 	return prevScope
 end
 
-local function addCleanup(e: EffectScope, cleanup: () -> ())
-	local prevCleanup = e.cleanup
+local function addCleanup(e: Effect | EffectScope, cleanup: () -> ())
+	local prevCleanup: (() -> ())? = e.cleanup
+	local nextCleanup = cleanup
+
 	if prevCleanup then
-		e.cleanup = function()
+		function cleanup()
 			prevCleanup()
-			cleanup()
+			nextCleanup()
 		end
-	else
-		e.cleanup = cleanup
 	end
+
+	(e :: Effect).cleanup = cleanup
 end
 
 local function onEffectCleanup(cleanup: () -> ())
@@ -124,19 +126,19 @@ local function onScopeDispose(cleanup: () -> ())
 	end
 end
 
-local function cleanupEffect(e: EffectScope)
+local function cleanupEffect(e: Effect | EffectScope)
 	if e.cleanup then
 		local prev = setCurrentSub(nil)
 		local success, result: unknown = pcall(e.cleanup)
-		setCurrentSub(prev)
-		e.cleanup = nil
+		setCurrentSub(prev);
+		(e :: Effect).cleanup = nil
 		if not success then
 			error(result)
 		end
 	end
 end
 
-local function run(e: EffectScope, flags: ReactiveFlags): ()
+local function run(e: Effect | EffectScope, flags: ReactiveFlags): ()
 	if bit32.band(flags, DIRTY) ~= 0 or (bit32.band(flags, PENDING) ~= 0 and checkDirty(e.deps :: Link, e)) then
 		cleanupEffect(e)
 		local prev = setCurrentSub(e)
@@ -149,7 +151,7 @@ local function run(e: EffectScope, flags: ReactiveFlags): ()
 		end
 		return
 	elseif bit32.band(flags, PENDING) ~= 0 then
-		e.flags = bit32.band(flags, bit32.bnot(PENDING))
+		(e :: Effect).flags = bit32.band(flags, bit32.bnot(PENDING))
 	end
 	local link = e.deps
 	while link do
@@ -158,7 +160,7 @@ local function run(e: EffectScope, flags: ReactiveFlags): ()
 		if bit32.band(depFlags, QUEUED) ~= 0 then
 			depFlags = bit32.band(depFlags, bit32.bnot(QUEUED))
 			dep.flags = depFlags
-			run(dep :: EffectScope, depFlags)
+			run(dep :: Effect, depFlags)
 		end
 		link = link.nextDep
 	end
@@ -167,7 +169,7 @@ end
 local function flush(): ()
 	while notifyIndex < queuedEffectsLength do
 		notifyIndex += 1
-		local effect = queuedEffects[notifyIndex] :: EffectScope
+		local effect = queuedEffects[notifyIndex] :: Effect
 		queuedEffects[notifyIndex] = nil
 		effect.flags = bit32.band(effect.flags, bit32.bnot(QUEUED))
 		run(effect, effect.flags)
@@ -302,7 +304,7 @@ local function computedGetter<T>(this: Computed<T>): () -> T
 	end
 end
 
-local function effectStop(this: EffectScope): () -> ()
+local function effectStop(this: Effect | EffectScope): () -> ()
 	return function()
 		local dep = this.deps
 		while dep do
@@ -313,7 +315,7 @@ local function effectStop(this: EffectScope): () -> ()
 			unlink(sub)
 			sub = this.subs
 		end
-		this.flags = NONE
+		(this :: Effect).flags = NONE
 		cleanupEffect(this)
 	end
 end
@@ -368,43 +370,43 @@ local function effectScope(fn: () -> (), detached: boolean?): () -> ()
 	return effectStop(e)
 end
 
-local function notify(e: ReactiveNode)
+local function update(signal: Signal | Computed): boolean
+	if (signal :: Computed).getter then
+		return updateComputed(signal :: Computed)
+	else
+		return updateSignal(signal :: Signal, signal.value)
+	end
+end
+
+local function notify(e: Effect | EffectScope)
 	local flags = e.flags
 	if bit32.band(flags, QUEUED) == 0 then
-		e.flags = bit32.bor(flags, QUEUED)
+		(e :: Effect).flags = bit32.bor(flags, QUEUED)
 		local subs = e.subs
 		if subs then
-			notify(subs.sub :: Effect | EffectScope)
+			notify(subs.sub :: Effect)
 		else
 			queuedEffectsLength += 1
-			queuedEffects[queuedEffectsLength] = e :: Effect
+			queuedEffects[queuedEffectsLength] = e
 		end
 	end
 end
 
-local function update(signal: ReactiveNode): boolean
-	if (signal :: Computed).getter then
-		return updateComputed(signal :: Computed)
-	else
-		return updateSignal(signal :: Signal, (signal :: Signal).value)
-	end
-end
-
-local function unwatched(signal: ReactiveNode)
+local function unwatched(signal: Signal | Effect | Computed)
 	local toRemove = signal.deps
 	if toRemove then
 		repeat
 			toRemove = unlink(toRemove, signal)
-		until not toRemove
-		signal.flags = bit32.bor(signal.flags, DIRTY)
+		until toRemove == nil
+		(signal :: Signal).flags = bit32.bor(signal.flags, DIRTY)
 	end
-	cleanupEffect(signal :: EffectScope)
+	cleanupEffect(signal :: Effect)
 end
 
 local reactiveSystem = createReactiveSystem({
-	update = update,
-	notify = notify,
-	unwatched = unwatched,
+	update = update :: (ReactiveNode) -> boolean,
+	notify = notify :: (ReactiveNode) -> (),
+	unwatched = unwatched :: (ReactiveNode) -> (),
 })
 
 link = reactiveSystem.link

pkgs/charm/src/spawn.luau

@@ -1,36 +0,0 @@
-local freeThread: thread?
-
-local function run<T...>(thread: thread, fn: (T...) -> (), ...: T...): boolean
-	fn(...)
-	if freeThread then
-		return true
-	else
-		freeThread = thread
-		return false
-	end
-end
-
-local function runner()
-	local quit = false
-	repeat
-		quit = run(coroutine.yield())
-	until quit
-end
-
-local function spawn<T...>(callback: (T...) -> (), ...: T...)
-	local thread
-
-	if freeThread then
-		thread = freeThread
-		freeThread = nil
-	elseif task then
-		thread = task.spawn(runner)
-	else
-		thread = coroutine.create(runner)
-		coroutine.resume(thread)
-	end
-
-	coroutine.resume(thread, thread, callback, ...)
-end
-
-return spawn

pkgs/charm/src/subscribe.luau

@@ -1,27 +1,27 @@
 local signals = require("./signals")
-local spawn = require("./spawn")
 
 local effect = signals.effect
 local pauseTracking = signals.pauseTracking
 local resumeTracking = signals.resumeTracking
 
-local function subscribe<T>(selector: () -> T, listener: (T, T) -> ()): () -> ()
+local function subscribe<T>(selector: () -> T, listener: (state: T, lastState: T) -> ()): () -> ()
 	local current: T
 	local mounted = true
 
 	return effect(function()
-		local value = selector()
+		local previous = current
+		current = selector()
 
 		if mounted then
-			current = value
 			mounted = false
-		elseif value ~= current then
-			local previous = current
-			current = value
-
+		elseif current ~= previous then
 			pauseTracking()
-			spawn(listener, value, previous)
+			local success, result: unknown = pcall(listener, current, previous)
 			resumeTracking()
+
+			if not success then
+				error(`subscribe() encountered an error: {result}`)
+			end
 		end
 	end)
 end

pkgs/charm/src/types.luau

@@ -2,8 +2,4 @@ export type Selector<T = any> = () -> T
 
 export type Atom<T = any> = ((((T) -> T) | T) -> T) & (() -> T)
 
-export type AtomOptions<T = any> = {
-	equals: ((prev: T, next: T) -> boolean)?,
-}
-
 return {}

tests/init.luau

@@ -1,14 +1,11 @@
 local suite = require("@tests/suite")
 
-local describe = suite.describe
 local finish = suite.finish
 
 if not table.isfrozen(_G) then
 	_G.__DEV__ = true
 end
 
-describe("charm", function()
-	require("@pkgs/charm/tests/spec")
-end)
+require("@pkgs/charm/tests/spec")
 
 finish()