Fix transfer of manually built enhancements

changelog/snippets/fix.7162.md

@@ -0,0 +1 @@
+- (#7162) Fix transfer of units with manually built enhancements.

engine/Sim/CAiBrain.lua

@@ -542,4 +542,12 @@ end
 function CAiBrain:TakeResource(type, amount)
 end
 
+--- Called by the engine when a unit fails unit transfer due to issues other than unit cap.
+---@type fun(self: moho.aibrain_methods)
+CAiBrain.OnFailedUnitTransfer = nil
+
+--- Called by the engine when a unit fails to create due to unit cap.
+---@type fun(self: moho.aibrain_methods)
+CAiBrain.OnUnitCapLimitReached = nil
+
 return CAiBrain

lua/SimUtils.lua

@@ -5,6 +5,8 @@
 -- upvalues for performance
 local ArmyBrains = ArmyBrains
 local GetCurrentCommandSource = GetCurrentCommandSource
+local TableEmpty = table.empty
+local KillThread = KillThread
 
 ------------------------------------------------------------------------------------------------------------------------
 --#region General Unit Transfer Scripts
@@ -326,11 +328,8 @@ function TransferUnitsOwnership(units, toArmy, captured, noRestrictions)
             local unitEnh = SimUnitEnhancements[unit.EntityId]
             if unitEnh then
                 activeEnhancements = {}
-                for i, enh in unitEnh do
-                    activeEnhancements[i] = enh
-                end
-                if not activeEnhancements[1] then
-                    activeEnhancements = nil
+                for slot, enh in unitEnh do
+                    activeEnhancements[slot] = enh
                 end
             end
         end
@@ -420,8 +419,20 @@ function TransferUnitsOwnership(units, toArmy, captured, noRestrictions)
         end
 
         if activeEnhancements then
-            for _, enh in activeEnhancements do
-                newUnit:CreateEnhancement(enh)
+            local thread = newUnit.OnStopBeingBuiltEnhancementsThread
+            if thread then KillThread(thread) end
+            if TableEmpty(activeEnhancements) then
+                newUnit.OnStopBeingBuiltEnhancementsThread = nil
+            else
+                newUnit.OnStopBeingBuiltEnhancementsThread = newUnit:ForkThread(function()
+                    WaitTicks(1)
+                    if newUnit and not newUnit.Dead and not IsDestroyed(newUnit) then
+                        for _, enh in activeEnhancements do
+                            newUnit:CreateEnhancement(enh)
+                        end
+                    end
+                    newUnit.OnStopBeingBuiltEnhancementsThread = nil
+                end)
             end
         end
 
@@ -738,6 +749,7 @@ function GiveUnitsToPlayer(data, units)
     if manualShare == 'none' or table.empty(units) then
         return
     end
+    ---@cast units -nil
     local toArmy = data.To
     local owner = units[1].Army
     if OkayToMessWithArmy(owner) and IsAlly(owner, toArmy) then

lua/defaultcomponents.lua

@@ -1,6 +1,8 @@
 local Buff = import("/lua/sim/buff.lua")
 local Entity = import("/lua/sim/entity.lua").Entity
 
+local IsDestroyed = IsDestroyed
+
 ---@class ShieldEffectsComponent : Unit
 ---@field Trash TrashBag
 ---@field ShieldEffectsBag TrashBag
@@ -251,6 +253,7 @@ IntelComponent = ClassSimple {
 
             --- display progress
             for k = 1, ticks do
+                if self.Dead or IsDestroyed(self) then return end
 
                 -- prevent changing work progress when we are doing work (such as an enhancement)
                 if not self.WorkItem then

lua/sim/Unit.lua

@@ -175,6 +175,7 @@ local cUnitGetBuildRate = cUnit.GetBuildRate
 ---@field ImmuneToStun? boolean
 ---@field Anims? Animator[] # Animators that get stopped when a unit is stunned. Not used in FAF.
 ---@field IsBeingTransferred? boolean
+---@field OnStopBeingBuiltEnhancementsThread thread?
 Unit = ClassUnit(moho.unit_methods, IntelComponent, VeterancyComponent, DebugUnitComponent, FastDecayComponent) {
 
     IsUnit = true,
@@ -2490,7 +2491,7 @@ Unit = ClassUnit(moho.unit_methods, IntelComponent, VeterancyComponent, DebugUni
         end
 
         if bp.EnhancementPresetAssigned then
-            self:ForkThread(self.CreatePresetEnhancementsThread)
+            self.OnStopBeingBuiltEnhancementsThread = self:ForkThread(self.CreatePresetEnhancementsThread)
         end
 
         -- Don't try sending a Notify message from here if we're an ACU
@@ -2614,10 +2615,7 @@ Unit = ClassUnit(moho.unit_methods, IntelComponent, VeterancyComponent, DebugUni
         local bp = self.Blueprint
         if bp.Enhancements and bp.EnhancementPresetAssigned and bp.EnhancementPresetAssigned.Enhancements then
             for k, v in bp.EnhancementPresetAssigned.Enhancements do
-                -- Enhancements may already have been created by SimUtils.TransferUnitsOwnership
-                if not self:HasEnhancement(v) then
-                    self:CreateEnhancement(v)
-                end
+                self:CreateEnhancement(v)
             end
         end
     end,
@@ -2630,6 +2628,7 @@ Unit = ClassUnit(moho.unit_methods, IntelComponent, VeterancyComponent, DebugUni
         if self and not self.Dead then
             self:CreatePresetEnhancements()
         end
+        self.OnStopBeingBuiltEnhancementsThread = nil
     end,
 
     ---@param self Unit
@@ -3324,6 +3323,7 @@ Unit = ClassUnit(moho.unit_methods, IntelComponent, VeterancyComponent, DebugUni
         end
 
         self:RequestRefreshUI()
+        return true
     end,
 
     ---@param self Unit