Added IDisabledTrait & rewrote upgrade code using a level-based approach.

Upgradeable traits are notified whenever an upgrade of their declared types are granted or revoked.  The traits maintain their own internal level counter, which is then used to enable or disable the trait functionality.  A trait can register for multiple upgrade types which then all affect the internal level counter.

	IDisabledTrait for identifying (and filtering) disabled traits
	UpgradableTrait provides an abstract base for traits to support upgrade levels
	Added IDisabledTrait support to GlobalButtonOrderGenerator

	Includes rework by pchote with alterations.

OpenRA.Mods.Common/CommonTraitsInterfaces.cs

@@ -21,4 +21,11 @@ namespace OpenRA.Mods.Common
 
 	public interface INotifyChat { bool OnChat(string from, string message); }
 	public interface IRenderActorPreviewInfo { IEnumerable<IActorPreview> RenderPreview (ActorPreviewInitializer init); }
+
+	public interface IUpgradable
+	{
+		IEnumerable<string> UpgradeTypes { get; }
+		bool AcceptsUpgradeLevel(Actor self, string type, int level);
+		void UpgradeLevelChanged(Actor self, string type, int oldLevel, int newLevel);
+	}
 }

OpenRA.Mods.Common/Modifiers/TimedUpgradeBar.cs

@@ -0,0 +1,58 @@
+#region Copyright & License Information
+/*
+ * Copyright 2007-2014 The OpenRA Developers (see AUTHORS)
+ * This file is part of OpenRA, which is free software. It is made
+ * available to you under the terms of the GNU General Public License
+ * as published by the Free Software Foundation. For more information,
+ * see COPYING.
+ */
+#endregion
+
+using System;
+using System.Drawing;
+using System.Linq;
+using OpenRA.Traits;
+
+namespace OpenRA.Mods.Common
+{
+	[Desc("Visualizes the remaining time for an upgrade.")]
+	class TimedUpgradeBarInfo : ITraitInfo, Requires<UpgradeManagerInfo>
+	{
+		[Desc("Upgrade that this bar corresponds to")]
+		public readonly string Upgrade = null;
+
+		public readonly Color Color = Color.Red;
+
+		public object Create(ActorInitializer init) { return new TimedUpgradeBar(init.self, this); }
+	}
+
+	class TimedUpgradeBar : ISelectionBar
+	{
+		readonly TimedUpgradeBarInfo info;
+		readonly Actor self;
+		float value;
+
+		public TimedUpgradeBar(Actor self, TimedUpgradeBarInfo info)
+		{
+			this.self = self;
+			this.info = info;
+
+			self.Trait<UpgradeManager>().RegisterWatcher(info.Upgrade, Update);
+		}
+
+		public void Update(int duration, int remaining)
+		{
+			value = remaining * 1f / duration;
+		}
+
+		public float GetValue()
+		{
+			if (!self.Owner.IsAlliedWith(self.World.RenderPlayer))
+				return 0;
+
+			return value;
+		}
+
+		public Color GetColor() { return info.Color; }
+	}
+}

OpenRA.Mods.Common/OpenRA.Mods.Common.csproj

@@ -80,6 +80,7 @@
     <Compile Include="LoadWidgetAtGameStart.cs" />
     <Compile Include="ModChooserLoadScreen.cs" />
     <Compile Include="NullLoadScreen.cs" />
+    <Compile Include="Modifiers\TimedUpgradeBar.cs" />
     <Compile Include="Orders\DeployOrderTargeter.cs" />
     <Compile Include="Orders\EnterAlliedActorTargeter.cs" />
     <Compile Include="Orders\UnitOrderTargeter.cs" />
@@ -136,6 +137,8 @@
     <Compile Include="Traits\World\SpawnMapActors.cs" />
     <Compile Include="Traits\World\StartGameNotification.cs" />
     <Compile Include="Traits\World\TerrainGeometryOverlay.cs" />
+    <Compile Include="Upgrades\UpgradableTrait.cs" />
+    <Compile Include="Upgrades\UpgradeManager.cs" />
     <Compile Include="UtilityCommands\ConvertPngToShpCommand.cs" />
     <Compile Include="UtilityCommands\ConvertSpriteToPngCommand.cs" />
     <Compile Include="UtilityCommands\ExtractFilesCommand.cs" />

OpenRA.Mods.Common/Traits/Modifiers/UpgradeOverlay.cs

@@ -15,26 +15,18 @@ using OpenRA.Traits;
 namespace OpenRA.Mods.Common.Traits
 {
 	[Desc("Display a colored overlay when a timed upgrade is active.")]
-	public class UpgradeOverlayInfo : ITraitInfo
+	public class UpgradeOverlayInfo : UpgradableTraitInfo, ITraitInfo
 	{
-		[Desc("Upgrade that is required before this overlay is rendered")]
-		public readonly string RequiresUpgrade = null;
-
 		[Desc("Palette to use when rendering the overlay")]
 		public readonly string Palette = "invuln";
 
 		public object Create(ActorInitializer init) { return new UpgradeOverlay(this); }
 	}
 
-	public class UpgradeOverlay : IRenderModifier, IUpgradable
+	public class UpgradeOverlay : UpgradableTrait<UpgradeOverlayInfo>, IRenderModifier
 	{
-		readonly UpgradeOverlayInfo info;
-		bool enabled;
-
 		public UpgradeOverlay(UpgradeOverlayInfo info)
-		{
-			this.info = info;
-		}
+			: base (info) { }
 
 		public IEnumerable<IRenderable> ModifyRender(Actor self, WorldRenderer wr, IEnumerable<IRenderable> r)
 		{
@@ -42,22 +34,11 @@ namespace OpenRA.Mods.Common.Traits
 			{
 				yield return a;
 
-				if (enabled && !a.IsDecoration)
-					yield return a.WithPalette(wr.Palette(info.Palette))
+				if (!IsTraitDisabled && !a.IsDecoration)
+					yield return a.WithPalette(wr.Palette(Info.Palette))
 						.WithZOffset(a.ZOffset + 1)
 						.AsDecoration();
 			}
 		}
-
-		public bool AcceptsUpgrade(string type)
-		{
-			return type == info.RequiresUpgrade;
-		}
-
-		public void UpgradeAvailable(Actor self, string type, bool available)
-		{
-			if (type == info.RequiresUpgrade)
-				enabled = available;
-		}
 	}
 }

OpenRA.Mods.Common/Upgrades/UpgradableTrait.cs

@@ -0,0 +1,86 @@
+#region Copyright & License Information
+/*
+ * Copyright 2007-2014 The OpenRA Developers (see AUTHORS)
+ * This file is part of OpenRA, which is free software. It is made
+ * available to you under the terms of the GNU General Public License
+ * as published by the Free Software Foundation. For more information,
+ * see COPYING.
+ */
+#endregion
+using System.Collections.Generic;
+using System.Linq;
+using OpenRA.Traits;
+
+namespace OpenRA.Mods.Common
+{
+	/// <summary>Use as base class for *Info to subclass of UpgradableTrait. (See UpgradableTrait.)</summary>
+	public abstract class UpgradableTraitInfo
+	{
+		[Desc("The upgrade types which can enable or disable this trait.")]
+		public readonly string[] UpgradeTypes = { };
+
+		[Desc("The minimum upgrade level at which this trait is enabled.", "Defaults to 0 (enabled by default).")]
+		public readonly int UpgradeMinEnabledLevel = 0;
+
+		[Desc("The maximum upgrade level at which the trait is enabled.",
+			"Defaults to UpgradeMaxAcceptedLevel (enabled for all levels greater than UpgradeMinEnabledLevel).",
+			"Set this to a value smaller than UpgradeMaxAcceptedLevel to disable the trait at higher levels.",
+			"Use UpgradeMaxAcceptedLevel: 2 (1 more) to be able to extend upgrade time.")]
+		public readonly int UpgradeMaxEnabledLevel = int.MaxValue;
+
+		[Desc("The maximum upgrade level that this trait will accept.")]
+		public readonly int UpgradeMaxAcceptedLevel = 1;
+	}
+
+	/// <summary>
+	/// Abstract base for enabling and disabling trait using upgrades.
+	/// Requires basing *Info on UpgradableTraitInfo and using base(info) constructor.
+	/// Note that EnabledByUpgrade is not called at creation even if this starts as enabled.
+	/// </summary>,
+	public abstract class UpgradableTrait<InfoType> : IUpgradable, IDisabledTrait, ISync where InfoType : UpgradableTraitInfo
+	{
+		public readonly InfoType Info;
+		public IEnumerable<string> UpgradeTypes { get { return Info.UpgradeTypes; } }
+		[Sync] public bool IsTraitDisabled { get; private set; }
+
+		public UpgradableTrait(InfoType info)
+		{
+			Info = info;
+			IsTraitDisabled = info.UpgradeTypes != null && info.UpgradeTypes.Length > 0 && info.UpgradeMinEnabledLevel > 0;
+		}
+
+		public bool AcceptsUpgradeLevel(Actor self, string type, int level)
+		{
+			return level > 0 && level <= Info.UpgradeMaxAcceptedLevel;
+		}
+
+		public void UpgradeLevelChanged(Actor self, string type, int oldLevel, int newLevel)
+		{
+			if (!Info.UpgradeTypes.Contains(type))
+				return;
+
+			// Restrict the levels to the allowed range
+			oldLevel = oldLevel.Clamp(0, Info.UpgradeMaxAcceptedLevel);
+			newLevel = newLevel.Clamp(0, Info.UpgradeMaxAcceptedLevel);
+			if (oldLevel == newLevel)
+				return;
+
+			var wasDisabled = IsTraitDisabled;
+			IsTraitDisabled = newLevel < Info.UpgradeMinEnabledLevel || newLevel > Info.UpgradeMaxEnabledLevel;
+			UpgradeLevelChanged(self, oldLevel, newLevel);
+
+			if (IsTraitDisabled != wasDisabled)
+			{
+				if (wasDisabled)
+					UpgradeEnabled(self);
+				else
+					UpgradeDisabled(self);
+			}
+		}
+
+		// Subclasses can add upgrade support by querying IsTraitDisabled and/or overriding these methods.
+		protected virtual void UpgradeLevelChanged(Actor self, int oldLevel, int newLevel) { }
+		protected virtual void UpgradeEnabled(Actor self) { }
+		protected virtual void UpgradeDisabled(Actor self) { }
+	}
+}

OpenRA.Mods.Common/Upgrades/UpgradeManager.cs

@@ -0,0 +1,160 @@
+#region Copyright & License Information
+/*
+ * Copyright 2007-2014 The OpenRA Developers (see AUTHORS)
+ * This file is part of OpenRA, which is free software. It is made
+ * available to you under the terms of the GNU General Public License
+ * as published by the Free Software Foundation. For more information,
+ * see COPYING.
+ */
+#endregion
+
+using System;
+using System.Collections.Generic;
+using System.Drawing;
+using System.Linq;
+using OpenRA.Graphics;
+using OpenRA.Primitives;
+using OpenRA.Traits;
+
+namespace OpenRA.Mods.Common
+{
+	public class UpgradeManagerInfo : ITraitInfo
+	{
+		public object Create(ActorInitializer init) { return new UpgradeManager(init); }
+	}
+
+	public class UpgradeManager : ITick
+	{
+		class TimedUpgrade
+		{
+			public readonly string Upgrade;
+			public readonly int Duration;
+			public int Remaining;
+
+			public TimedUpgrade(string upgrade, int duration)
+			{
+				Upgrade = upgrade;
+				Duration = duration;
+				Remaining = duration;
+			}
+
+			public void Tick() { Remaining--; }
+		}
+
+		class UpgradeState
+		{
+			public readonly List<IUpgradable> Traits = new List<IUpgradable>();
+			public readonly List<object> Sources = new List<object>();
+			public readonly List<Action<int, int>> Watchers = new List<Action<int, int>>();
+		}
+
+		readonly List<TimedUpgrade> timedUpgrades = new List<TimedUpgrade>();
+		readonly Lazy<Dictionary<string, UpgradeState>> upgrades;
+		readonly Dictionary<IUpgradable, int> levels = new Dictionary<IUpgradable, int>();
+
+		public UpgradeManager(ActorInitializer init)
+		{
+			upgrades = Exts.Lazy(() =>
+			{
+				var ret = new Dictionary<string, UpgradeState>();
+				foreach (var up in init.self.TraitsImplementing<IUpgradable>())
+					foreach (var t in up.UpgradeTypes)
+						ret.GetOrAdd(t).Traits.Add(up);
+
+				return ret;
+			});
+		}
+
+		public void GrantTimedUpgrade(Actor self, string upgrade, int duration)
+		{
+			var timed = timedUpgrades.FirstOrDefault(u => u.Upgrade == upgrade);
+			if (timed == null)
+			{
+				timed = new TimedUpgrade(upgrade, duration);
+				timedUpgrades.Add(timed);
+				GrantUpgrade(self, upgrade, timed);
+			}
+			else
+				timed.Remaining = Math.Max(duration, timed.Remaining);
+		}
+
+		// Different upgradeable traits may define (a) different level ranges for the same upgrade type,
+		// and (b) multiple upgrade types for the same trait. The unrestricted level for each trait is
+		// tracked independently so that we can can correctly revoke levels without adding the burden of
+		// tracking both the overall (unclamped) and effective (clamped) levels on each individual trait.
+		void NotifyUpgradeLevelChanged(IEnumerable<IUpgradable> traits, Actor self, string upgrade, int levelAdjust)
+		{
+			foreach (var up in traits)
+			{
+				var oldLevel = levels.GetOrAdd(up);
+				var newLevel = levels[up] = oldLevel + levelAdjust;
+
+				// This will internally clamp the levels to its own restricted range
+				up.UpgradeLevelChanged(self, upgrade, oldLevel, newLevel);
+			}
+		}
+
+		int GetOverallLevel(IUpgradable upgradable)
+		{
+			int level;
+			return levels.TryGetValue(upgradable, out level) ? level : 0;
+		}
+
+		public void GrantUpgrade(Actor self, string upgrade, object source)
+		{
+			UpgradeState s;
+			if (!upgrades.Value.TryGetValue(upgrade, out s))
+				return;
+
+			// Track the upgrade source so that the upgrade can be removed without conflicts
+			s.Sources.Add(source);
+
+			NotifyUpgradeLevelChanged(s.Traits, self, upgrade, 1);
+		}
+
+		public void RevokeUpgrade(Actor self, string upgrade, object source)
+		{
+			UpgradeState s;
+			if (!upgrades.Value.TryGetValue(upgrade, out s))
+				return;
+
+			if (!s.Sources.Remove(source))
+				throw new InvalidOperationException("Object <{0}> revoked more levels of upgrade {1} than it granted for {2}.".F(source, upgrade, self));
+
+			NotifyUpgradeLevelChanged(s.Traits, self, upgrade, -1);
+		}
+
+		public bool AcceptsUpgrade(Actor self, string upgrade)
+		{
+			UpgradeState s;
+			if (!upgrades.Value.TryGetValue(upgrade, out s))
+				return false;
+
+			return s.Traits.Any(up => up.AcceptsUpgradeLevel(self, upgrade, GetOverallLevel(up) + 1));
+		}
+
+		public void RegisterWatcher(string upgrade, Action<int, int> action)
+		{
+			UpgradeState s;
+			if (!upgrades.Value.TryGetValue(upgrade, out s))
+				return;
+
+			s.Watchers.Add(action);
+		}
+
+		public void Tick(Actor self)
+		{
+			foreach (var u in timedUpgrades)
+			{
+				u.Tick();
+				if (u.Remaining <= 0)
+					RevokeUpgrade(self, u.Upgrade, u);
+
+				foreach (var a in upgrades.Value[u.Upgrade].Watchers)
+					a(u.Duration, u.Remaining);
+			}
+
+			timedUpgrades.RemoveAll(u => u.Remaining <= 0);
+		}
+	}
+}

OpenRA.Mods.Common/UtilityCommands/UpgradeRules.cs

@@ -664,6 +664,33 @@ namespace OpenRA.Mods.Common.UtilityCommands
 					}
 				}
 
+				if (engineVersion < 20141121)
+				{
+					if (depth == 1)
+					{
+						if (node.Value.Nodes.Exists(n => n.Key == "RestrictedByUpgrade"))
+						{
+							node.Value.Nodes.Add(new MiniYamlNode("UpgradeMaxEnabledLevel", "0"));
+							node.Value.Nodes.Add(new MiniYamlNode("UpgradeMaxAcceptedLevel", "1"));
+						}
+						else if (node.Value.Nodes.Exists(n => n.Key == "RequiresUpgrade"))
+							node.Value.Nodes.Add(new MiniYamlNode("UpgradeMinEnabledLevel", "1"));
+
+						if (node.Key.StartsWith("DisableUpgrade") && !node.Value.Nodes.Any(n => n.Key == "RequiresUpgrade" || n.Key == "UpgradeTypes"))
+							node.Value.Nodes.Add(new MiniYamlNode("UpgradeTypes", "disable"));
+
+						if (node.Key.StartsWith("InvulnerabilityUpgrade") && !node.Value.Nodes.Any(n => n.Key == "RequiresUpgrade" || n.Key == "UpgradeTypes"))
+							node.Value.Nodes.Add(new MiniYamlNode("UpgradeTypes", "invulnerability"));
+					}
+					else if (depth == 2)
+					{
+						if (node.Key == "RequiresUpgrade" || node.Key == "RestrictedByUpgrade")
+							node.Key = "UpgradeTypes";
+						else if (node.Key == "-RequiresUpgrade" || node.Key == "-RestrictedByUpgrade")
+							node.Key = "-UpgradeTypes";
+					}
+				}
+
 				UpgradeActorRules(engineVersion, ref node.Value.Nodes, node, depth + 1);
 			}
 		}