Add Upgrade support to Power/CanPowerDown

OpenRA.Mods.Common/Traits/Power/CanPowerDown.cs

@@ -14,17 +14,18 @@ using OpenRA.Traits;
 namespace OpenRA.Mods.Common.Power
 {
 	[Desc("The player can disable the power individually on this actor.")]
-	public class CanPowerDownInfo : ITraitInfo, Requires<PowerInfo>
+	public class CanPowerDownInfo : UpgradableTraitInfo, ITraitInfo, Requires<PowerInfo>
 	{
-		public object Create(ActorInitializer init) { return new CanPowerDown(init.self); }
+		public object Create(ActorInitializer init) { return new CanPowerDown(init.self, this); }
 	}
 
-	public class CanPowerDown : IPowerModifier, IResolveOrder, IDisable, ISync
+	public class CanPowerDown : UpgradableTrait<CanPowerDownInfo>, IPowerModifier, IResolveOrder, IDisable
 	{
 		[Sync] bool disabled = false;
 		readonly Power power;
 
-		public CanPowerDown(Actor self)
+		public CanPowerDown(Actor self, CanPowerDownInfo info)
+			: base(info)
 		{
 			power = self.Trait<Power>();
 		}
@@ -33,7 +34,7 @@ namespace OpenRA.Mods.Common.Power
 
 		public void ResolveOrder(Actor self, Order order)
 		{
-			if (order.OrderString == "PowerDown")
+			if (!IsTraitDisabled && order.OrderString == "PowerDown")
 			{
 				disabled = !disabled;
 				Sound.PlayNotification(self.World.Map.Rules, self.Owner, "Sounds", disabled ? "EnablePower" : "DisablePower", self.Owner.Country.Race);
@@ -46,7 +47,7 @@ namespace OpenRA.Mods.Common.Power
 
 		public int GetPowerModifier()
 		{
-			return disabled ? 0 : 100;
+			return !IsTraitDisabled && disabled ? 0 : 100;
 		}
 	}
 }

OpenRA.Mods.Common/Traits/Power/Player/PowerManager.cs

@@ -45,59 +45,38 @@ namespace OpenRA.Mods.Common.Power
 			this.self = self;
 			this.info = info;
 
-			self.World.ActorAdded += UpdateActor;
-			self.World.ActorRemoved += RemoveActor;
-
 			devMode = self.Trait<DeveloperMode>();
 			wasHackEnabled = devMode.UnlimitedPower;
 		}
 
 		public void UpdateActor(Actor a)
 		{
-			UpdateActors(new[] { a });
-		}
-
-		public void UpdateActors(IEnumerable<Actor> actors)
-		{
-			foreach (var a in actors)
-			{
-				if (a.Owner != self.Owner)
-					return;
-
-				var power = a.TraitOrDefault<Power>();
-				if (power == null)
-					return;
-
-				powerDrain[a] = power.GetCurrentPower();
-			}
-			UpdateTotals();
-		}
-
-		void RemoveActor(Actor a)
-		{
-			if (a.Owner != self.Owner || !a.HasTrait<Power>())
+			int old;
+			powerDrain.TryGetValue(a, out old); // old is 0 if a is not in powerDrain
+			var amount = a.TraitsImplementing<Power>().Where(t => !t.IsTraitDisabled).Aggregate(0, (v, p) => v + p.GetEnabledPower());
+			powerDrain[a] = amount;
+			if (amount == old || devMode.UnlimitedPower)
 				return;
-
-			powerDrain.Remove(a);
-			UpdateTotals();
+			if (old > 0)
+				totalProvided -= old;
+			else if (old < 0)
+				totalDrained += old;
+			if (amount > 0)
+				totalProvided += amount;
+			else if (amount < 0)
+				totalDrained -= amount;
 		}
 
-		public void UpdateTotals()
+		public void RemoveActor(Actor a)
 		{
-			totalProvided = 0;
-			totalDrained = 0;
-
-			foreach (var kv in powerDrain)
-			{
-				var p = kv.Value;
-				if (p > 0)
-					totalProvided += p;
-				else
-					totalDrained -= p;
-			}
-
-			if (devMode.UnlimitedPower)
-				totalProvided = 1000000;
+			int amount;
+			if (!powerDrain.TryGetValue(a, out amount))
+				return;
+			if (amount > 0)
+				totalProvided -= amount;
+			else if (amount < 0)
+				totalDrained += amount;
+			powerDrain.Remove(a);
 		}
 
 		int nextPowerAdviceTime = 0;
@@ -108,7 +87,16 @@ namespace OpenRA.Mods.Common.Power
 		{
 			if (wasHackEnabled != devMode.UnlimitedPower)
 			{
-				UpdateTotals();
+				totalProvided = 0;
+				totalDrained = 0;
+
+				if (!devMode.UnlimitedPower)
+					foreach (var kv in powerDrain)
+						if (kv.Value > 0)
+							totalProvided += kv.Value;
+						else if (kv.Value < 0)
+							totalDrained -= kv.Value;
+
 				wasHackEnabled = devMode.UnlimitedPower;
 			}
 
@@ -150,7 +138,8 @@ namespace OpenRA.Mods.Common.Power
 				.Select(tp => tp.Actor)
 				.Where(a => !a.IsDead && a.IsInWorld && a.Owner == self.Owner);
 
-			UpdateActors(actors);
+			foreach (var a in actors)
+				UpdateActor(a);
 		}
 	}
 }

OpenRA.Mods.Common/Traits/Power/Power.cs

@@ -14,7 +14,7 @@ using OpenRA.Traits;
 
 namespace OpenRA.Mods.Common.Power
 {
-	public class PowerInfo : ITraitInfo
+	public class PowerInfo : UpgradableTraitInfo, ITraitInfo
 	{
 		[Desc("If negative, it will drain power. If positive, it will provide power.")]
 		public readonly int Amount = 0;
@@ -22,28 +22,33 @@ namespace OpenRA.Mods.Common.Power
 		public object Create(ActorInitializer init) { return new Power(init.self, this); }
 	}
 
-	public class Power : INotifyOwnerChanged
+	public class Power : UpgradableTrait<PowerInfo>, INotifyAddedToWorld, INotifyRemovedFromWorld, INotifyOwnerChanged
 	{
-		readonly PowerInfo info;
 		readonly Lazy<IPowerModifier[]> powerModifiers;
 
 		public PowerManager PlayerPower { get; private set; }
 
-		public int GetCurrentPower()
+		public int GetEnabledPower()
 		{
-			return Util.ApplyPercentageModifiers(info.Amount, powerModifiers.Value.Select(m => m.GetPowerModifier()));
+			return Util.ApplyPercentageModifiers(Info.Amount, powerModifiers.Value.Select(m => m.GetPowerModifier()));
 		}
 
 		public Power(Actor self, PowerInfo info)
+			: base(info)
 		{
-			this.info = info;
 			PlayerPower = self.Owner.PlayerActor.Trait<PowerManager>();
 			powerModifiers = Exts.Lazy(() => self.TraitsImplementing<IPowerModifier>().ToArray());
 		}
 		
+		protected override void UpgradeEnabled(Actor self) { PlayerPower.UpdateActor(self); }
+		protected override void UpgradeDisabled(Actor self) { PlayerPower.UpdateActor(self); }
+		public void AddedToWorld(Actor self) { PlayerPower.UpdateActor(self); }
+		public void RemovedFromWorld(Actor self) { PlayerPower.RemoveActor(self); }
 		public void OnOwnerChanged(Actor self, Player oldOwner, Player newOwner)
 		{
+			PlayerPower.RemoveActor(self);
 			PlayerPower = newOwner.PlayerActor.Trait<PowerManager>();
+			PlayerPower.UpdateActor(self);
 		}
 	}
 }