fix(data structures): use maps instead of lists for data structures

Author: atticusofsparta

src/common/main.lua

@@ -14,12 +14,12 @@ main.init = function()
 	-- Example ANT structure
 	-- ANTS["antId"] = {
 	--     Owner = "userId",
-	--     Controllers = {"userId1", "userId2"},
+	--     Controllers = {"userId1" = true, "userId2" = true},
 	-- }
 	ANTS = ANTS or {}
 
 	-- Example ADDRESSES structure - keyed references to the above ANTS structure
-	-- ADDRESSES["userAddress"] = {["antProcessId1"] = ANTS["antProcessId1"], ANTS["antProcessId2"]}
+	-- ADDRESSES["userAddress"] = {"antProcessId1" = true, "antProcessId2" = true}
 	ADDRESSES = ADDRESSES or {}
 
 	--[[
@@ -71,11 +71,14 @@ main.init = function()
 		assert(type(antId) == "string", "Process-Id tag is required")
 		assert(ANTS[antId] == nil, "ANT is already registered")
 
-		ANTS[antId] = {
-			Owner = nil,
-			Controllers = {},
-		}
-
+		utils.register({
+			id = antId,
+			timestamp = tonumber(msg.Timestamp),
+		})
+		--[[
+			Send a request message for current ANT state to the process. Expect back 
+			a State-Notice so that we can update the registered ANT settings.
+		]]
 		ao.send({
 			Target = antId,
 			Action = "State",
@@ -93,12 +96,12 @@ main.init = function()
 
 		-- Register the ANT if not already registered
 		if not isRegistered then
-			ANTS[msg.From] = {
-				Owner = stateRes.Owner,
-				Controllers = {},
-				-- for cleaning function
-				RegisteredAt = tonumber(msg.Timestamp),
-			}
+			utils.register({
+				id = msg.From,
+				timestamp = tonumber(msg.Timestamp),
+				owner = stateRes.Owner,
+				controllers = stateRes.Controllers,
+			})
 		end
 
 		utils.updateAssociations(msg.From, stateRes)
@@ -117,7 +120,12 @@ main.init = function()
 		local address = msg.Tags["Address"]
 		assert(type(address) == "string", "Address is required")
 
-		local antIds = ADDRESSES[address] or {}
+		local antIds = {}
+		if ADDRESSES[address] then
+			for antId, _ in pairs(ADDRESSES[address]) do
+				table.insert(antIds, antId)
+			end
+		end
 
 		-- Send the list of ant_ids as a JSON array
 		ao.send({

src/common/utils.lua

@@ -242,54 +242,78 @@ function utils.indexOf(t, value)
 	return -1
 end
 
-function utils.updateUserAssociations(antId, state, user)
-	-- remove previous associations
-	if ADDRESSES[user] and user ~= state.Owner and not utils.includes(user, state.Controllers) then
-		local antIndex = utils.indexOf(ADDRESSES[user], antId)
-		if antIndex then
-			table.remove(ADDRESSES[user], antIndex)
-		end
-	end
+--- func description
+-- @param {table} params
+-- @param {string} params.id
+-- @param {number} params.timestamp
+-- @param {string | nil} params.owner
+-- @param {table | nil} params.controllers
+function utils.register(params)
+	assert(type(params) == "table", "Register Params must be a table")
+	local id, timestamp, owner, controllers = params.id, params.timestamp, params.owner, params.controllers
+	assert(type(id) == "string", "ANT ID must be a string")
+	assert(type(timestamp) == "number", "Timestamp must be a number")
+	assert(type(owner) == "string" or owner == nil, "Owner must be a string")
+	assert(type(controllers) == "table" or controllers == nil, "Controllers must be a table")
+	ANTS[id] = {
+		Owner = owner,
+		Controllers = controllers or {},
+		RegisteredAt = timestamp,
+	}
+end
 
-	-- add new associations
-	if user == state.Owner or utils.includes(user, state.Controllers) then
-		if not ADDRESSES[user] then
-			ADDRESSES[user] = {}
-		end
-		if not utils.includes(antId, ADDRESSES[user]) then
-			table.insert(ADDRESSES[user], antId)
-		end
+function utils.controllerTableFromArray(t)
+	assert(type(t) == "table", "argument needs to be a table")
+	local map = {}
+	for _, v in ipairs(t) do
+		map[v] = true
 	end
+	return map
 end
 
 function utils.updateAssociations(antId, state)
 	-- Remove previous associations for old owner and controllers
-	local previousOwner = ANTS[antId].Owner
-	local previousControllers = ANTS[antId].Controllers
+	local oldAnt = ANTS[antId]
+
+	local newOwner = state.Owner or "nilOwner"
+	local newControllers = utils.controllerTableFromArray(state.Controllers)
 
-	if previousOwner and previousOwner ~= state.Owner and type(ANTS[antId]) == "table" then
-		utils.updateUserAssociations(antId, ANTS[antId], previousOwner)
+	local newAffliates = {}
+	newAffliates[newOwner] = true
+	for user, _ in pairs(newControllers) do
+		newAffliates[user] = true
 	end
 
-	for _, user in ipairs(previousControllers) do
-		if not utils.includes(user, state.Controllers) then
-			utils.updateUserAssociations(antId, ANTS[antId], user)
+	if oldAnt ~= nil then
+		local previousOwner = ANTS[antId].Owner or "nilOwner"
+		local previousControllers = ANTS[antId].Controllers
+
+		local oldAffliates = {}
+		oldAffliates[previousOwner] = true
+		for user, _ in pairs(previousControllers) do
+			oldAffliates[user] = true
 		end
-	end
 
-	-- Add new associations for new owner and controllers
-	local users = { state.Owner }
-	for _, user in ipairs(state.Controllers) do
-		table.insert(users, user)
+		for oldAffliate, _ in pairs(oldAffliates) do
+			if not newAffliates[oldAffliate] and ADDRESSES[oldAffliate] then
+				ADDRESSES[oldAffliate][antId] = nil
+			end
+		end
 	end
 
-	for _, user in ipairs(users) do
-		utils.updateUserAssociations(antId, state, user)
+	for user, _ in pairs(newAffliates) do
+		ADDRESSES[user] = ADDRESSES[user] or {}
+		ADDRESSES[user][antId] = true
 	end
-
 	-- Update the ANTS table
-	ANTS[antId].Owner = state.Owner
-	ANTS[antId].Controllers = state.Controllers
+	if not newOwner and #newControllers == 0 then
+		ANTS[antId] = nil
+	else -- remove ant from registry if it has no owner or controllers
+		ANTS[antId] = {
+			Owner = newOwner,
+			Controllers = newControllers,
+		}
+	end
 end
 
 -- it is possible to register an ANT that does not respond to the state request, so we need to clean up the ANTS table

test/register.test.ts

@@ -68,7 +68,7 @@ describe('ANT Registration Cases', async () => {
     );
 
     // if we have messages in this case we have errors
-    console.log('stateNoticeResult', stateNoticeResult);
+
     assert.strictEqual(stateNoticeResult.Messages.length, 0);
 
     const allAntsResult = await sendMessage(
@@ -197,6 +197,6 @@ describe('ANT Registration Cases', async () => {
     const failureNotice = stateNoticeResult.Messages[0].Tags.find(
       (tag) => tag.name === 'Action',
     );
-    assert.strictEqual(failureNotice.value, 'State-Notice-Failure');
+    assert.strictEqual(failureNotice.value, 'Invalid-State-Notice-Notice');
   });
 });