Merge INotifyRearm/Repair into INotifyResupply

And streamline its notify methods.
Also cache INotifyResupply traits at beginning
of Resupply activity.

OpenRA.Mods.Common/Activities/Resupply.cs

@@ -26,12 +26,11 @@ namespace OpenRA.Mods.Common.Activities
 		readonly Repairable repairable;
 		readonly RepairableNear repairableNear;
 		readonly Rearmable rearmable;
+		readonly INotifyResupply[] notifyResupplies;
 
 		int remainingTicks;
 		bool played;
-		bool paused;
+		ResupplyType activeResupplyTypes = ResupplyType.None;
-		bool repairComplete;
-		bool rearmComplete;
 
 		public Resupply(Actor self, Actor host, WDist closeEnough)
 		{
@@ -42,13 +41,19 @@ namespace OpenRA.Mods.Common.Activities
 			repairable = self.TraitOrDefault<Repairable>();
 			repairableNear = self.TraitOrDefault<RepairableNear>();
 			rearmable = self.TraitOrDefault<Rearmable>();
+			notifyResupplies = host.TraitsImplementing<INotifyResupply>().ToArray();
 
-			repairComplete = health == null || health.DamageState == DamageState.Undamaged
+			var cannotRepairAtHost = health == null || health.DamageState == DamageState.Undamaged
 				|| !allRepairsUnits.Any()
 				|| ((repairable == null || !repairable.Info.RepairActors.Contains(host.Info.Name))
 					&& (repairableNear == null || !repairableNear.Info.RepairActors.Contains(host.Info.Name)));
 
-			rearmComplete = rearmable == null || !rearmable.Info.RearmActors.Contains(host.Info.Name) || rearmable.RearmableAmmoPools.All(p => p.FullAmmo());
+			if (!cannotRepairAtHost)
+				activeResupplyTypes |= ResupplyType.Repair;
+
+			var cannotRearmAtHost = rearmable == null || !rearmable.Info.RearmActors.Contains(host.Info.Name) || rearmable.RearmableAmmoPools.All(p => p.FullAmmo());
+			if (!cannotRearmAtHost)
+				activeResupplyTypes |= ResupplyType.Rearm;
 		}
 
 		protected override void OnFirstRun(Actor self)
@@ -56,41 +61,39 @@ namespace OpenRA.Mods.Common.Activities
 			if (host.Type == TargetType.Invalid)
 				return;
 
-			if (!repairComplete)
+			if (activeResupplyTypes > 0)
-				foreach (var notifyRepair in host.Actor.TraitsImplementing<INotifyRepair>())
+				foreach (var notifyResupply in notifyResupplies)
-					notifyRepair.BeforeRepair(host.Actor, self);
+					notifyResupply.BeforeResupply(host.Actor, self, activeResupplyTypes);
-
-			if (!rearmComplete)
-			{
-				foreach (var notifyRearm in host.Actor.TraitsImplementing<INotifyRearm>())
-					notifyRearm.RearmingStarted(host.Actor, self);
 
 			// Reset the ReloadDelay to avoid any issues with early cancellation
 			// from previous reload attempts (explicit order, host building died, etc).
 			// HACK: this really shouldn't be managed from here
+			if (activeResupplyTypes.HasFlag(ResupplyType.Rearm))
 				foreach (var pool in rearmable.RearmableAmmoPools)
 					pool.RemainingTicks = pool.Info.ReloadDelay;
 		}
-		}
 
 		public override Activity Tick(Actor self)
 		{
 			if (IsCanceling)
 				return NextActivity;
 
-			if (host.Type == TargetType.Invalid || health == null)
+			if (host.Type == TargetType.Invalid)
 				return NextActivity;
 
 			if (closeEnough.LengthSquared > 0 && !host.IsInRange(self.CenterPosition, closeEnough))
 				return NextActivity;
 
-			if (!repairComplete)
+			if (activeResupplyTypes.HasFlag(ResupplyType.Repair))
 				RepairTick(self);
 
-			if (!rearmComplete)
+			if (activeResupplyTypes.HasFlag(ResupplyType.Rearm))
 				RearmTick(self);
 
-			if (repairComplete && rearmComplete)
+			foreach (var notifyResupply in notifyResupplies)
+				notifyResupply.ResupplyTick(host.Actor, self, activeResupplyTypes);
+
+			if (activeResupplyTypes == 0)
 				return NextActivity;
 
 			return this;
@@ -99,26 +102,11 @@ namespace OpenRA.Mods.Common.Activities
 		void RepairTick(Actor self)
 		{
 			// First active.
-			RepairsUnits repairsUnits = null;
+			var repairsUnits = allRepairsUnits.FirstOrDefault(r => !r.IsTraitDisabled && !r.IsTraitPaused);
-			paused = false;
-			foreach (var r in allRepairsUnits)
-			{
-				if (!r.IsTraitDisabled)
-				{
-					if (r.IsTraitPaused)
-						paused = true;
-					else
-					{
-						repairsUnits = r;
-						break;
-					}
-				}
-			}
-
 			if (repairsUnits == null)
 			{
-				if (!paused)
+				if (!allRepairsUnits.Any(r => r.IsTraitPaused))
-					repairComplete = true;
+					activeResupplyTypes &= ~ResupplyType.Repair;
 
 				return;
 			}
@@ -134,10 +122,7 @@ namespace OpenRA.Mods.Common.Activities
 
 				Game.Sound.PlayNotification(self.World.Map.Rules, self.Owner, "Speech", repairsUnits.Info.FinishRepairingNotification, self.Owner.Faction.InternalName);
 
-				foreach (var notifyRepair in host.Actor.TraitsImplementing<INotifyRepair>())
+				activeResupplyTypes &= ~ResupplyType.Repair;
-					notifyRepair.AfterRepair(host.Actor, self);
-
-				repairComplete = true;
 				return;
 			}
 
@@ -163,10 +148,6 @@ namespace OpenRA.Mods.Common.Activities
 				}
 
 				self.InflictDamage(host.Actor, new Damage(-hpToRepair));
-
-				foreach (var depot in host.Actor.TraitsImplementing<INotifyRepair>())
-					depot.RepairTick(host.Actor, self);
-
 				remainingTicks = repairsUnits.Info.Interval;
 			}
 			else
@@ -175,34 +156,26 @@ namespace OpenRA.Mods.Common.Activities
 
 		void RearmTick(Actor self)
 		{
-			rearmComplete = true;
+			var rearmComplete = true;
-			foreach (var pool in rearmable.RearmableAmmoPools)
+			foreach (var ammoPool in rearmable.RearmableAmmoPools)
 			{
-				if (!pool.FullAmmo())
+				if (!ammoPool.FullAmmo())
-				{
-					Reload(self, host.Actor, pool);
-					rearmComplete = false;
-				}
-			}
-
-			if (rearmComplete)
-				foreach (var notifyRearm in host.Actor.TraitsImplementing<INotifyRearm>())
-					notifyRearm.RearmingFinished(host.Actor, self);
-		}
-
-		void Reload(Actor self, Actor host, AmmoPool ammoPool)
 				{
 					if (--ammoPool.RemainingTicks <= 0)
 					{
-				foreach (var notify in host.TraitsImplementing<INotifyRearm>())
-					notify.Rearming(host, self);
-
 						ammoPool.RemainingTicks = ammoPool.Info.ReloadDelay;
 						if (!string.IsNullOrEmpty(ammoPool.Info.RearmSound))
 							Game.Sound.PlayToPlayer(SoundType.World, self.Owner, ammoPool.Info.RearmSound, self.CenterPosition);
 
 						ammoPool.GiveAmmo(self, ammoPool.Info.ReloadCount);
 					}
+
+					rearmComplete = false;
+				}
+			}
+
+			if (rearmComplete)
+				activeResupplyTypes &= ~ResupplyType.Rearm;
 		}
 	}
 }

OpenRA.Mods.Common/Traits/Render/WithRepairOverlay.cs

@@ -14,6 +14,7 @@ using OpenRA.Traits;
 
 namespace OpenRA.Mods.Common.Traits.Render
 {
+	// TODO: Refactor this trait into WithResupplyOverlay
 	[Desc("Displays an overlay when the building is being repaired by the player.")]
 	public class WithRepairOverlayInfo : PausableConditionalTraitInfo, Requires<RenderSpritesInfo>, Requires<BodyOrientationInfo>
 	{
@@ -38,10 +39,11 @@ namespace OpenRA.Mods.Common.Traits.Render
 		public override object Create(ActorInitializer init) { return new WithRepairOverlay(init.Self, this); }
 	}
 
-	public class WithRepairOverlay : PausableConditionalTrait<WithRepairOverlayInfo>, INotifyDamageStateChanged, INotifyRepair
+	public class WithRepairOverlay : PausableConditionalTrait<WithRepairOverlayInfo>, INotifyDamageStateChanged, INotifyResupply
 	{
 		readonly Animation overlay;
 		bool visible;
+		bool repairing;
 
 		public WithRepairOverlay(Actor self, WithRepairOverlayInfo info)
 			: base(info)
@@ -65,8 +67,12 @@ namespace OpenRA.Mods.Common.Traits.Render
 			overlay.ReplaceAnim(RenderSprites.NormalizeSequence(overlay, e.DamageState, overlay.CurrentSequence.Name));
 		}
 
-		void INotifyRepair.BeforeRepair(Actor self, Actor host)
+		void INotifyResupply.BeforeResupply(Actor self, Actor target, ResupplyType types)
 		{
+			repairing = types.HasFlag(ResupplyType.Repair);
+			if (!repairing)
+				return;
+
 			if (Info.StartSequence != null)
 			{
 				visible = true;
@@ -75,18 +81,18 @@ namespace OpenRA.Mods.Common.Traits.Render
 			}
 		}
 
-		void INotifyRepair.RepairTick(Actor self, Actor host)
+		void INotifyResupply.ResupplyTick(Actor self, Actor target, ResupplyType types)
 		{
-			if (Info.StartSequence == null)
+			var wasRepairing = repairing;
+			repairing = types.HasFlag(ResupplyType.Repair);
+
+			if (repairing && Info.StartSequence == null && !visible)
 			{
 				visible = true;
 				overlay.PlayThen(overlay.CurrentSequence.Name, () => visible = false);
 			}
-		}
 
-		void INotifyRepair.AfterRepair(Actor self, Actor target)
+			if (!repairing && wasRepairing && Info.EndSequence != null)
-		{
-			if (Info.EndSequence != null)
 			{
 				visible = true;
 				overlay.PlayThen(Info.EndSequence, () => visible = false);

OpenRA.Mods.Common/Traits/Render/WithResupplyAnimation.cs

@@ -15,13 +15,6 @@ using OpenRA.Traits;
 
 namespace OpenRA.Mods.Common.Traits.Render
 {
-	[Flags]
-	public enum ResupplyType
-	{
-		Rearm = 1,
-		Repair = 2
-	}
-
 	[Desc("Replaces the default animation when actor resupplies a unit.")]
 	public class WithResupplyAnimationInfo : ConditionalTraitInfo, Requires<WithSpriteBodyInfo>
 	{
@@ -37,7 +30,7 @@ namespace OpenRA.Mods.Common.Traits.Render
 		public override object Create(ActorInitializer init) { return new WithResupplyAnimation(init.Self, this); }
 	}
 
-	public class WithResupplyAnimation : ConditionalTrait<WithResupplyAnimationInfo>, INotifyRepair, INotifyRearm, ITick
+	public class WithResupplyAnimation : ConditionalTrait<WithResupplyAnimationInfo>, INotifyResupply, ITick
 	{
 		readonly WithSpriteBody wsb;
 		bool animPlaying;
@@ -71,28 +64,16 @@ namespace OpenRA.Mods.Common.Traits.Render
 			}
 		}
 
-		void INotifyRepair.BeforeRepair(Actor self, Actor target)
+		void INotifyResupply.BeforeResupply(Actor self, Actor target, ResupplyType types)
 		{
-			repairing = true;
+			repairing = types.HasFlag(ResupplyType.Repair);
+			rearming = types.HasFlag(ResupplyType.Rearm);
 		}
 
-		void INotifyRepair.RepairTick(Actor self, Actor target) { }
+		void INotifyResupply.ResupplyTick(Actor self, Actor target, ResupplyType types)
-
-		void INotifyRepair.AfterRepair(Actor self, Actor target)
 		{
-			repairing = false;
+			repairing = types.HasFlag(ResupplyType.Repair);
-		}
+			rearming = types.HasFlag(ResupplyType.Rearm);
-
-		void INotifyRearm.RearmingStarted(Actor self, Actor target)
-		{
-			rearming = true;
-		}
-
-		void INotifyRearm.Rearming(Actor self, Actor target) { }
-
-		void INotifyRearm.RearmingFinished(Actor self, Actor target)
-		{
-			rearming = false;
 		}
 
 		protected override void TraitDisabled(Actor self)

OpenRA.Mods.Common/TraitsInterfaces.cs

@@ -24,6 +24,14 @@ namespace OpenRA.Mods.Common.Traits
 
 	public enum AttackDelayType { Preparation, Attack }
 
+	[Flags]
+	public enum ResupplyType
+	{
+		None = 0,
+		Rearm = 1,
+		Repair = 2
+	}
+
 	public interface IQuantizeBodyOrientationInfo : ITraitInfo
 	{
 		int QuantizedBodyFacings(ActorInfo ai, SequenceProvider sequenceProvider, string race);
@@ -116,11 +124,10 @@ namespace OpenRA.Mods.Common.Traits
 	public interface INotifyAppliedDamage { void AppliedDamage(Actor self, Actor damaged, AttackInfo e); }
 
 	[RequireExplicitImplementation]
-	public interface INotifyRepair
+	public interface INotifyResupply
 	{
-		void BeforeRepair(Actor self, Actor target);
+		void BeforeResupply(Actor host, Actor target, ResupplyType types);
-		void RepairTick(Actor self, Actor target);
+		void ResupplyTick(Actor host, Actor target, ResupplyType types);
-		void AfterRepair(Actor self, Actor target);
 	}
 
 	[RequireExplicitImplementation]