From 160ade10606f5b210a0edcff48cb7e0f2929a8cc Mon Sep 17 00:00:00 2001
From: Paul Chote <pchote@users.noreply.github.com>
Date: Mon, 23 Apr 2018 03:22:31 +0000
Subject: [PATCH] Add CaptureTypes to GiveCashOnCapture.

---
 .../Traits/GivesCashOnCapture.cs              | 23 ++++++++++++++++++-
 1 file changed, 22 insertions(+), 1 deletion(-)

diff --git a/OpenRA.Mods.Common/Traits/GivesCashOnCapture.cs b/OpenRA.Mods.Common/Traits/GivesCashOnCapture.cs
index 2ca0929314..0dc5c04cd3 100644
--- a/OpenRA.Mods.Common/Traits/GivesCashOnCapture.cs
+++ b/OpenRA.Mods.Common/Traits/GivesCashOnCapture.cs
@@ -10,6 +10,8 @@
 #endregion
 
 using System;
+using System.Collections.Generic;
+using System.Linq;
 using OpenRA.Mods.Common.Effects;
 using OpenRA.Traits;
 
@@ -27,6 +29,9 @@ namespace OpenRA.Mods.Common.Traits
 		[Desc("Amount of money awarded for capturing the actor.")]
 		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 override object Create(ActorInitializer init) { return new GivesCashOnCapture(this); }
 	}
 
@@ -42,7 +47,7 @@ namespace OpenRA.Mods.Common.Traits
 
 		void INotifyCapture.OnCapture(Actor self, Actor captor, Player oldOwner, Player newOwner)
 		{
-			if (IsTraitDisabled)
+			if (IsTraitDisabled || !IsValidCaptor(captor))
 				return;
 
 			var resources = newOwner.PlayerActor.Trait<PlayerResources>();
@@ -66,5 +71,21 @@ namespace OpenRA.Mods.Common.Traits
 			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>();
+			if (capturesInfo != null && info.CaptureTypes.Overlaps(capturesInfo.CaptureTypes))
+				return true;
+
+			var externalCapturesInfo = captor.Info.TraitInfoOrDefault<ExternalCapturesInfo>();
+			if (externalCapturesInfo != null && info.CaptureTypes.Overlaps(externalCapturesInfo.CaptureTypes))
+				return true;
+
+			return false;
+		}
 	}
 }