From 508c99655c0c583d41104893c8e86d9723463fb6 Mon Sep 17 00:00:00 2001
From: reaperrr <thereaperrr@aol.com>
Date: Sun, 9 Aug 2015 02:13:14 +0200
Subject: [PATCH] Made Armor upgradable

---
 OpenRA.Mods.Common/Traits/Armor.cs                    | 10 ++++++++--
 OpenRA.Mods.Common/Warheads/DamageWarhead.cs          | 11 ++++++-----
 .../Warheads/HealthPercentageDamageWarhead.cs         |  2 +-
 3 files changed, 15 insertions(+), 8 deletions(-)

diff --git a/OpenRA.Mods.Common/Traits/Armor.cs b/OpenRA.Mods.Common/Traits/Armor.cs
index 04064cd913..e637ff52a1 100644
--- a/OpenRA.Mods.Common/Traits/Armor.cs
+++ b/OpenRA.Mods.Common/Traits/Armor.cs
@@ -13,10 +13,16 @@ using OpenRA.Traits;
 namespace OpenRA.Mods.Common.Traits
 {
 	[Desc("Used to define weapon efficiency modifiers with different percentages per Type.")]
-	public class ArmorInfo : TraitInfo<Armor>
+	public class ArmorInfo : UpgradableTraitInfo
 	{
 		public readonly string Type = null;
+
+		public override object Create(ActorInitializer init) { return new Armor(init.Self, this); }
 	}
 
-	public class Armor { }
+	public class Armor : UpgradableTrait<ArmorInfo>
+	{
+		public Armor(Actor self, ArmorInfo info)
+			: base(info) { }
+	}
 }
\ No newline at end of file
diff --git a/OpenRA.Mods.Common/Warheads/DamageWarhead.cs b/OpenRA.Mods.Common/Warheads/DamageWarhead.cs
index 460d019f52..5cbac65fe5 100644
--- a/OpenRA.Mods.Common/Warheads/DamageWarhead.cs
+++ b/OpenRA.Mods.Common/Warheads/DamageWarhead.cs
@@ -9,6 +9,7 @@
 #endregion
 
 using System.Collections.Generic;
+using System.Linq;
 using OpenRA.Mods.Common.Traits;
 using OpenRA.Traits;
 
@@ -34,13 +35,13 @@ namespace OpenRA.Mods.Common.Warheads
 				: new Dictionary<string, int>();
 		}
 
-		public int DamageVersus(ActorInfo victim)
+		public int DamageVersus(Actor victim)
 		{
-			var armor = victim.Traits.GetOrDefault<ArmorInfo>();
-			if (armor != null && armor.Type != null)
+			var armor = victim.TraitsImplementing<Armor>().Where(a => !a.IsTraitDisabled && a.Info.Type != null);
+			foreach (var a in armor)
 			{
 				int versus;
-				if (Versus.TryGetValue(armor.Type, out versus))
+				if (Versus.TryGetValue(a.Info.Type, out versus))
 					return versus;
 			}
 
@@ -60,7 +61,7 @@ namespace OpenRA.Mods.Common.Warheads
 
 		public virtual void DoImpact(Actor victim, Actor firedBy, IEnumerable<int> damageModifiers)
 		{
-			var damage = Util.ApplyPercentageModifiers(Damage, damageModifiers.Append(DamageVersus(victim.Info)));
+			var damage = Util.ApplyPercentageModifiers(Damage, damageModifiers.Append(DamageVersus(victim)));
 			victim.InflictDamage(firedBy, damage, this);
 		}
 	}
diff --git a/OpenRA.Mods.Common/Warheads/HealthPercentageDamageWarhead.cs b/OpenRA.Mods.Common/Warheads/HealthPercentageDamageWarhead.cs
index c1d5c5ab1f..62d252d849 100644
--- a/OpenRA.Mods.Common/Warheads/HealthPercentageDamageWarhead.cs
+++ b/OpenRA.Mods.Common/Warheads/HealthPercentageDamageWarhead.cs
@@ -38,7 +38,7 @@ namespace OpenRA.Mods.Common.Warheads
 				return;
 
 			// Damage is measured as a percentage of the target health
-			var damage = Util.ApplyPercentageModifiers(healthInfo.HP, damageModifiers.Append(Damage, DamageVersus(victim.Info)));
+			var damage = Util.ApplyPercentageModifiers(healthInfo.HP, damageModifiers.Append(Damage, DamageVersus(victim)));
 			victim.InflictDamage(firedBy, damage, this);
 		}
 	}