chore(PE-8188): add patch balances work

.gitignore

@@ -17,3 +17,4 @@ coverage
 
 mainnet/**/outputs/*
 mainnet/outputs/*
+._*

src/balances.lua

@@ -91,9 +91,36 @@ end
 --- Checks if a wallet has a sufficient balance
 --- @param wallet string The address of the wallet
 --- @param quantity number The amount to check against the balance
---- @return boolean True if the wallet has a sufficient balance, false otherwise
+--- @return table<string, boolean> A table of addresses and whether they have a sufficient balance
 function balances.walletHasSufficientBalance(wallet, quantity)
 	return Balances[wallet] ~= nil and Balances[wallet] >= quantity
 end
 
+---@param oldBalances table<string, number> A table of addresses and their balances
+---@param newBalances table<string, number> A table of addresses and their balances
+function balances.patchBalances(oldBalances, newBalances)
+	assert(type(oldBalances) == "table", "Old balances must be a table")
+	assert(type(newBalances) == "table", "New balances must be a table")
+	local affectedBalancesAddresses = {}
+	for address, _ in pairs(oldBalances) do
+		if Balances[address] ~= oldBalances[address] then
+			affectedBalancesAddresses[address] = true
+		end
+	end
+	for address, _ in pairs(newBalances) do
+		if oldBalances[address] ~= newBalances[address] then
+			affectedBalancesAddresses[address] = true
+		end
+	end
+
+	--- For simplicity we always include the protocol balance in the patch message
+	local patchMessage = { device = "[email protected]", balances = { [ao.id] = Balances[ao.id] or 0 } }
+	for address, _ in pairs(affectedBalancesAddresses) do
+		patchMessage.balances[address] = Balances[address] or 0
+	end
+
+	ao.send(patchMessage)
+	return affectedBalancesAddresses
+end
+
 return balances

src/main.lua

@@ -471,6 +471,8 @@ local function addEventingHandler(handlerName, pattern, handleFn, critical, prin
 	critical = critical or false
 	printEvent = printEvent == nil and true or printEvent
 	Handlers.add(handlerName, pattern, function(msg)
+		-- Store the old balances to compare after the handler has run for patching state
+		local oldBalances = utils.deepCopy(Balances)
 		-- add an ARIOEvent to the message if it doesn't exist
 		msg.ioEvent = msg.ioEvent or ARIOEvent(msg)
 		-- global handler for all eventing errors, so we can log them and send a notice to the sender for non critical errors and discard the memory on critical errors
@@ -494,6 +496,11 @@ local function addEventingHandler(handlerName, pattern, handleFn, critical, prin
 			error(errorWithEvent, 0) -- 0 ensures not to include this line number in the error message
 		end
 
+		-- Send patch message to HB
+		local affectedBalances = balances.patchBalances(oldBalances, Balances)
+		-- Add affected balances addresses to the event before cleaning up
+		msg.ioEvent:addField("Affected-Balances-Addresses", utils.getTableKeys(affectedBalances))
+
 		msg.ioEvent:addField("Handler-Memory-KiB-Used", collectgarbage("count"), false)
 		collectgarbage("collect")
 		msg.ioEvent:addField("Final-Memory-KiB-Used", collectgarbage("count"), false)
@@ -625,6 +632,7 @@ addEventingHandler(ActionMap.Transfer, utils.hasMatchingTag("Action", ActionMap.
 	msg.ioEvent:addField("RecipientFormatted", recipient)
 
 	local result = balances.transfer(recipient, msg.From, quantity, allowUnsafeAddresses)
+
 	if result ~= nil then
 		local senderNewBalance = result[msg.From]
 		local recipientNewBalance = result[recipient]
@@ -1034,6 +1042,7 @@ addEventingHandler(
 			},
 			Data = json.encode(fundFrom and result or recordResult),
 		})
+		-- TODO patch other ARIO locations based on FUND
 	end
 )
 
@@ -1782,6 +1791,7 @@ end, function(msg)
 	]]
 	--
 	print("Ticking from " .. lastCreatedEpochIndex .. " to " .. targetCurrentEpochIndex)
+
 	for epochIndexToTick = lastCreatedEpochIndex, targetCurrentEpochIndex do
 		local tickResult = tick.tickEpoch(msg.Timestamp, blockHeight, hashchain, msgId, epochIndexToTick)
 		if tickResult.pruneGatewaysResult ~= nil then
@@ -2137,17 +2147,12 @@ end)
 -- Pagination handlers
 
 addEventingHandler("paginatedRecords", function(msg)
-        return msg.Action == "Paginated-Records" or msg.Action == ActionMap.Records
+	return msg.Action == "Paginated-Records" or msg.Action == ActionMap.Records
 end, function(msg)
-        local page = utils.parsePaginationTags(msg)
-        local result = arns.getPaginatedRecords(
-                page.cursor,
-                page.limit,
-                page.sortBy or "startTimestamp",
-                page.sortOrder,
-                page.filters
-        )
-        Send(msg, { Target = msg.From, Action = "Records-Notice", Data = json.encode(result) })
+	local page = utils.parsePaginationTags(msg)
+	local result =
+		arns.getPaginatedRecords(page.cursor, page.limit, page.sortBy or "startTimestamp", page.sortOrder, page.filters)
+	Send(msg, { Target = msg.From, Action = "Records-Notice", Data = json.encode(result) })
 end)
 
 addEventingHandler("paginatedGateways", function(msg)