From 50498520794b729aac5b4ffaa44bb04a14ec3dc3 Mon Sep 17 00:00:00 2001
From: OldManAlpha <60587722+OldManAlpha@users.noreply.github.com>
Date: Tue, 2 Sep 2025 16:44:07 -0700
Subject: [PATCH] Fix losing custom units in GUID roster

---
 Puppeteer.lua       | 34 ++++++++++++++++++++++++++--------
 core/Command.lua    |  4 ++--
 libs/GuidRoster.lua |  7 ++++++-
 libs/PTUnit.lua     |  3 ++-
 4 files changed, 36 insertions(+), 12 deletions(-)

diff --git a/Puppeteer.lua b/Puppeteer.lua
index ef91d63..ad5230f 100644
--- a/Puppeteer.lua
+++ b/Puppeteer.lua
@@ -262,18 +262,36 @@ function OnAddonLoaded()
         -- Older versions of SuperWoW had an issue where units that aren't part of normal units wouldn't receive events,
         -- so updates are done manually
         local needsManualUpdates = util.SuperWoWFeatureLevel < util.SuperWoW_v1_4
+        local existing = {}
+        local nextCleanup = GetTime() + 10
         customUnitUpdater:SetScript("OnUpdate", function()
             if GetTime() > nextUpdate then
                 nextUpdate = GetTime() + 0.25
 
-                for unit, guid in pairs(CustomUnitGUIDMap) do
-                    if needsManualUpdates or not UnitExists(guid) then
-                        PTUnit.Get(unit):UpdateAuras()
-                        for ui in UnitFrames(unit) do
-                            ui:UpdateHealth()
-                            ui:UpdatePower()
-                            ui:UpdateAuras()
-                            ui:UpdateIncomingHealing()
+                for guid, units in pairs(GUIDCustomUnitMap) do
+                    local exists = UnitExists(guid) == 1
+                    local needsUpdate = false
+                    if (existing[guid] ~= nil) ~= exists then
+                        existing[guid] = exists or nil
+                        needsUpdate = true
+                    end
+
+                    if needsManualUpdates or needsUpdate then
+                        PTUnit.Get(guid):UpdateAuras()
+                        for _, unit in ipairs(units) do
+                            for ui in UnitFrames(unit) do
+                                ui:UpdateAll()
+                                ui:UpdateIncomingHealing()
+                            end
+                        end
+                    end
+                end
+
+                if GetTime() > nextCleanup then
+                    nextCleanup = GetTime() + 10
+                    for guid in pairs(existing) do
+                        if not GUIDCustomUnitMap[guid] then
+                            existing[guid] = nil
                         end
                     end
                 end
diff --git a/core/Command.lua b/core/Command.lua
index 94e4fb9..1dcaad2 100644
--- a/core/Command.lua
+++ b/core/Command.lua
@@ -16,7 +16,7 @@ SlashCmdList["PUPPETEER"] = function(args)
     elseif args == "check" then
         Puppeteer.CheckGroup()
     elseif args == "update" then
-        for _, ui in pairs(Puppeteer.AllUnitFrames) do
+        for _, ui in ipairs(Puppeteer.AllUnitFrames) do
             ui:SizeElements()
             ui:UpdateAll()
         end
@@ -28,7 +28,7 @@ SlashCmdList["PUPPETEER"] = function(args)
         PTOptions.TestUI = not PTOptions.TestUI
         Puppeteer.TestUI = PTOptions.TestUI
         if PTOptions.TestUI then
-            for _, ui in pairs(Puppeteer.AllUnitFrames) do
+            for _, ui in ipairs(Puppeteer.AllUnitFrames) do
                 ui.fakeStats = ui.GenerateFakeStats()
                 ui:Show()
             end
diff --git a/libs/GuidRoster.lua b/libs/GuidRoster.lua
index 97efc94..b603dea 100644
--- a/libs/GuidRoster.lua
+++ b/libs/GuidRoster.lua
@@ -24,12 +24,17 @@ function ResetRoster()
 end
 
 function PopulateRoster()
-    for _, unit in ipairs(util.AllUnits) do
+    for _, unit in ipairs(util.AllRealUnits) do
         local exists, guid = UnitExists(unit)
         if exists then
             AddUnit(guid, unit)
         end
     end
+    for guid, units in pairs(PTUnitProxy.GUIDCustomUnitMap) do
+        for _, unit in ipairs(units) do
+            AddUnit(guid, unit)
+        end
+    end
 end
 
 function AddUnit(guid, unit)
diff --git a/libs/PTUnit.lua b/libs/PTUnit.lua
index 3319b69..8525a83 100644
--- a/libs/PTUnit.lua
+++ b/libs/PTUnit.lua
@@ -8,6 +8,7 @@ local _G = getfenv(0)
 local util = PTUtil
 local GetAuraInfo = util.GetAuraInfo
 local AllUnits = util.AllUnits
+local AllRealUnits = util.AllRealUnits
 local AllUnitsSet = util.AllUnitsSet
 local superwow = util.IsSuperWowPresent()
 local canGetAuraIDs = util.CanClientGetAuraIDs()
@@ -54,7 +55,7 @@ end
 function UpdateGuidCaches()
     local cached = PTUnit.Cached
     local prevCached = PTUtil.CloneTableCompost(cached)
-    for _, unit in ipairs(AllUnits) do
+    for _, unit in ipairs(AllRealUnits) do
         local exists, guid = UnitExists(unit)
         if exists then
             if not cached[guid] then