Annotate sync of enhancements

changelog/snippets/other.6626.md

@@ -0,0 +1 @@
+- (#6626) Annotate sync of enhancements.

lua/ScenarioFramework.lua

@@ -802,7 +802,7 @@ function PlayVoiceOver(voSound)
 end
 
 --- Sets enhancement restrictions from the names of the enhancements you do not want the player to build
----@param enhancements string[]
+---@param enhancements Enhancement[]
 function RestrictEnhancements(enhancements)
     local restrict = {}
     for _, enh in enhancements do

lua/SimSync.lua

@@ -1,7 +1,9 @@
 ---@declare-global
+
 -- The global sync table is copied to the user layer every time the main and sim threads are
 -- synchronized on the sim beat (which is like a tick but happens even when the game is paused)
-
+---@class SyncTable: table
+---@field EnhanceRestrict table<Enhancement, true>
 Sync = { }
 
 local SyncDefaults = {
@@ -19,6 +21,10 @@ end
 -- focus army we can resync the data.
 UnitData = {}
 
+---@alias EnhancementSyncData table<EnhancementSlot, Enhancement>
+---@alias EnhancementSyncTable table<EntityId, EnhancementSyncData>
+
+---@type EnhancementSyncTable
 SimUnitEnhancements = {}
 
 function ResetSyncTable()
@@ -39,6 +45,9 @@ function ResetSyncTable()
     end
 end
 
+---@param unit Unit
+---@param enhancement Enhancement
+---@param slot EnhancementSlot
 function AddUnitEnhancement(unit, enhancement, slot)
     if not slot then return end
     local id = unit.EntityId
@@ -47,6 +56,8 @@ function AddUnitEnhancement(unit, enhancement, slot)
     SyncUnitEnhancements()
 end
 
+---@param unit Unit
+---@param enhancement Enhancement
 function RemoveUnitEnhancement(unit, enhancement)
     if not unit or unit.Dead then return end
     local id = unit.EntityId
@@ -68,13 +79,15 @@ function RemoveUnitEnhancement(unit, enhancement)
     SyncUnitEnhancements()
 end
 
+---@param unit Unit
 function RemoveAllUnitEnhancements(unit)
     local id = unit.EntityId
     if not SimUnitEnhancements[id] then return end
     SimUnitEnhancements[id] = nil
     SyncUnitEnhancements()
 end
 
+--- Syncs enhancements of units allied to us (or all units if we're observer)
 function SyncUnitEnhancements()
     import("/lua/enhancementcommon.lua").SetEnhancementTable(SimUnitEnhancements)
     local sync = {}

lua/enhancementcommon.lua

@@ -36,22 +36,28 @@
 --| UEFSCUEnhancementBuffName # there are none
 ---| SeraphimSCUEnhancementBuffName
 
-
+---@type EnhancementSyncTable
 local enhancementTable = {}
+---@type table<Enhancement, true>
 local restrictedList = {}
 
-function SetEnhancementTable(entry)
-    enhancementTable = entry
+---@param table EnhancementSyncTable
+function SetEnhancementTable(table)
+    enhancementTable = table
 end
 
+---@param table table<Enhancement, true>
 function RestrictList(table)
     restrictedList = table
 end
 
+---@return table<Enhancement, true>
 function GetRestricted()
     return restrictedList
 end
 
+---@param entityID EntityId
+---@return EnhancementSyncData?
 function GetEnhancements(entityID)
     return enhancementTable[tostring(entityID)]
 end

lua/sim/Unit.lua

@@ -3205,7 +3205,7 @@ Unit = ClassUnit(moho.unit_methods, IntelComponent, VeterancyComponent, DebugUni
     end;
 
     ---@param self Unit
-    ---@param enh string
+    ---@param enh Enhancement
     ---@return boolean
     CreateEnhancement = function(self, enh)
         local bp = self.Blueprint.Enhancements[enh]
@@ -3241,7 +3241,7 @@ Unit = ClassUnit(moho.unit_methods, IntelComponent, VeterancyComponent, DebugUni
     end,
 
     ---@param self Unit
-    ---@param enhancement string
+    ---@param enhancement Enhancement
     CreateEnhancementEffects = function(self, enhancement)
         local bp = self.Blueprint.Enhancements[enhancement]
         local effects = TrashBag()
@@ -3277,7 +3277,7 @@ Unit = ClassUnit(moho.unit_methods, IntelComponent, VeterancyComponent, DebugUni
     end,
 
     ---@param self Unit
-    ---@param enh string
+    ---@param enh Enhancement
     ---@return boolean
     HasEnhancement = function(self, enh)
         local unitEnh = SimUnitEnhancements[self.EntityId]

lua/ui/game/construction.lua

@@ -466,6 +466,10 @@ function CreateTabs(type)
                 sortedOptions[slotName] = {}
                 for enhName, enhTable in enhancements do
                     if enhTable.Slot == slotName then
+                        ---@class EnhQueueData: UnitBlueprintEnhancement
+                        ---@field ID Enhancement
+                        ---@field UnitID UnitId
+
                         enhTable.ID = enhName
                         enhTable.UnitID = selection[1]:GetBlueprint().BlueprintId
                         table.insert(sortedOptions[slotName], enhTable)

lua/ui/notify/enhancementqueue.lua

@@ -2,8 +2,12 @@
 
 local SetIgnoreSelection = import("/lua/ui/game/gamemain.lua").SetIgnoreSelection
 
+--- Queue for enhancements. The currently upgrading enhancement is index 1.
+---@alias EnhancementQueue table<EntityId, Enhancement[] | EnhQueueData[]>
 local enhancementQueue = {}
 
+---@param units UserUnit
+---@param enhancement Enhancement
 function enqueueEnhancement(units, enhancement)
     if not units[1] then return end
 
@@ -22,17 +26,22 @@ function enqueueEnhancement(units, enhancement)
     end
 end
 
+---@return EnhancementQueue
 function getEnhancementQueue()
     return enhancementQueue
 end
 
+--- Removes the currently upgrading enhancement from the unit's enhancement queue
+---@param unit UserUnit
 function removeEnhancement(unit)
     local id = unit:GetEntityId()
     if enhancementQueue[id] and not table.empty(enhancementQueue[id]) then
         table.remove(enhancementQueue[id], 1)
     end
 end
 
+--- Removes all queued enhancements for the units
+---@param units UserUnit[]
 function clearEnhancements(units)
     for _, unit in units do
         local id = unit:GetEntityId()
@@ -42,13 +51,19 @@ function clearEnhancements(units)
     end
 end
 
+---@param unit UserUnit
+---@return boolean
 function currentlyUpgrading(unit)
     local currentCommand = unit:GetCommandQueue()[1]
     local queue = enhancementQueue[unit:GetEntityId()][1]
 
-    return currentCommand.type == 'Script' and queue and not string.find(queue.ID, 'Remove')
+    return currentCommand.type == 'Script' and queue and not string.find(queue.ID, 'Remove') -- ex: "StabilitySuppressantRemove"
 end
 
+--- Returns buildable categories for the selection updated based on queued tech suite enhancements
+---@param originalBuildables EntityCategory
+---@param selection UserUnit[]
+---@return EntityCategory
 function ModifyBuildablesForACU(originalBuildables, selection)
     local newBuildableCategories
     local upgradingACUFound = false