Fix and cleanup OtherArmyUnitCountBuildConditions.lua (FAForever#4310)

lua/editor/OtherArmyUnitCountBuildConditions.lua

@@ -9,128 +9,127 @@
 --**  Copyright © 2005 Gas Powered Games, Inc.  All rights reserved.
 --****************************************************************************
 
+---@alias CompareType
+---| ">="
+---| "<="
+---| "=="
+---| '>'
+---| '<'
+
+
 ---@param aiBrain AIBrain
----@param targetBrain string
----@param numReq integer
+---@param targetBrains string[]
+---@param numReq number
 ---@param category EntityCategory
+---@param compareType CompareType? defaults to `">="`
 ---@return boolean
-function BrainGreaterThanNumCategory( aiBrain, targetBrain, numReq, category )
-    local testBrain = ArmyBrains[1]
-    for k,v in ArmyBrains do
-        if v.Name == targetBrain then
-            testBrain = v
-            break
+function BrainsCompareNumCategory(aiBrain, targetBrains, numReq, category, compareType)
+    local num = 0
+    local targetBrainSet = {}
+    local armySetup = ScenarioInfo.ArmySetup
+    for _, brain in targetBrains do
+        if brain == 'HumanPlayers' then
+            local tblArmy = ListArmies()
+            for _, strArmy in ipairs(tblArmy) do
+                if armySetup[strArmy].Human then
+                    targetBrainSet[armySetup[strArmy].ArmyName] = true
+                end
+            end
+        else
+            targetBrainSet[brain] = true
         end
     end
-    local numUnits = testBrain:GetCurrentUnits(category)
-    if numUnits > numReq then
-        return true
+
+    for _, testBrain in ipairs(ArmyBrains) do
+        if targetBrainSet[testBrain.Name] then
+            num = num + testBrain:GetCurrentUnits(category)
+        end
+    end
+
+    if not compareType or compareType == '>=' then
+        return num >= numReq
+    elseif compareType == '==' then
+        return num == numReq
+    elseif compareType == '<=' then
+        return num <= numReq
+    elseif compareType == '>' then
+        return num > numReq
+    elseif compareType == '<' then
+        return num < numReq
     else
         return false
     end
 end
 
 ---@param aiBrain AIBrain
----@param targetBrain string
+---@param targetBrain string[]
 ---@param numReq integer
 ---@param category EntityCategory
 ---@return boolean
-function BrainLessThanNumCategory( aiBrain, targetBrain, numReq, category )
-    local testBrain = ArmyBrains[1]
-    for k,v in ArmyBrains do
-        if v.Name == targetBrain then
-            testBrain = v
-            break
-        end
-    end
-    local numUnits = testBrain:GetCurrentUnits(category)
-    if numUnits < numReq then
-        return true
-    else
-        return false
-    end
+function BrainGreaterThanNumCategory(aiBrain, targetBrain, numReq, category)
+    return BrainsCompareNumCategory(aiBrain, targetBrain, numReq, category, ">")
 end
 
 ---@param aiBrain AIBrain
----@param targetBrain string
+---@param targetBrain string[]
 ---@param numReq integer
 ---@param category EntityCategory
 ---@return boolean
-function BrainGreaterThanOrEqualNumCategory( aiBrain, targetBrain, numReq, category )
-    local testBrain = ArmyBrains[1]
-    for k,v in ArmyBrains do
-        if v.Name == targetBrain then
-            testBrain = v
-            break
-        end
-    end
-    local numUnits = testBrain:GetCurrentUnits(category)
-    if numUnits >= numReq then
-        return true
-    else
-        return false
-    end
+function BrainLessThanNumCategory(aiBrain, targetBrain, numReq, category)
+    return BrainsCompareNumCategory(aiBrain, targetBrain, numReq, category, "<")
 end
 
 ---@param aiBrain AIBrain
----@param targetBrain string
+---@param targetBrain string[]
 ---@param numReq integer
 ---@param category EntityCategory
 ---@return boolean
-function BrainLessThanOrEqualNumCategory( aiBrain, targetBrain, numReq, category )
-    local testBrain = ArmyBrains[1]
-    for k,v in ArmyBrains do
-        if v.Name == targetBrain then
-            testBrain = v
-            break
-        end
-    end
-    local numUnits = testBrain:GetCurrentUnits(category)
-    if numUnits <= numReq then
-        return true
-    else
-        return false
-    end
+function BrainGreaterThanOrEqualNumCategory(aiBrain, targetBrain, numReq, category)
+    return BrainsCompareNumCategory(aiBrain, targetBrain, numReq, category, ">=")
+end
+
+---@param aiBrain AIBrain
+---@param targetBrain string[]
+---@param numReq integer
+---@param category EntityCategory
+---@return boolean
+function BrainLessThanOrEqualNumCategory(aiBrain, targetBrain, numReq, category)
+    return BrainsCompareNumCategory(aiBrain, targetBrain, numReq, category, "<=")
 end
 
 ---@param aiBrain AIBrain
 ---@param numReq integer
 ---@param categories EntityCategory
----@param compareType string
+---@param compareType CompareType? defaults to `">="`
 ---@return boolean
-function FocusBrainBeingBuiltOrActiveCategoryCompare( aiBrain, numReq, categories, compareType )
-    local testBrain = ArmyBrains[GetFocusArmy()]
+function FocusBrainBeingBuiltOrActiveCategoryCompare(aiBrain, numReq, categories, compareType)
     local num = 0
-    for k,v in categories do
-        num = num + testBrain:GetBlueprintStat('Units_BeingBuilt', v)
-        num = num + testBrain:GetBlueprintStat('Units_Active', v)
+    local tblArmy = ListArmies()
+    for iArmy, strArmy in pairs(tblArmy) do
+        if ScenarioInfo.ArmySetup[strArmy].Human then
+            local testBrain = GetArmyBrain(strArmy)
+            for k, v in categories do
+                num = num + testBrain:GetBlueprintStat('Units_BeingBuilt', v)
+                num = num + testBrain:GetBlueprintStat('Units_Active', v)
+            end
+        end
     end
-
     if not compareType or compareType == '>=' then
-        if num >= numReq then
-            return true
-        end
+        return num >= numReq
     elseif compareType == '==' then
-        if num == numReq then
-            return true
-        end
+        return num == numReq
     elseif compareType == '<=' then
-        if num <= numReq then
-            return true
-        end
+        return num <= numReq
     elseif compareType == '>' then
-        if num > numReq then
-            return true
-        end
+        return num > numReq
     elseif compareType == '<' then
-        if num < numReq then
-            return true
-        end
+        return num < numReq
+    else
+        return false
     end
-    return false
 end
 
 -- Moved unsed Imports to bottom for mod compatibilty
 local AIUtils = import('/lua/ai/aiutilities.lua')
 local ScenarioFramework = import('/lua/scenarioframework.lua')
-local ScenarioUtils = import('/lua/sim/ScenarioUtilities.lua')
+local ScenarioUtils = import('/lua/sim/scenarioutilities.lua')