From 9905f217ef7abd678fe27c33c4cc78e40e465672 Mon Sep 17 00:00:00 2001
From: atlimit8 <atlimit8-vcs@gmx.com>
Date: Mon, 6 Jul 2015 00:17:58 -0500
Subject: [PATCH] Make TargetableUnit upgradable

---
 .../Traits/Air/TargetableAircraft.cs             |  4 ++--
 OpenRA.Mods.Common/Traits/TargetableUnit.cs      | 16 +++++++++-------
 OpenRA.Mods.RA/Traits/TargetableSubmarine.cs     |  3 ++-
 3 files changed, 13 insertions(+), 10 deletions(-)

diff --git a/OpenRA.Mods.Common/Traits/Air/TargetableAircraft.cs b/OpenRA.Mods.Common/Traits/Air/TargetableAircraft.cs
index 9067f61630..e8d9924197 100644
--- a/OpenRA.Mods.Common/Traits/Air/TargetableAircraft.cs
+++ b/OpenRA.Mods.Common/Traits/Air/TargetableAircraft.cs
@@ -35,8 +35,8 @@ namespace OpenRA.Mods.Common.Traits
 		{
 			get
 			{
-				return (self.CenterPosition.Z > 0) ? info.TargetTypes
-					: info.GroundedTargetTypes;
+				return IsTraitDisabled ? None
+					: (self.CenterPosition.Z > 0 ? info.TargetTypes : info.GroundedTargetTypes);
 			}
 		}
 	}
diff --git a/OpenRA.Mods.Common/Traits/TargetableUnit.cs b/OpenRA.Mods.Common/Traits/TargetableUnit.cs
index a38111173a..733c792b10 100644
--- a/OpenRA.Mods.Common/Traits/TargetableUnit.cs
+++ b/OpenRA.Mods.Common/Traits/TargetableUnit.cs
@@ -14,7 +14,7 @@ using OpenRA.Traits;
 namespace OpenRA.Mods.Common.Traits
 {
 	[Desc("Actor can be targeted.")]
-	public class TargetableUnitInfo : ITraitInfo, ITargetableInfo
+	public class TargetableUnitInfo : UpgradableTraitInfo, ITargetableInfo
 	{
 		[Desc("Target type. Used for filtering (in)valid targets.")]
 		public readonly string[] TargetTypes = { };
@@ -22,35 +22,37 @@ namespace OpenRA.Mods.Common.Traits
 
 		public bool RequiresForceFire = false;
 
-		public virtual object Create(ActorInitializer init) { return new TargetableUnit(init.Self, this); }
+		public override object Create(ActorInitializer init) { return new TargetableUnit(init.Self, this); }
 	}
 
-	public class TargetableUnit : ITargetable
+	public class TargetableUnit : UpgradableTrait<TargetableUnitInfo>, ITargetable
 	{
-		readonly TargetableUnitInfo info;
+		protected static readonly string[] None = new string[] { };
 		protected Cloak cloak;
 
 		public TargetableUnit(Actor self, TargetableUnitInfo info)
+			: base(info)
 		{
-			this.info = info;
 			cloak = self.TraitOrDefault<Cloak>();
 		}
 
 		public virtual bool TargetableBy(Actor self, Actor viewer)
 		{
+			if (IsTraitDisabled)
+				return false;
 			if (cloak == null || (!viewer.IsDead && viewer.HasTrait<IgnoresCloak>()))
 				return true;
 
 			return cloak.IsVisible(self, viewer.Owner);
 		}
 
-		public virtual string[] TargetTypes { get { return info.TargetTypes; } }
+		public virtual string[] TargetTypes { get { return IsTraitDisabled ? None : Info.TargetTypes; } }
 
 		public virtual IEnumerable<WPos> TargetablePositions(Actor self)
 		{
 			yield return self.CenterPosition;
 		}
 
-		public bool RequiresForceFire { get { return info.RequiresForceFire; } }
+		public bool RequiresForceFire { get { return Info.RequiresForceFire; } }
 	}
 }
diff --git a/OpenRA.Mods.RA/Traits/TargetableSubmarine.cs b/OpenRA.Mods.RA/Traits/TargetableSubmarine.cs
index c5aad54b27..f90d1d79aa 100644
--- a/OpenRA.Mods.RA/Traits/TargetableSubmarine.cs
+++ b/OpenRA.Mods.RA/Traits/TargetableSubmarine.cs
@@ -34,7 +34,8 @@ namespace OpenRA.Mods.RA.Traits
 		{
 			get
 			{
-				return cloak.Cloaked ? info.CloakedTargetTypes : info.TargetTypes;
+				return IsTraitDisabled ? None
+					: (cloak.Cloaked ? info.CloakedTargetTypes : info.TargetTypes);
 			}
 		}
 	}