StoreUnitColor

library StoreUnitColor/*requieres a unit indexer (this uses the Bribe's unit indexer)*/

    //! novjass
    /*This library is to store the actual color of a unit and also if is changed (Except if you use */SetPlayerColorBJEnum/*).
    I recommend use better the function */SetPlayerColorEx/* instead of */SetPlayerColorBJ/* because this last have a reference leak.
    Also add a constant */PLAYER_COLOR_BLACK/* that you can use and a function */GetUnitColor/* to have the saved unit color.*/
    //! endnovjass
  
    globals
        constant playercolor PLAYER_COLOR_BLACK = ConvertPlayerColor(PLAYER_NEUTRAL_AGGRESSIVE)
    endglobals
  
    private module M
  
        static playercolor array UnitColor
  
        static method StoreColorEnum takes nothing returns nothing
            set thistype.UnitColor[GetUnitUserData(GetEnumUnit())] = bj_setPlayerTargetColor
        endmethod
  
        static method StoreColorIfPlayerChange takes player whichPlayer,  playercolor color,  boolean changeExisting returns nothing
            if changeExisting then
                set bj_setPlayerTargetColor = color
                set bj_lastCreatedGroup = CreateGroup()
                call GroupEnumUnitsOfPlayer(bj_lastCreatedGroup, whichPlayer, null)
                call ForGroup(bj_lastCreatedGroup, function thistype.StoreColorEnum)
                call DestroyGroup(bj_lastCreatedGroup)
            endif
        endmethod
  
        static method StoreColorIfUnitChange takes unit whichUnit,  playercolor color returns nothing
            set thistype.UnitColor[GetUnitUserData(whichUnit)] = color
        endmethod
  
        static method StoreColorIfUnitChangeOwner takes unit whichUnit,  player whichPlayer,  boolean changeColor returns nothing
            if changeColor then
                set thistype.UnitColor[GetUnitUserData(whichUnit)] = GetPlayerColor(whichPlayer)
            endif
        endmethod
      
        //If you use another unit indexer,  edit this
        static method SaveColor takes nothing returns nothing
            set UnitColor[udg_UDex] = GetPlayerColor(GetOwningPlayer(udg_UDexUnits[udg_UDex]))
        endmethod
      
        static method onInit takes nothing returns nothing
            local trigger t = CreateTrigger()
            call TriggerRegisterVariableEvent(t, "udg_UnitIndexEvent", EQUAL, 1.00)
            call TriggerAddAction(t, function thistype.SaveColor)
            set t = null
        endmethod
        //-- --
      
    endmodule
  
    private struct S extends array
        implement M
    endstruct
  
    hook SetPlayerColorBJ S.StoreColorIfPlayerChange
    hook SetUnitColor S.StoreColorIfUnitChange
    hook SetUnitOwner S.StoreColorIfUnitChangeOwner
      
    function GetUnitColor takes unit whichUnit returns playercolor
        return S.UnitColor[GetUnitUserData(whichUnit)]
    endfunction
  
    function SetPlayerColorEx takes player whichPlayer,  playercolor color,  boolean changeExisting returns nothing
        local group g
        local unit u
        call SetPlayerColor(whichPlayer, color)
        if changeExisting then
            set g = CreateGroup()
            call GroupEnumUnitsOfPlayer(g, whichPlayer, null)
            loop
                set u = FirstOfGroup(g)
                exitwhen u == null
                call SetUnitColor(u, color)
                call GroupRemoveUnit(g, u)
            endloop
            call DestroyGroup(g)
        endif
        set g = null
    endfunction
endlibrary

OnInit(function () -- https://www.hiveworkshop.com/threads/global-initialization.317099/
    local AddHook = Require.optional "AddHook"

    --This library is to store the actual color of a unit and also if is changed.
    --Also add a constant PLAYER_COLOR_BLACK that you can use and a function GetUnitColor to have the saved unit color.

    PLAYER_COLOR_BLACK = ConvertPlayerColor(PLAYER_NEUTRAL_AGGRESSIVE)

    local UnitColor = {}

    local re = CreateRegion()
    local r = GetWorldBounds()
    RegionAddRect(re, r)
    local b = Filter(function()
        UnitColor[GetFilterUnit()] = GetPlayerColor(GetOwningPlayer(GetFilterUnit()))
        return false
    end)
    TriggerRegisterEnterRegion(CreateTrigger(), re, b)
    for i = 0, bj_MAX_PLAYER_SLOTS do
        GroupEnumUnitsOfPlayer(bj_lastCreatedGroup, Player(i), b)
    end
    RemoveRect(r)

    --I just had to change these 2 functions.

    if AddHook then
        AddHook("SetUnitColor", function (whichUnit, color)
            UnitColor[whichUnit] = color
            SetUnitColor.old(whichUnit, color)
        end)
        AddHook("SetUnitOwner", function (whichUnit, whichPlayer, changeColor)
            if changeColor then
                UnitColor[whichUnit] = GetPlayerColor(whichPlayer)
            end
            SetUnitOwner.old(whichUnit, whichPlayer, changeColor)
        end)
    else
        local OldSetUnitColor = SetUnitColor
        function SetUnitColor(whichUnit, color)
            UnitColor[whichUnit] = color
            OldSetUnitColor(whichUnit, color)
        end

        local OldSetUnitOwner = SetUnitOwner
        function SetUnitOwner(whichUnit, whichPlayer, changeColor)
            if changeColor then
                UnitColor[whichUnit] = GetPlayerColor(whichPlayer)
            end
            OldSetUnitOwner(whichUnit, whichPlayer, changeColor)
        end
    end

    function GetUnitColor(whichUnit)
        return UnitColor[whichUnit]
    end
end)