Support multiple Power traits

OpenRA.Mods.Common/Effects/PowerdownIndicator.cs

@@ -19,10 +19,12 @@ namespace OpenRA.Mods.Common.Effects
 	{
 		readonly Actor a;
 		readonly Animation anim;
+		readonly CanPowerDown canPowerDown;
 
 		public PowerdownIndicator(Actor a)
 		{
 			this.a = a;
+			canPowerDown = a.Trait<CanPowerDown>();
 
 			anim = new Animation(a.World, "poweroff");
 			anim.PlayRepeating("offline");
@@ -30,7 +32,7 @@ namespace OpenRA.Mods.Common.Effects
 
 		public void Tick(World world)
 		{
-			if (!a.IsInWorld || a.IsDead || !a.Trait<CanPowerDown>().Disabled)
+			if (!a.IsInWorld || a.IsDead || !canPowerDown.Disabled)
 				world.AddFrameEndTask(w => w.Remove(this));
 
 			anim.Tick();

OpenRA.Mods.Common/Scripting/Properties/PowerProperties.cs

@@ -9,6 +9,8 @@
 #endregion
 
 using System;
+using System.Collections.Generic;
+using System.Linq;
 using Eluant;
 using OpenRA.Mods.Common.Power;
 using OpenRA.Scripting;
@@ -56,18 +58,18 @@ namespace OpenRA.Mods.Common.Scripting
 	[ScriptPropertyGroup("Power")]
 	public class ActorPowerProperties : ScriptActorProperties, Requires<PowerInfo>
 	{
-		readonly PowerInfo pi;
+		readonly IEnumerable<Power.Power> power;
 
 		public ActorPowerProperties(ScriptContext context, Actor self)
 			: base(context, self)
 		{
-			pi = self.Info.Traits.GetOrDefault<PowerInfo>();
+			power = self.TraitsImplementing<Power.Power>();
 		}
 
 		[Desc("Returns the power drained/provided by this actor.")]
 		public int Power
 		{
-			get { return pi.Amount; }
+			get { return power.Sum(p => p.GetEnabledPower()); }
 		}
 	}
 }

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

@@ -16,18 +16,21 @@ namespace OpenRA.Mods.Common.Power
 	[Desc("The player can disable the power individually on this actor.")]
 	public class CanPowerDownInfo : UpgradableTraitInfo, ITraitInfo, Requires<PowerInfo>
 	{
+		[Desc("Restore power when this trait is disabled.")]
+		public readonly bool CancelWhenDisabled = false;
+
 		public object Create(ActorInitializer init) { return new CanPowerDown(init.self, this); }
 	}
 
-	public class CanPowerDown : UpgradableTrait<CanPowerDownInfo>, IPowerModifier, IResolveOrder, IDisable
+	public class CanPowerDown : UpgradableTrait<CanPowerDownInfo>, IPowerModifier, IResolveOrder, IDisable, INotifyOwnerChanged
 	{
 		[Sync] bool disabled = false;
-		readonly Power power;
+		PowerManager power;
 
 		public CanPowerDown(Actor self, CanPowerDownInfo info)
 			: base(info)
 		{
-			power = self.Trait<Power>();
+			power = self.Owner.PlayerActor.Trait<PowerManager>();
 		}
 
 		public bool Disabled { get { return disabled; } }
@@ -38,7 +41,7 @@ namespace OpenRA.Mods.Common.Power
 			{
 				disabled = !disabled;
 				Sound.PlayNotification(self.World.Map.Rules, self.Owner, "Sounds", disabled ? "EnablePower" : "DisablePower", self.Owner.Country.Race);
-				power.PlayerPower.UpdateActor(self);
+				power.UpdateActor(self);
 
 				if (disabled)
 					self.World.AddFrameEndTask(w => w.Add(new PowerdownIndicator(self)));
@@ -49,5 +52,19 @@ namespace OpenRA.Mods.Common.Power
 		{
 			return !IsTraitDisabled && disabled ? 0 : 100;
 		}
+
+		public void OnOwnerChanged(Actor self, Player oldOwner, Player newOwner)
+		{
+			power = newOwner.PlayerActor.Trait<PowerManager>();
+		}
+
+		protected override void UpgradeDisabled(Actor self)
+		{
+			if (!disabled || !Info.CancelWhenDisabled)
+				return;
+			disabled = false;
+			Sound.PlayNotification(self.World.Map.Rules, self.Owner, "Sounds", "EnablePower", self.Owner.Country.Race);
+			power.UpdateActor(self);
+		}
 	}
 }

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

@@ -18,14 +18,14 @@ namespace OpenRA.Mods.Common.Power
 		public object Create(ActorInitializer init) { return new ScalePowerWithHealth(init.self); }
 	}
 
-	public class ScalePowerWithHealth : IPowerModifier, INotifyDamage
+	public class ScalePowerWithHealth : IPowerModifier, INotifyDamage, INotifyOwnerChanged
 	{
-		readonly Power power;
 		readonly Health health;
+		PowerManager power;
 
 		public ScalePowerWithHealth(Actor self)
 		{
-			power = self.Trait<Power>();
+			power = self.Owner.PlayerActor.Trait<PowerManager>();
 			health = self.Trait<Health>();
 		}
 
@@ -34,9 +34,10 @@ namespace OpenRA.Mods.Common.Power
 			return 100 * health.HP / health.MaxHP;
 		}
 
-		public void Damaged(Actor self, AttackInfo e)
+		public void Damaged(Actor self, AttackInfo e) { power.UpdateActor(self); }
+		public void OnOwnerChanged(Actor self, Player oldOwner, Player newOwner)
 		{
-			power.PlayerPower.UpdateActor(self);
+			power = newOwner.PlayerActor.Trait<PowerManager>();
 		}
 	}
 }

OpenRA.Mods.D2k/Widgets/BuildPaletteWidget.cs

@@ -508,10 +508,11 @@ namespace OpenRA.Mods.D2k.Widgets
 				* ((lowpower) ? CurrentQueue.Info.LowPowerSlowdown : 1);
 			DrawRightAligned(WidgetUtils.FormatTime(time), pos + new int2(-5, 35), lowpower ? Color.Red : Color.White);
 
-			var pi = info.Traits.GetOrDefault<PowerInfo>();
+			var pis = info.Traits.WithInterface<PowerInfo>().Where(i => i.UpgradeMinEnabledLevel < 1);
-			if (pi != null)
+			var amount = pis.Sum(i => i.Amount);
-				DrawRightAligned("{1}{0}".F(pi.Amount, pi.Amount > 0 ? "+" : ""), pos + new int2(-5, 20),
+			if (pis != null)
-					((power.PowerProvided - power.PowerDrained) >= -pi.Amount || pi.Amount > 0) ? Color.White : Color.Red);
+				DrawRightAligned("{1}{0}".F(amount, amount > 0 ? "+" : ""), pos + new int2(-5, 20),
+					((power.PowerProvided - power.PowerDrained) >= -amount || amount > 0) ? Color.White : Color.Red);
 
 			p += new int2(5, 35);
 			if (!canBuildThis)

OpenRA.Mods.RA/AI/BaseBuilder.cs

@@ -140,8 +140,8 @@ namespace OpenRA.Mods.RA.AI
 			// First priority is to get out of a low power situation
 			if (playerPower.ExcessPower < 0)
 			{
-				var power = GetProducibleBuilding("Power", buildableThings, a => a.Traits.Get<PowerInfo>().Amount);
+				var power = GetProducibleBuilding("Power", buildableThings, a => a.Traits.WithInterface<PowerInfo>().Where(i => i.UpgradeMinEnabledLevel < 1).Sum(p => p.Amount));
-				if (power != null && power.Traits.Get<PowerInfo>().Amount > 0)
+				if (power != null && power.Traits.WithInterface<PowerInfo>().Where(i => i.UpgradeMinEnabledLevel < 1).Sum(p => p.Amount) > 0)
 				{
 					// TODO: Handle the case when of when we actually do need a power plant because we don't have enough but are also suffering from a power outage
 					if (playerPower.PowerOutageRemainingTicks <= 0)
@@ -204,12 +204,12 @@ namespace OpenRA.Mods.RA.AI
 
 				// Will this put us into low power?
 				var actor = world.Map.Rules.Actors[frac.Key];
-				var pi = actor.Traits.GetOrDefault<PowerInfo>();
+				var pis = actor.Traits.WithInterface<PowerInfo>().Where(i => i.UpgradeMinEnabledLevel < 1);
-				if (playerPower.ExcessPower < 0 || (pi != null && playerPower.ExcessPower < pi.Amount))
+				if (playerPower.ExcessPower < 0 || playerPower.ExcessPower < pis.Sum(pi => pi.Amount))
 				{
 					// Try building a power plant instead
-					var power = GetProducibleBuilding("Power", buildableThings, a => a.Traits.Get<PowerInfo>().Amount);
+					var power = GetProducibleBuilding("Power", buildableThings, a => a.Traits.WithInterface<PowerInfo>().Where(i => i.UpgradeMinEnabledLevel < 1).Sum(pi => pi.Amount));
-					if (power != null && power.Traits.Get<PowerInfo>().Amount > 0)
+					if (power != null && power.Traits.WithInterface<PowerInfo>().Where(i => i.UpgradeMinEnabledLevel < 1).Sum(pi => pi.Amount) > 0)
 					{
 						// TODO: Handle the case when of when we actually do need a power plant because we don't have enough but are also suffering from a power outage
 						if (playerPower.PowerOutageRemainingTicks > 0)

OpenRA.Mods.RA/Widgets/Logic/ProductionTooltipLogic.cs

@@ -58,7 +58,6 @@ namespace OpenRA.Mods.RA.Widgets.Logic
 				var tooltip = actor.Traits.Get<TooltipInfo>();
 				var buildable = actor.Traits.Get<BuildableInfo>();
 				var cost = actor.Traits.Get<ValuedInfo>().Cost;
-				var pi = actor.Traits.GetOrDefault<PowerInfo>();
 
 				nameLabel.GetText = () => tooltip.Name;
 
@@ -74,7 +73,7 @@ namespace OpenRA.Mods.RA.Widgets.Logic
 				var requiresString = prereqs.Any() ? requiresLabel.Text.F(prereqs.JoinWith(", ")) : "";
 				requiresLabel.GetText = () => requiresString;
 
-				var power = pi != null ? pi.Amount : 0;
+				var power = actor.Traits.WithInterface<PowerInfo>().Where(i => i.UpgradeMinEnabledLevel < 1).Sum(i => i.Amount);
 				var powerString = power.ToString();
 				powerLabel.GetText = () => powerString;
 				powerLabel.GetColor = () => ((pm.PowerProvided - pm.PowerDrained) >= -power || power > 0)