From 346e670563aa85e6e08932dc3b6dc7d5caf10078 Mon Sep 17 00:00:00 2001
From: Paul Chote <pchote@users.noreply.github.com>
Date: Fri, 26 Oct 2018 20:35:06 +0100
Subject: [PATCH] Simplify type filtering in
 GivesCashOnCapture/TransformOnCapture.

---
 .../Traits/GivesCashOnCapture.cs               | 18 +++++-------------
 .../Traits/TransformOnCapture.cs               | 13 ++-----------
 2 files changed, 7 insertions(+), 24 deletions(-)

diff --git a/OpenRA.Mods.Common/Traits/GivesCashOnCapture.cs b/OpenRA.Mods.Common/Traits/GivesCashOnCapture.cs
index 7c84330bea..583a812a82 100644
--- a/OpenRA.Mods.Common/Traits/GivesCashOnCapture.cs
+++ b/OpenRA.Mods.Common/Traits/GivesCashOnCapture.cs
@@ -29,7 +29,7 @@ namespace OpenRA.Mods.Common.Traits
 		public readonly int Amount = 0;
 
 		[Desc("Award cash only if the capturer's CaptureTypes overlap with these types. Leave empty to allow all types.")]
-		public readonly HashSet<string> CaptureTypes = new HashSet<string>();
+		public readonly BitSet<CaptureType> CaptureTypes = default(BitSet<CaptureType>);
 
 		public override object Create(ActorInitializer init) { return new GivesCashOnCapture(this); }
 	}
@@ -46,27 +46,19 @@ namespace OpenRA.Mods.Common.Traits
 
 		void INotifyCapture.OnCapture(Actor self, Actor captor, Player oldOwner, Player newOwner, BitSet<CaptureType> captureTypes)
 		{
-			if (IsTraitDisabled || !IsValidCaptor(captor))
+			if (IsTraitDisabled)
+				return;
+
+			if (!info.CaptureTypes.IsEmpty && !info.CaptureTypes.Overlaps(captureTypes))
 				return;
 
 			var resources = newOwner.PlayerActor.Trait<PlayerResources>();
-
 			var amount = resources.ChangeCash(info.Amount);
-
 			if (!info.ShowTicks && amount != 0)
 				return;
 
 			self.World.AddFrameEndTask(w => w.Add(
 				new FloatingText(self.CenterPosition, self.Owner.Color.RGB, FloatingText.FormatCashTick(amount), info.DisplayDuration)));
 		}
-
-		bool IsValidCaptor(Actor captor)
-		{
-			if (!info.CaptureTypes.Any())
-				return true;
-
-			var capturesInfo = captor.Info.TraitInfoOrDefault<CapturesInfo>();
-			return capturesInfo != null && info.CaptureTypes.Overlaps(capturesInfo.CaptureTypes);
-		}
 	}
 }
diff --git a/OpenRA.Mods.Common/Traits/TransformOnCapture.cs b/OpenRA.Mods.Common/Traits/TransformOnCapture.cs
index 01fc4e37c5..7c45116f94 100644
--- a/OpenRA.Mods.Common/Traits/TransformOnCapture.cs
+++ b/OpenRA.Mods.Common/Traits/TransformOnCapture.cs
@@ -25,7 +25,7 @@ namespace OpenRA.Mods.Common.Traits
 		public readonly bool SkipMakeAnims = true;
 
 		[Desc("Transform only if the capturer's CaptureTypes overlap with these types. Leave empty to allow all types.")]
-		public readonly HashSet<string> CaptureTypes = new HashSet<string>();
+		public readonly BitSet<CaptureType> CaptureTypes = default(BitSet<CaptureType>);
 
 		public virtual object Create(ActorInitializer init) { return new TransformOnCapture(init, this); }
 	}
@@ -43,7 +43,7 @@ namespace OpenRA.Mods.Common.Traits
 
 		void INotifyCapture.OnCapture(Actor self, Actor captor, Player oldOwner, Player newOwner, BitSet<CaptureType> captureTypes)
 		{
-			if (!IsValidCaptor(captor))
+			if (!info.CaptureTypes.IsEmpty && !info.CaptureTypes.Overlaps(captureTypes))
 				return;
 
 			var facing = self.TraitOrDefault<IFacing>();
@@ -53,14 +53,5 @@ namespace OpenRA.Mods.Common.Traits
 			self.CancelActivity();
 			self.QueueActivity(transform);
 		}
-
-		bool IsValidCaptor(Actor captor)
-		{
-			if (!info.CaptureTypes.Any())
-				return true;
-
-			var capturesInfo = captor.Info.TraitInfoOrDefault<CapturesInfo>();
-			return capturesInfo != null && info.CaptureTypes.Overlaps(capturesInfo.CaptureTypes);
-		}
 	}
 }