Add a UnitUpgradeManager trait.

This introduces support for timed upgrades, starting with crate buffs.

OpenRA.Mods.Common/UtilityCommands/UpgradeRules.cs

@@ -575,6 +575,13 @@ namespace OpenRA.Mods.Common.UtilityCommands
 					}
 				}
 
+				// Routed unit upgrades via the UnitUpgradeManager trait
+				if (engineVersion < 20141001)
+				{
+					if (depth == 0 && node.Value.Nodes.Any(n => n.Key.StartsWith("GainsStatUpgrades")))
+						node.Value.Nodes.Add(new MiniYamlNode("UnitUpgradeManager", new MiniYaml("")));
+				}
+
 				UpgradeActorRules(engineVersion, ref node.Value.Nodes, node, depth + 1);
 			}
 		}

OpenRA.Mods.RA/Crates/UnitUpgradeCrateAction.cs

@@ -16,9 +16,12 @@ namespace OpenRA.Mods.RA.Crates
 	[Desc("Grants an upgrade to the collector.")]
 	public class UnitUpgradeCrateActionInfo : CrateActionInfo
 	{
-		[Desc("The upgrade to grant.")]
+		[Desc("The upgrades to apply.")]
 		public readonly string[] Upgrades = { };
 
+		[Desc("Duration of the upgrade (in ticks). Set to 0 for a permanent upgrade.")]
+		public readonly int Duration = 0;
+
 		[Desc("The range to search for extra collectors in.", "Extra collectors will also be granted the crate action.")]
 		public readonly WRange Range = new WRange(3);
 
@@ -42,16 +45,8 @@ namespace OpenRA.Mods.RA.Crates
 
 		bool AcceptsUpgrade(Actor a)
 		{
-			return a.TraitsImplementing<IUpgradable>()
+			var um = a.TraitOrDefault<UpgradeManager>();
-				.Any(up => info.Upgrades.Any(u => up.AcceptsUpgrade(u)));
+			return um != null && info.Upgrades.Any(u => um.AcceptsUpgrade(a, u));
-		}
-
-		void GrantActorUpgrades(Actor a)
-		{
-			foreach (var up in a.TraitsImplementing<IUpgradable>())
-				foreach (var u in info.Upgrades)
-					if (up.AcceptsUpgrade(u))
-						up.UpgradeAvailable(a, u, true);
 		}
 
 		public override int GetSelectionShares(Actor collector)
@@ -61,25 +56,32 @@ namespace OpenRA.Mods.RA.Crates
 
 		public override void Activate(Actor collector)
 		{
-			collector.World.AddFrameEndTask(w => GrantActorUpgrades(collector));
-
 			var actorsInRange = self.World.FindActorsInCircle(self.CenterPosition, info.Range)
-				.Where(a => a != self && a.Owner == collector.Owner && AcceptsUpgrade(a));
+				.Where(a => a != self && a != collector && a.Owner == collector.Owner && AcceptsUpgrade(a));
 
-			if (actorsInRange.Any())
+			if (info.MaxExtraCollectors > -1)
+				actorsInRange = actorsInRange.Take(info.MaxExtraCollectors);
+
+			collector.World.AddFrameEndTask(w =>
 			{
-				if (info.MaxExtraCollectors > -1)
+				foreach (var a in actorsInRange.Append(collector))
-					actorsInRange = actorsInRange.Take(info.MaxExtraCollectors);
-
-				collector.World.AddFrameEndTask(w =>
 				{
-					foreach (var a in actorsInRange)
+					if (!a.IsInWorld || a.IsDead())
+						continue;
+
+					var um = a.TraitOrDefault<UpgradeManager>();
+					foreach (var u in info.Upgrades)
 					{
-						if (!a.IsDead() && a.IsInWorld)
+						if (!um.AcceptsUpgrade(a, u))
-							GrantActorUpgrades(a);
+							continue;
+
+						if (info.Duration > 0)
+							um.GrantTimedUpgrade(a, u, info.Duration);
+						else
+							um.GrantUpgrade(a, u, this);
 					}
-				});
+				}
-			}
+			});
 
 			base.Activate(collector);
 		}

OpenRA.Mods.RA/GainsExperience.cs

@@ -102,10 +102,10 @@ namespace OpenRA.Mods.RA
 
 				Level++;
 
-				foreach (var up in self.TraitsImplementing<IUpgradable>())
+				var um = self.TraitOrDefault<UpgradeManager>();
+				if (um != null)
 					foreach (var u in upgrades)
-						if (up.AcceptsUpgrade(u))
+						um.GrantUpgrade(self, u, this);
-							up.UpgradeAvailable(self, u, true);
 
 				Sound.PlayNotification(self.World.Map.Rules, self.Owner, "Sounds", "LevelUp", self.Owner.Country.Race);
 				self.World.AddFrameEndTask(w => w.Add(new CrateEffect(self, "levelup", info.LevelUpPalette)));

OpenRA.Mods.RA/GlobalUpgradable.cs

@@ -17,7 +17,7 @@ using OpenRA.Traits;
 
 namespace OpenRA.Mods.RA
 {
-	public class GlobalUpgradableInfo : ITraitInfo
+	public class GlobalUpgradableInfo : ITraitInfo, Requires<UpgradeManagerInfo>
 	{
 		public readonly string[] Upgrades = { };
 		public readonly string[] Prerequisites = { };
@@ -28,34 +28,36 @@ namespace OpenRA.Mods.RA
 	public class GlobalUpgradable : INotifyAddedToWorld, INotifyRemovedFromWorld
 	{
 		readonly GlobalUpgradableInfo info;
-		readonly GlobalUpgradeManager manager;
+		readonly GlobalUpgradeManager globalManager;
+		readonly UpgradeManager manager;
 
-		public GlobalUpgradable(Actor actor, GlobalUpgradableInfo info)
+		public GlobalUpgradable(Actor self, GlobalUpgradableInfo info)
 		{
 			this.info = info;
-			manager = actor.Owner.PlayerActor.Trait<GlobalUpgradeManager>();
+			globalManager = self.Owner.PlayerActor.Trait<GlobalUpgradeManager>();
+			manager = self.Trait<UpgradeManager>();
 		}
 
 		public void AddedToWorld(Actor self)
 		{
 			if (info.Prerequisites.Any())
-				manager.Register(self, this, info.Prerequisites);
+				globalManager.Register(self, this, info.Prerequisites);
 		}
 
 		public void RemovedFromWorld(Actor self)
 		{
 			if (info.Prerequisites.Any())
-				manager.Unregister(self, this, info.Prerequisites);
+				globalManager.Unregister(self, this, info.Prerequisites);
 		}
 
 		public void PrerequisitesUpdated(Actor self, bool available)
 		{
-			var upgrades = self.TraitsImplementing<IUpgradable>();
+			foreach (var u in info.Upgrades)
-			foreach (var u in upgrades)
 			{
-				foreach (var t in info.Upgrades)
+				if (available)
-					if (u.AcceptsUpgrade(t))
+					manager.GrantUpgrade(self, u, this);
-						u.UpgradeAvailable(self, t, available);
+				else
+					manager.RevokeUpgrade(self, u, this);
 			}
 		}
 	}

OpenRA.Mods.RA/OpenRA.Mods.RA.csproj

@@ -564,6 +564,13 @@
     <Compile Include="Scripting\Properties\HarvesterProperties.cs" />
     <Compile Include="Scripting\Properties\HelicopterProperties.cs" />
     <Compile Include="Scripting\Properties\PlaneProperties.cs" />
+    <Compile Include="SupportPowers\GrantUpgradePower.cs" />
+    <Compile Include="Modifiers\UpgradeOverlay.cs" />
+    <Compile Include="TimedUpgradeBar.cs" />
+    <Compile Include="InvulnerabilityUpgrade.cs" />
+    <Compile Include="Warheads\TimedUpgradeWarhead.cs" />
+    <Compile Include="DisableUpgrade.cs" />
+    <Compile Include="UpgradeManager.cs" />
   </ItemGroup>
   <ItemGroup>
     <ProjectReference Include="..\OpenRA.Game\OpenRA.Game.csproj">

OpenRA.Mods.RA/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.RA
+{
+	[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.RA/UpgradeManager.cs

@@ -0,0 +1,134 @@
+#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.RA
+{
+	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--; }
+		}
+
+		readonly List<TimedUpgrade> timedUpgrades = new List<TimedUpgrade>();
+		readonly Dictionary<string, List<object>> sources = new Dictionary<string, List<object>>();
+		readonly Dictionary<string, List<Action<int, int>>> watchers = new Dictionary<string, List<Action<int, int>>>();
+		readonly Lazy<IEnumerable<IUpgradable>> upgradable;
+
+		public UpgradeManager(ActorInitializer init)
+		{
+			upgradable = Exts.Lazy(() => init.self.TraitsImplementing<IUpgradable>());
+		}
+
+		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);
+		}
+
+		public void GrantUpgrade(Actor self, string upgrade, object source)
+		{
+			List<object> ss;
+			if (!sources.TryGetValue(upgrade, out ss))
+			{
+				ss = new List<object>();
+				sources.Add(upgrade, ss);
+
+				foreach (var up in upgradable.Value)
+					if (up.AcceptsUpgrade(upgrade))
+						up.UpgradeAvailable(self, upgrade, true);
+			}
+
+			// Track the upgrade source so that the upgrade can be removed without conflicts
+			ss.Add(source);
+		}
+
+		public void RevokeUpgrade(Actor self, string upgrade, object source)
+		{
+			// This upgrade may have been granted by multiple sources
+			// We must be careful to only remove the upgrade after all
+			// sources have been revoked
+			List<object> ss;
+			if (!sources.TryGetValue(upgrade, out ss))
+				return;
+
+			ss.Remove(source);
+			if (!ss.Any())
+			{
+				foreach (var up in upgradable.Value)
+					if (up.AcceptsUpgrade(upgrade))
+						up.UpgradeAvailable(self, upgrade, false);
+
+				sources.Remove(upgrade);
+			}
+		}
+
+		public bool AcceptsUpgrade(Actor self, string upgrade)
+		{
+			return upgradable.Value.Any(up => up.AcceptsUpgrade(upgrade));
+		}
+
+		public void RegisterWatcher(string upgrade, Action<int, int> action)
+		{
+			if (!watchers.ContainsKey(upgrade))
+				watchers.Add(upgrade, new List<Action<int, int>>());
+
+			watchers[upgrade].Add(action);
+		}
+
+		public void Tick(Actor self)
+		{
+			foreach (var u in timedUpgrades)
+			{
+				u.Tick();
+				if (u.Remaining <= 0)
+					RevokeUpgrade(self, u.Upgrade, u);
+
+				List<Action<int, int>> actions;
+				if (watchers.TryGetValue(u.Upgrade, out actions))
+					foreach (var a in actions)
+						a(u.Duration, u.Remaining);
+			}
+
+			timedUpgrades.RemoveAll(u => u.Remaining <= 0);
+		}
+	}
+}

mods/cnc/rules/defaults.yaml

@@ -47,6 +47,7 @@
 		HealIfBelow: 1
 		DamageCooldown: 125
 		RequiresUpgrade: selfheal
+	UpgradeManager:
 
 ^Tank:
 	AppearsOnRadar:
@@ -100,6 +101,7 @@
 		HealIfBelow: 1
 		DamageCooldown: 125
 		RequiresUpgrade: selfheal
+	UpgradeManager:
 
 ^Helicopter:
 	AppearsOnRadar:
@@ -138,6 +140,7 @@
 		HealIfBelow: 1
 		DamageCooldown: 125
 		RequiresUpgrade: selfheal
+	UpgradeManager:
 
 ^Infantry:
 	AppearsOnRadar:
@@ -218,6 +221,7 @@
 		HealIfBelow: 1
 		DamageCooldown: 125
 		RequiresUpgrade: selfheal
+	UpgradeManager:
 
 ^CivInfantry:
 	Inherits: ^Infantry
@@ -324,6 +328,7 @@
 		HealIfBelow: 1
 		DamageCooldown: 125
 		RequiresUpgrade: selfheal
+	UpgradeManager:
 
 ^Ship:
 	AppearsOnRadar:
@@ -357,6 +362,7 @@
 		HealIfBelow: 1
 		DamageCooldown: 125
 		RequiresUpgrade: selfheal
+	UpgradeManager:
 
 ^Building:
 	AppearsOnRadar:

mods/d2k/rules/defaults.yaml

@@ -46,6 +46,7 @@
 		HealIfBelow: 1
 		DamageCooldown: 125
 		RequiresUpgrade: selfheal
+	UpgradeManager:
 
 ^Tank:
 	AppearsOnRadar:
@@ -95,6 +96,7 @@
 		HealIfBelow: 1
 		DamageCooldown: 125
 		RequiresUpgrade: selfheal
+	UpgradeManager:
 
 ^Husk:
 	Health:
@@ -221,6 +223,7 @@
 		HealIfBelow: 1
 		DamageCooldown: 125
 		RequiresUpgrade: selfheal
+	UpgradeManager:
 
 ^Plane:
 	AppearsOnRadar:
@@ -253,6 +256,7 @@
 		HealIfBelow: 1
 		DamageCooldown: 125
 		RequiresUpgrade: selfheal
+	UpgradeManager:
 
 ^Helicopter:
 	Inherits: ^Plane

mods/ra/rules/defaults.yaml

@@ -60,6 +60,7 @@
 		HealIfBelow: 1
 		DamageCooldown: 125
 		RequiresUpgrade: selfheal
+	UpgradeManager:
 
 ^Tank:
 	AppearsOnRadar:
@@ -123,6 +124,7 @@
 		HealIfBelow: 1
 		DamageCooldown: 125
 		RequiresUpgrade: selfheal
+	UpgradeManager:
 
 ^Infantry:
 	AppearsOnRadar:
@@ -207,6 +209,7 @@
 		HealIfBelow: 1
 		DamageCooldown: 125
 		RequiresUpgrade: selfheal
+	UpgradeManager:
 
 ^Ship:
 	AppearsOnRadar:
@@ -248,6 +251,7 @@
 		DamageCooldown: 125
 		RequiresUpgrade: selfheal
 	IronCurtainable:
+	UpgradeManager:
 
 ^Plane:
 	AppearsOnRadar:
@@ -292,6 +296,7 @@
 		DamageCooldown: 125
 		RequiresUpgrade: selfheal
 	IronCurtainable:
+	UpgradeManager:
 
 ^Helicopter:
 	Inherits: ^Plane
@@ -348,6 +353,7 @@
 	Demolishable:
 	ScriptTriggers:
 	IronCurtainable:
+	UpgradeManager:
 
 ^Defense:
 	Inherits: ^Building

mods/ts/rules/defaults.yaml

@@ -157,6 +157,7 @@
 	DeathSounds@ZAPPED:
 		DeathSound: Zapped
 		DeathTypes: 6
+	UpgradeManager:
 
 ^CivilianInfantry:
 	Inherits: ^Infantry
@@ -243,6 +244,7 @@
 	Explodes:
 		Weapon: UnitExplodeSmall
 		EmptyWeapon: UnitExplodeSmall
+	UpgradeManager:
 
 ^Tank:
 	AppearsOnRadar:
@@ -302,6 +304,7 @@
 	Explodes:
 		Weapon: UnitExplodeSmall
 		EmptyWeapon: UnitExplodeSmall
+	UpgradeManager:
 
 ^Helicopter:
 	AppearsOnRadar:
@@ -348,6 +351,7 @@
 	ScriptTriggers:
 	Guard:
 	Guardable:
+	UpgradeManager:
 
 ^BlossomTree:
 	Tooltip: