Fix compacting

compact_atc_chains.py

@@ -45,10 +45,13 @@ def compact_chains(data: object) -> list[list[object]]:
 
     for chain in chains:
         if isinstance(chain, list):
-            if len(chain) != 3:
+            if len(chain) == 2 and isinstance(chain[1], list):
-                raise ValueError("Invalid compact chain entry")
                 compact.append(chain)
                 continue
+            if len(chain) == 3 and isinstance(chain[1], int) and isinstance(chain[2], list):
+                compact.append([chain[0], chain[2]])
+                continue
+            raise ValueError("Invalid compact chain entry")
 
         if not isinstance(chain, dict):
             raise ValueError("Expected verbose or compact chain entries as input")
@@ -58,16 +61,15 @@ def compact_chains(data: object) -> list[list[object]]:
             continue
 
         base_id = chain["base_id"]
-        base_icon = crop_icon(items[0].get("icon_nfp_16x16"))
+        compact_items: list[list[object]] = []
-        compact_items = []
 
-        for item in items[1:]:
+        for item in items:
             compact_items.append([
                 item["item_id"],
                 crop_icon(item.get("icon_nfp_16x16")),
             ])
 
-        compact.append([base_id, base_icon, compact_items])
+        compact.append([base_id, compact_items])
 
     return compact
 
@@ -76,17 +78,19 @@ def main() -> int:
     parser = argparse.ArgumentParser(
         description="Rewrite atc_chains.json into a compact format used by compcount.lua."
     )
-    parser.add_argument("path", nargs="?", default="atc_chains.json")
+    parser.add_argument("source", nargs="?", default="atc_chains_uncompressed.json")
+    parser.add_argument("destination", nargs="?", default="atc_chains.json")
     args = parser.parse_args()
 
-    path = Path(args.path)
+    source_path = Path(args.source)
-    original_text = path.read_text(encoding="utf-8")
+    destination_path = Path(args.destination)
+    original_text = source_path.read_text(encoding="utf-8")
     data = json.loads(original_text)
     compact = {"chains": compact_chains(data)}
     compact_text = json.dumps(compact, separators=(",", ":"))
-    path.write_text(compact_text, encoding="utf-8")
+    destination_path.write_text(compact_text, encoding="utf-8")
 
-    print(f"Rewrote {path}")
+    print(f"Compacted {source_path} -> {destination_path}")
     print(f"Old size: {len(original_text.encode('utf-8'))} bytes")
     print(f"New size: {len(compact_text.encode('utf-8'))} bytes")
     return 0

compcount.lua

@@ -45,6 +45,7 @@ local function ensureStorageBridge()
 end
 
 local chainItemsByBaseId
+local chainAllItemsByBaseId
 local baseItemIds
 local allChainItemIds
 local itemById
@@ -117,6 +118,33 @@ local function getCompressionLevel(itemId, stage)
   return 0
 end
 
+local function createChainItem(itemId, iconNfp, stage)
+  return {
+    id = itemId,
+    icon_nfp = iconNfp,
+    stage = stage,
+    compression_level = getCompressionLevel(itemId, stage),
+    equivalent_factor = 1,
+  }
+end
+
+local function applyChainEquivalentFactors(baseId, orderedItems)
+  local baseIndex = nil
+
+  for i = 1, #orderedItems do
+    if orderedItems[i].id == baseId then
+      baseIndex = i
+      break
+    end
+  end
+
+  assert(baseIndex, "Missing base item in chain: " .. tostring(baseId))
+
+  for i = 1, #orderedItems do
+    orderedItems[i].equivalent_factor = 9 ^ (i - baseIndex)
+  end
+end
+
 local function ensureChainsLoaded()
   if not chainItemsByBaseId then
     ensureWhitelistLoaded()
@@ -134,6 +162,7 @@ local function ensureChainsLoaded()
     assert(type(chains) == "table", "Invalid JSON data")
 
     chainItemsByBaseId = {}
+    chainAllItemsByBaseId = {}
     baseItemIds = {}
     allChainItemIds = {}
     itemById = {}
@@ -148,60 +177,50 @@ local function ensureChainsLoaded()
 
         baseItemIds[#baseItemIds + 1] = baseId
 
-        local items = {}
+        local orderedItems = {}
-        local itemCount = 0
+        local pageItems = {}
 
         if chain.base_id then
           for _, item in ipairs(chain.items or {}) do
-            local compressionLevel = getCompressionLevel(item.item_id, item.stage)
+            local chainItem = createChainItem(item.item_id, item.icon_nfp_16x16, item.stage)
-
+            orderedItems[#orderedItems + 1] = chainItem
-            if item.stage == "item" then
+            itemById[item.item_id] = chainItem
-              itemById[item.item_id] = {
-                id = item.item_id,
-                icon_nfp = item.icon_nfp_16x16,
-                compression_level = compressionLevel,
-                compression_factor = 9 ^ compressionLevel,
-              }
-              allChainItemIds[#allChainItemIds + 1] = item.item_id
-            else
-              itemCount = itemCount + 1
-              items[itemCount] = {
-                id = item.item_id,
-                icon_nfp = item.icon_nfp_16x16,
-                compression_level = compressionLevel,
-                compression_factor = 9 ^ compressionLevel,
-              }
-              itemById[item.item_id] = items[itemCount]
             allChainItemIds[#allChainItemIds + 1] = item.item_id
           end
-          end
-        else
-          local baseCompressionLevel = getCompressionLevel(baseId, "item")
 
-          itemById[baseId] = {
+        else
-            id = baseId,
+          local compactItems = nil
-            icon_nfp = chain[2],
+
-            compression_level = baseCompressionLevel,
+          if type(chain[2]) == "number" then
-            compression_factor = 9 ^ baseCompressionLevel,
+            compactItems = chain[3] or {}
-          }
+          elseif type(chain[2]) == "table" then
+            compactItems = chain[2]
+          elseif type(chain[2]) == "string" and type(chain[3]) == "table" then
+            orderedItems[1] = createChainItem(baseId, chain[2], "item")
+            itemById[baseId] = orderedItems[1]
             allChainItemIds[#allChainItemIds + 1] = baseId
+            compactItems = chain[3]
+          else
+            error("Invalid compact chain entry for " .. tostring(baseId))
+          end
 
-          for i = 1, #(chain[3] or {}) do
+          for i = 1, #compactItems do
-            local compactItem = chain[3][i]
+            local compactItem = compactItems[i]
-            local compressionLevel = getCompressionLevel(compactItem[1], nil)
+            local chainItem = createChainItem(compactItem[1], compactItem[2], compactItem[3])
-            itemCount = itemCount + 1
+            orderedItems[#orderedItems + 1] = chainItem
-            items[itemCount] = {
+            itemById[compactItem[1]] = chainItem
-              id = compactItem[1],
-              icon_nfp = compactItem[2],
-              compression_level = compressionLevel,
-              compression_factor = 9 ^ compressionLevel,
-            }
-            itemById[compactItem[1]] = items[itemCount]
             allChainItemIds[#allChainItemIds + 1] = compactItem[1]
           end
         end
 
-        chainItemsByBaseId[baseId] = items
+        applyChainEquivalentFactors(baseId, orderedItems)
+
+        for i = 1, #orderedItems do
+          pageItems[#pageItems + 1] = orderedItems[i]
+        end
+
+        chainAllItemsByBaseId[baseId] = orderedItems
+        chainItemsByBaseId[baseId] = pageItems
       end
     end
   end
@@ -213,6 +232,12 @@ local function getPageItems(base_id)
   return chainItemsByBaseId[base_id or defaultBaseId] or {}
 end
 
+local function getAllChainItems(base_id)
+  ensureChainsLoaded()
+
+  return chainAllItemsByBaseId[base_id or defaultBaseId] or {}
+end
+
 local function getBaseItemIds()
   ensureChainsLoaded()
 
@@ -447,16 +472,11 @@ end
 
 local function getBaseEquivalentCount(baseId, itemCounts)
   local total = 0
-  local baseItem = getItemById(baseId)
+  local chainItems = getAllChainItems(baseId)
-  local pageItems = getPageItems(baseId)
 
-  if baseItem then
+  for i = 1, #chainItems do
-    total = total + (tonumber(itemCounts[baseId]) or 0) * (baseItem.compression_factor or 1)
+    local item = chainItems[i]
-  end
+    total = total + (tonumber(itemCounts[item.id]) or 0) * (item.equivalent_factor or 1)
-
-  for i = 1, #pageItems do
-    local item = pageItems[i]
-    total = total + (tonumber(itemCounts[item.id]) or 0) * (item.compression_factor or 1)
   end
 
   return total
@@ -468,7 +488,7 @@ local function getEquivalentLevelCounts(baseId, pageItems, itemCounts)
 
   for i = 1, #pageItems do
     local item = pageItems[i]
-    local factor = item.compression_factor or 1
+    local factor = item.equivalent_factor or 1
     equivalentCounts[item.id] = math.floor(baseEquivalentCount / factor)
   end