Further refinements for factory transfer (#6821)

Author: lL1l1

changelog/snippets/fix.6821.md

@@ -0,0 +1,3 @@
+- (#6821) Fix factory transfer when units that are meant to be rebuilt are restricted by HQ requirements.
+
+- (#6821) Fix transfer in the middle of factory transfer causing rebuilding to fail.

lua/SimUtils.lua

@@ -19,16 +19,41 @@ local buildersCategory = categories.ALLUNITS - categories.CONSTRUCTION - categor
 
 ---@alias FactoryRebuildDataTable table<UnitId, (FactoryUnit | FactoryRebuildData)[]>
 
+--- Clear data for a factory so transferring it doesn't try to rebuild units again
+---@param factory FactoryUnit | FactoryRebuildData
+local function clearFactoryRebuildData(factory)
+    factory.FacRebuild_UnitId = nil
+    factory.FacRebuild_Progress = nil
+    factory.FacRebuild_BuildTime = nil
+    factory.FacRebuild_Health = nil
+    factory.FacRebuild_OldBuildRate = nil
+end
+
 ---@param factoryRebuildDataTable FactoryRebuildDataTable
 function FactoryRebuildUnits(factoryRebuildDataTable)
     for buildUnitId, factories in factoryRebuildDataTable do
+        -- Remove support factories that can't build their unit due to lacking an HQ
+        local noFactories = false
+        for i, factory in factories do
+            if not factory:CanBuild(buildUnitId) then
+                clearFactoryRebuildData(factory)
+                factories[i] = nil
+                if table.empty(factories) then
+                    factoryRebuildDataTable[buildUnitId] = nil
+                    noFactories = true
+                end
+                continue
+            end
+        end
+        if noFactories then continue end
+
         IssueClearCommands(factories)
         IssueBuildFactory(factories, buildUnitId, 1)
     end
     -- wait for build order to start and then rebuild the units for free
     WaitTicks(1)
-    for k, factories in pairs(factoryRebuildDataTable) do
-        for i, factory in pairs(factories) do
+    for k, factories in factoryRebuildDataTable do
+        for i, factory in factories do
             if factory.Dead then
                 factories[i] = nil
                 if table.empty(factories) then
@@ -45,8 +70,8 @@ function FactoryRebuildUnits(factoryRebuildDataTable)
     end
     -- wait for buildpower to apply then return the factories to normal and pause them
     WaitTicks(1)
-    for k, factories in pairs(factoryRebuildDataTable) do
-        for i, factory in pairs(factories) do
+    for k, factories in factoryRebuildDataTable do
+        for i, factory in factories do
             if factory.Dead then
                 factories[i] = nil
                 if table.empty(factories) then
@@ -97,11 +122,7 @@ Health: %f
                 rebuiltUnit:SetHealth(nil, factory.FacRebuild_Health)
             end
 
-            -- clean up after the rebuilding
-            factory.FacRebuild_Progress = nil
-            factory.FacRebuild_BuildTime = nil
-            factory.FacRebuild_Health = nil
-            factory.FacRebuild_OldBuildRate = nil
+            clearFactoryRebuildData(factory)
         end
     end
 end
@@ -185,10 +206,10 @@ function TransferUnitsOwnership(units, toArmy, captured, noRestrictions)
         local defaultBuildRate
         local upgradeBuildTimeComplete
         local exclude
-        local FacRebuild_UnitId
-        local FacRebuild_Progress
-        local FacRebuild_BuildTime
-        local FacRebuild_Health
+        local FacRebuild_UnitId = unit.FacRebuild_UnitId
+        local FacRebuild_Progress = unit.FacRebuild_Progress
+        local FacRebuild_BuildTime = unit.FacRebuild_BuildTime
+        local FacRebuild_Health = unit.FacRebuild_Health
 
         local shield = unit.MyShield
         if shield then
@@ -380,6 +401,9 @@ function TransferUnitsOwnership(units, toArmy, captured, noRestrictions)
             else
                 table.insert(data, newFactoryUnit)
             end
+            -- store data for rebuilding
+            -- unit id is not needed during rebuild but is needed if transferred again in the middle of rebuild
+            newFactoryUnit.FacRebuild_UnitId = FacRebuild_UnitId
             newFactoryUnit.FacRebuild_Progress = FacRebuild_Progress
             newFactoryUnit.FacRebuild_BuildTime = FacRebuild_BuildTime
             newFactoryUnit.FacRebuild_Health = FacRebuild_Health