From 2d332d0a130f63e3d2e6c02773f7876389d82882 Mon Sep 17 00:00:00 2001
From: Gustas <gustas.kazukauskas@gmail.com>
Date: Wed, 15 Nov 2023 20:47:34 +0200
Subject: [PATCH] Fix pillbox not uncloaking upon firing

---
 OpenRA.Mods.Common/Traits/Armament.cs         | 26 ++++++++++++-----
 .../Traits/Attack/AttackGarrisoned.cs         | 29 ++++++++++++++++---
 2 files changed, 43 insertions(+), 12 deletions(-)

diff --git a/OpenRA.Mods.Common/Traits/Armament.cs b/OpenRA.Mods.Common/Traits/Armament.cs
index d9ec8fb65f..283b6b677f 100644
--- a/OpenRA.Mods.Common/Traits/Armament.cs
+++ b/OpenRA.Mods.Common/Traits/Armament.cs
@@ -118,7 +118,7 @@ namespace OpenRA.Mods.Common.Traits
 
 		BodyOrientation coords;
 		INotifyBurstComplete[] notifyBurstComplete;
-		INotifyAttack[] notifyAttacks;
+		List<(Actor NotifyActor, INotifyAttack Notify)> notifyAttacks;
 
 		int conditionToken = Actor.InvalidConditionToken;
 
@@ -174,7 +174,7 @@ namespace OpenRA.Mods.Common.Traits
 			hovers = self.TraitOrDefault<Hovers>();
 			coords = self.Trait<BodyOrientation>();
 			notifyBurstComplete = self.TraitsImplementing<INotifyBurstComplete>().ToArray();
-			notifyAttacks = self.TraitsImplementing<INotifyAttack>().ToArray();
+			notifyAttacks = self.TraitsImplementing<INotifyAttack>().Select(a => (self, a)).ToList();
 
 			rangeModifiers = self.TraitsImplementing<IRangeModifier>().ToArray().Select(m => m.GetRangeModifier());
 			reloadModifiers = self.TraitsImplementing<IReloadModifier>().ToArray().Select(m => m.GetReloadModifier());
@@ -184,6 +184,16 @@ namespace OpenRA.Mods.Common.Traits
 			base.Created(self);
 		}
 
+		public void AddNotifyAttacks(Actor attacker, INotifyAttack[] notifyAttacks)
+		{
+			this.notifyAttacks.AddRange(notifyAttacks.Select(a => (attacker, a)));
+		}
+
+		public void RemoveNotifyAttacks(INotifyAttack[] notifyAttacks)
+		{
+			this.notifyAttacks.RemoveAll(pair => notifyAttacks.Any(notify => notify == pair.Notify));
+		}
+
 		void UpdateCondition(Actor self)
 		{
 			if (string.IsNullOrEmpty(Info.ReloadingCondition))
@@ -278,8 +288,8 @@ namespace OpenRA.Mods.Common.Traits
 				UpdateBurst(self, target);
 
 				if (notifyAttacking)
-					foreach (var notify in notifyAttacks)
-						notify.Attacking(self, target, this, barrel);
+					foreach (var (notifyActor, notify) in notifyAttacks)
+						notify.Attacking(notifyActor, target, this, barrel);
 			}
 			while (FireDelay == 0 && CanFire(self, target));
 
@@ -288,8 +298,8 @@ namespace OpenRA.Mods.Common.Traits
 
 		protected virtual void FireBarrel(Actor self, IFacing facing, in Target target, Barrel barrel)
 		{
-			foreach (var na in notifyAttacks)
-				na.PreparingAttack(self, target, this, barrel);
+			foreach (var (notifyActor, notify) in notifyAttacks)
+				notify.PreparingAttack(notifyActor, target, this, barrel);
 
 			WPos MuzzlePosition() => self.CenterPosition + MuzzleOffset(self, barrel);
 			WAngle MuzzleFacing() => MuzzleOrientation(self, barrel).Yaw;
@@ -347,8 +357,8 @@ namespace OpenRA.Mods.Common.Traits
 					if (burst == args.Weapon.Burst && args.Weapon.StartBurstReport != null && args.Weapon.StartBurstReport.Length > 0)
 						Game.Sound.Play(SoundType.World, args.Weapon.StartBurstReport, self.World, self.CenterPosition);
 
-					foreach (var na in notifyAttacks)
-						na.Attacking(self, delayedTarget, this, barrel);
+					foreach (var (notifyActor, notify) in notifyAttacks)
+						notify.Attacking(notifyActor, delayedTarget, this, barrel);
 
 					Recoil = Info.Recoil;
 				}
diff --git a/OpenRA.Mods.Common/Traits/Attack/AttackGarrisoned.cs b/OpenRA.Mods.Common/Traits/Attack/AttackGarrisoned.cs
index 5a82a3a1db..3bf38c9ed2 100644
--- a/OpenRA.Mods.Common/Traits/Attack/AttackGarrisoned.cs
+++ b/OpenRA.Mods.Common/Traits/Attack/AttackGarrisoned.cs
@@ -78,6 +78,7 @@ namespace OpenRA.Mods.Common.Traits
 	public class AttackGarrisoned : AttackFollow, INotifyPassengerEntered, INotifyPassengerExited, IRender
 	{
 		public new readonly AttackGarrisonedInfo Info;
+		INotifyAttack[] notifyAttacks;
 		readonly Lazy<BodyOrientation> coords;
 		readonly List<Armament> armaments;
 		readonly List<AnimationWithOffset> muzzles;
@@ -97,6 +98,12 @@ namespace OpenRA.Mods.Common.Traits
 			paxRender = new Dictionary<Actor, RenderSprites>();
 		}
 
+		protected override void Created(Actor self)
+		{
+			notifyAttacks = self.TraitsImplementing<INotifyAttack>().ToArray();
+			base.Created(self);
+		}
+
 		protected override Func<IEnumerable<Armament>> InitializeGetArmaments(Actor self)
 		{
 			return () => armaments;
@@ -107,9 +114,15 @@ namespace OpenRA.Mods.Common.Traits
 			paxFacing.Add(passenger, passenger.Trait<IFacing>());
 			paxPos.Add(passenger, passenger.Trait<IPositionable>());
 			paxRender.Add(passenger, passenger.Trait<RenderSprites>());
-			armaments.AddRange(
-				passenger.TraitsImplementing<Armament>()
-				.Where(a => Info.Armaments.Contains(a.Info.Name)));
+
+			foreach (var a in passenger.TraitsImplementing<Armament>())
+			{
+				if (Info.Armaments.Contains(a.Info.Name))
+				{
+					a.AddNotifyAttacks(self, notifyAttacks);
+					armaments.Add(a);
+				}
+			}
 		}
 
 		void INotifyPassengerExited.OnPassengerExited(Actor self, Actor passenger)
@@ -117,7 +130,15 @@ namespace OpenRA.Mods.Common.Traits
 			paxFacing.Remove(passenger);
 			paxPos.Remove(passenger);
 			paxRender.Remove(passenger);
-			armaments.RemoveAll(a => a.Actor == passenger);
+
+			foreach (var a in armaments.ToList())
+			{
+				if (a.Actor == passenger)
+				{
+					a.RemoveNotifyAttacks(notifyAttacks);
+					armaments.Remove(a);
+				}
+			}
 		}
 
 		FirePort SelectFirePort(Actor self, WAngle targetYaw)