From 8e726f413e62e746daeeedecfcaa4108bfee95c7 Mon Sep 17 00:00:00 2001
From: lL1l1 <82986251+lL1l1@users.noreply.github.com>
Date: Tue, 14 May 2024 18:16:52 -0700
Subject: [PATCH] Fix wrong splash radius decal appearing for certain units
 (#6028)

---
 changelog/snippets/fix.6028.md |  3 ++
 lua/ui/controls/worldview.lua  | 68 +++++++++++++++++++---------------
 2 files changed, 42 insertions(+), 29 deletions(-)
 create mode 100644 changelog/snippets/fix.6028.md

diff --git a/changelog/snippets/fix.6028.md b/changelog/snippets/fix.6028.md
new file mode 100644
index 0000000000..75edc26668
--- /dev/null
+++ b/changelog/snippets/fix.6028.md
@@ -0,0 +1,3 @@
+- (#6028) Fix attack cursor splash radius decals appearing with the wrong size for units that display splash radius and have multiple weapons.
+
+In FAF, this affects the Ythotha, Mongoose, and Soul Ripper.
diff --git a/lua/ui/controls/worldview.lua b/lua/ui/controls/worldview.lua
index 2d730c4c12..a4b9dff618 100644
--- a/lua/ui/controls/worldview.lua
+++ b/lua/ui/controls/worldview.lua
@@ -32,7 +32,7 @@ local KeyCodeAlt = 18
 local KeyCodeCtrl = 17
 local KeyCodeShift = 16
 
-local weaponsCached = { }
+local unitsToWeaponsCached = { }
 
 ---@class Renderable : Destroyable
 ---@field OnRender fun(self:Renderable, worldView:WorldView)
@@ -41,52 +41,57 @@ local weaponsCached = { }
 ---@field texture string
 ---@field scale number
 
--- If all selected units with the SHOWATTACKRETICLE flag set are of the same type, return the weapon
--- table from their blueprint. Otherwise returns null.
-
---- Returns all weapon blueprints that match the predicate of units with the `SHOWATTACKRETICLE` category set
+--- Returns all unique weapon blueprints that match the predicate, and are from units with the `SHOWATTACKRETICLE` category set
 ---@param predicate function<WeaponBlueprint>
----@return WeaponBlueprint[]
+---@return table<UnitId, WeaponBlueprint[] | false> unitsToWeapons
 local function GetSelectedWeaponsWithReticules(predicate)
     local selectedUnits = GetSelectedUnits()
 
     -- clear out the cache
-    local weapons = weaponsCached
-    for k, other in weapons do
-        weapons[k] = false
+    local unitsToWeapons = unitsToWeaponsCached
+    for k, other in unitsToWeapons do
+        unitsToWeapons[k] = false
     end
 
     -- find valid units
     if selectedUnits then
         for i, u in selectedUnits do
             local bp = u:GetBlueprint()
-            if bp.CategoriesHash['SHOWATTACKRETICLE'] and (not weapons[bp.BlueprintId]) then
-                for k, v in bp.Weapon do
-                    if predicate(v) then
-                        weapons[bp.BlueprintId] = v
-                        break
+            local bpId = bp.BlueprintId
+            if bp.CategoriesHash['SHOWATTACKRETICLE'] and (not unitsToWeapons[bpId]) then
+                for _, wep in bp.Weapon do
+                    if predicate(wep) then
+                        if not unitsToWeapons[bpId] then
+                            unitsToWeapons[bpId] = { wep }
+                        else
+                            table.insert(unitsToWeapons[bpId], wep)
+                        end
                     end
                 end
             end
         end
     end
 
-    return weapons
+    return unitsToWeapons
 end
 
 --- A generic decal texture / size computation function that uses the damage or spread radius
 ---@param predicate function<WeaponBlueprint>
 ---@return WorldViewDecalData[]
 local function RadiusDecalFunction(predicate)
-    local weapons = GetSelectedWeaponsWithReticules(predicate)
+    local unitsToWeapons = GetSelectedWeaponsWithReticules(predicate)
 
     -- The maximum damage radius of a selected missile weapon.
     local maxRadius = 0
-    for _, w in weapons do
-        if w.FixedSpreadRadius and w.FixedSpreadRadius > maxRadius then
-            maxRadius = w.FixedSpreadRadius
-        elseif w.DamageRadius > maxRadius then
-            maxRadius = w.DamageRadius
+    for _, weapons in unitsToWeapons do
+        if weapons then
+            for _, w in weapons do
+                if w.FixedSpreadRadius and w.FixedSpreadRadius + w.DamageRadius > maxRadius then
+                    maxRadius = w.FixedSpreadRadius + w.DamageRadius
+                elseif w.DamageRadius > maxRadius then
+                    maxRadius = w.DamageRadius
+                end
+            end
         end
     end
 
@@ -105,7 +110,7 @@ end
 --- A decal texture / size computation function for `RULEUCC_Nuke`
 ---@return WorldViewDecalData[]
 local function NukeDecalFunc()
-    local weapons = GetSelectedWeaponsWithReticules(
+    local unitsToWeapons = GetSelectedWeaponsWithReticules(
         function(w)
             return w.NukeWeapon
         end
@@ -113,13 +118,17 @@ local function NukeDecalFunc()
 
     local inner = 0
     local outer = 0
-    for _, w in weapons do
-        if w.NukeOuterRingRadius > outer then
-            outer = w.NukeOuterRingRadius
-        end
+    for _, weapons in unitsToWeapons do
+        if weapons then
+            for _, w in weapons do
+                if w.NukeOuterRingRadius > outer then
+                    outer = w.NukeOuterRingRadius
+                end
 
-        if w.NukeInnerRingRadius > inner then
-            inner = w.NukeInnerRingRadius
+                if w.NukeInnerRingRadius > inner then
+                    inner = w.NukeInnerRingRadius
+                end
+            end
         end
     end
 
@@ -151,7 +160,8 @@ end
 local function AttackDecalFunc(mode)
     return RadiusDecalFunction(
         function(w)
-            return w.ManualFire == false and w.WeaponCategory ~= 'Teleport' and w.WeaponCategory ~= "Death" and w.WeaponCategory ~= "Anti Air"
+            return w.ManualFire == false and w.WeaponCategory ~= 'Teleport' and w.WeaponCategory ~= "Death"
+                and w.WeaponCategory ~= "Anti Air" and w.DamageType ~= 'Overcharge' and not w.EnabledByEnhancement
         end
 )
 end