From 42f3fa74395c28e2fcf4f7d3b464a4b9cfb6d75b Mon Sep 17 00:00:00 2001
From: Paul Chote <pchote@users.noreply.github.com>
Date: Sun, 29 Mar 2015 18:44:27 +0100
Subject: [PATCH] Add upgrade support to WithMuzzleFlash.

---
 OpenRA.Mods.Common/Traits/Render/WithMuzzleFlash.cs | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/OpenRA.Mods.Common/Traits/Render/WithMuzzleFlash.cs b/OpenRA.Mods.Common/Traits/Render/WithMuzzleFlash.cs
index 8e8a4eea81..253739ce95 100644
--- a/OpenRA.Mods.Common/Traits/Render/WithMuzzleFlash.cs
+++ b/OpenRA.Mods.Common/Traits/Render/WithMuzzleFlash.cs
@@ -17,7 +17,7 @@ using OpenRA.Traits;
 namespace OpenRA.Mods.Common.Traits
 {
 	[Desc("Renders the MuzzleSequence from the Armament trait.")]
-	class WithMuzzleFlashInfo : ITraitInfo, Requires<RenderSpritesInfo>, Requires<AttackBaseInfo>, Requires<ArmamentInfo>
+	class WithMuzzleFlashInfo : UpgradableTraitInfo, ITraitInfo, Requires<RenderSpritesInfo>, Requires<AttackBaseInfo>, Requires<ArmamentInfo>
 	{
 		[Desc("Ignore the weapon position, and always draw relative to the center of the actor")]
 		public readonly bool IgnoreOffset = false;
@@ -25,13 +25,14 @@ namespace OpenRA.Mods.Common.Traits
 		public object Create(ActorInitializer init) { return new WithMuzzleFlash(init.Self, this); }
 	}
 
-	class WithMuzzleFlash : INotifyAttack, IRender, ITick
+	class WithMuzzleFlash : UpgradableTrait<WithMuzzleFlashInfo>, INotifyAttack, IRender, ITick
 	{
 		Dictionary<Barrel, bool> visible = new Dictionary<Barrel, bool>();
 		Dictionary<Barrel, AnimationWithOffset> anims = new Dictionary<Barrel, AnimationWithOffset>();
 		Func<int> getFacing;
 
 		public WithMuzzleFlash(Actor self, WithMuzzleFlashInfo info)
+			: base(info)
 		{
 			var render = self.Trait<RenderSprites>();
 			var facing = self.TraitOrDefault<IFacing>();
@@ -64,7 +65,7 @@ namespace OpenRA.Mods.Common.Traits
 					anims.Add(barrel,
 						new AnimationWithOffset(muzzleFlash,
 							() => info.IgnoreOffset ? WVec.Zero : armClosure.MuzzleOffset(self, barrel),
-							() => !visible[barrel],
+							() => IsTraitDisabled || !visible[barrel],
 							() => false,
 							p => WithTurret.ZOffsetFromCenter(self, p, 2)));
 				}