fix: remove old keys when updating key in mapped()

packages/charm/src/init.luau

@@ -874,62 +874,76 @@ end
 	of an existing atom. If the key is omitted from the mapper's return value,
 	the original key is preserved.
 
+	This function is optimized to minimize unnecessary updates by only calling
+	the mapper for entries that have changed since the last run.
+
 	@param getter A signal or function that returns a table
 	@param mapper Function that maps each entry from the source to a new entry in the mapped table
 	@see https://github.com/littensy/charm?tab=readme-ov-file#mappedgetter-mapper
 ]]
 local function mapped<KI, VI, KO, VO>(getter: Getter<{ [KI]: VI }>, mapper: (VI, KI) -> (VO, KO)): Getter<{ [KO]: VO }>
-	local currentInput: { [KI]: VI } = {}
-	local keyCache: { [KI]: KO } = {}
+	local nextInput: { [KI]: VI } = {}
+	local keyMap: { [KI]: KO } = {}
 
-	local function map(value: VI, key: KI): (VO, KO)
+	local function callMapper(value: VI, key: KI): (VO, KO)
 		local prevSub = setActiveSub(nil)
 		local success, valueOut, keyOut = pcall(mapper, value, key)
 		activeSub = prevSub
 		if not success then
 			error(`mapped() transform function failed: {valueOut}`)
 		end
-		return valueOut, if keyOut == nil then key else keyOut :: any
+		return valueOut, if keyOut == nil then key :: any else keyOut
 	end
 
-	return computed(function(previousOutput)
-		local currentOutput: { [KO]: VO } = previousOutput or {}
-		local previousInput = currentInput
-		currentInput = getter()
+	return computed(function(lastOutput)
+		local nextOutput = lastOutput or {}
+		local prevInput = nextInput
+		nextInput = getter()
 
-		for keyIn in previousInput do
-			if currentInput[keyIn] ~= nil then
+		-- For keys that were removed from the input, remove the corresponding
+		-- keys from the output.
+		for prevKeyIn in prevInput do
+			if nextInput[prevKeyIn] == nil then
+				local prevKeyOut = keyMap[prevKeyIn]
+				if prevKeyOut ~= nil then
+					if nextOutput == lastOutput then
+						nextOutput = table.clone(nextOutput)
+					end
+					nextOutput[prevKeyOut] = nil
+					keyMap[prevKeyIn] = nil
+				end
+			end
+		end
+
+		-- For keys that are new or changed, pass it to the mapper and update
+		-- the output. If the mapper returns a new key that is different from
+		-- the previous call, the old mapped key should be removed from the
+		-- output.
+		for nextKeyIn, nextValueIn in nextInput do
+			if nextValueIn == prevInput[nextKeyIn] then
 				continue
 			end
 
-			local prevKeyOut = keyCache[keyIn]
+			local nextValueOut, nextKeyOut = callMapper(nextValueIn, nextKeyIn)
+			local prevKeyOut = keyMap[nextKeyIn]
 
-			if prevKeyOut then
-				if currentOutput == previousOutput then
-					currentOutput = table.clone(currentOutput)
+			if prevKeyOut ~= nil and prevKeyOut ~= nextKeyOut then
+				if nextOutput == lastOutput then
+					nextOutput = table.clone(nextOutput)
 				end
-				currentOutput[prevKeyOut] = nil
-				keyCache[keyIn] = nil
+				nextOutput[prevKeyOut] = nil
+				nextOutput[nextKeyOut] = nextValueOut
+				keyMap[nextKeyIn] = nextKeyOut
+			elseif nextOutput[nextKeyOut] ~= nextValueOut then
+				if nextOutput == lastOutput then
+					nextOutput = table.clone(nextOutput)
+				end
+				nextOutput[nextKeyOut] = nextValueOut
+				keyMap[nextKeyIn] = nextKeyOut
 			end
 		end
 
-		for keyIn, valueIn in currentInput do
-			if valueIn == previousInput[keyIn] then
-				continue
-			end
-
-			local valueOut, keyOut = map(valueIn, keyIn)
-
-			if currentOutput[keyOut] ~= valueOut then
-				if currentOutput == previousOutput then
-					currentOutput = table.clone(currentOutput)
-				end
-				currentOutput[keyOut] = valueOut
-				keyCache[keyIn] = keyOut
-			end
-		end
-
-		return currentOutput
+		return nextOutput
 	end)
 end
 

packages/charm/test/mapped.test.luau

@@ -104,4 +104,25 @@ test("should map keys", function()
 	assert(state.D == 4, "should add d")
 end)
 
+test("should remove old key if mapper returns a new key", function()
+	local src, setSrc = signal({ 1 })
+	local map = mapped(src, function(value: number, index: number)
+		return true, `{index}={value}`
+	end)
+
+	local state = map()
+	assert(state["1=1"], "should have mapped first element to key 1=1")
+
+	setSrc({ 2 })
+	state = map()
+	assert(state["1=1"] == nil, "should have removed old key 1=1")
+	assert(state["1=2"], "should have remapped element to key 1=2")
+
+	setSrc({ 3, 1 })
+	state = map()
+	assert(state["1=2"] == nil, "should have removed old key 1=2")
+	assert(state["1=3"], "should have remapped element to key 1=3")
+	assert(state["2=1"], "should have mapped second element to key 2=1")
+end)
+
 return {}