From c89d90e4b0ce84603b44f6f47aa96bf9aa7e1484 Mon Sep 17 00:00:00 2001
From: Mustafa Alperen Seki <mustafaoyunda37@gmail.com>
Date: Sat, 13 Apr 2019 20:30:59 +0300
Subject: [PATCH] Add Types to EngineerRepair(able).

---
 OpenRA.Mods.Common/Traits/EngineerRepair.cs     | 16 ++++++++++++++--
 OpenRA.Mods.Common/Traits/EngineerRepairable.cs |  9 ++++++++-
 2 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/OpenRA.Mods.Common/Traits/EngineerRepair.cs b/OpenRA.Mods.Common/Traits/EngineerRepair.cs
index 1499758a93..44637f7d5e 100644
--- a/OpenRA.Mods.Common/Traits/EngineerRepair.cs
+++ b/OpenRA.Mods.Common/Traits/EngineerRepair.cs
@@ -10,6 +10,7 @@
 #endregion
 
 using System.Collections.Generic;
+using System.Linq;
 using OpenRA.Mods.Common.Activities;
 using OpenRA.Mods.Common.Orders;
 using OpenRA.Primitives;
@@ -20,6 +21,9 @@ namespace OpenRA.Mods.Common.Traits
 	[Desc("Can instantly repair other actors, but gets consumed afterwards.")]
 	class EngineerRepairInfo : ITraitInfo
 	{
+		[Desc("Uses the \"EngineerRepairable\" trait to determine repairability.")]
+		public readonly BitSet<EngineerRepairType> Types = default(BitSet<EngineerRepairType>);
+
 		[VoiceReference] public readonly string Voice = "Action";
 
 		[Desc("Behaviour when entering the structure.",
@@ -104,7 +108,11 @@ namespace OpenRA.Mods.Common.Traits
 
 			public override bool CanTargetActor(Actor self, Actor target, TargetModifiers modifiers, ref string cursor)
 			{
-				if (!target.Info.HasTraitInfo<EngineerRepairableInfo>())
+				var engineerRepairable = target.Info.TraitInfoOrDefault<EngineerRepairableInfo>();
+				if (engineerRepairable == null)
+					return false;
+
+				if (!engineerRepairable.Types.IsEmpty && !engineerRepairable.Types.Overlaps(info.Types))
 					return false;
 
 				if (!info.ValidStances.HasStance(self.Owner.Stances[target.Owner]))
@@ -118,7 +126,11 @@ namespace OpenRA.Mods.Common.Traits
 
 			public override bool CanTargetFrozenActor(Actor self, FrozenActor target, TargetModifiers modifiers, ref string cursor)
 			{
-				if (!target.Info.HasTraitInfo<EngineerRepairableInfo>())
+				var engineerRepairable = target.Info.TraitInfoOrDefault<EngineerRepairableInfo>();
+				if (engineerRepairable == null)
+					return false;
+
+				if (!engineerRepairable.Types.IsEmpty && !engineerRepairable.Types.Overlaps(info.Types))
 					return false;
 
 				if (!info.ValidStances.HasStance(self.Owner.Stances[target.Owner]))
diff --git a/OpenRA.Mods.Common/Traits/EngineerRepairable.cs b/OpenRA.Mods.Common/Traits/EngineerRepairable.cs
index 31c49a2a0a..116bc4adb5 100644
--- a/OpenRA.Mods.Common/Traits/EngineerRepairable.cs
+++ b/OpenRA.Mods.Common/Traits/EngineerRepairable.cs
@@ -9,12 +9,19 @@
  */
 #endregion
 
+using OpenRA.Primitives;
 using OpenRA.Traits;
 
 namespace OpenRA.Mods.Common.Traits
 {
+	public class EngineerRepairType { }
+
 	[Desc("Eligible for instant repair.")]
-	class EngineerRepairableInfo : TraitInfo<EngineerRepairable> { }
+	class EngineerRepairableInfo : TraitInfo<EngineerRepairable>
+	{
+		[Desc("Actors with these Types under EngineerRepair trait can repair me.")]
+		public readonly BitSet<EngineerRepairType> Types = default(BitSet<EngineerRepairType>);
+	}
 
 	class EngineerRepairable { }
 }