Fix GrantConditionOnDeploy to support multiple sprite bodies

OpenRA.Mods.Common/Traits/Render/WithMakeAnimation.cs

@@ -10,28 +10,32 @@
 #endregion
 
 using System;
+using System.Linq;
 using OpenRA.Activities;
 using OpenRA.Traits;
 
 namespace OpenRA.Mods.Common.Traits.Render
 {
-	[Desc("Replaces the sprite during construction.")]
+	[Desc("Replaces the sprite during construction/deploy/undeploy.")]
 	public class WithMakeAnimationInfo : ITraitInfo, Requires<WithSpriteBodyInfo>
 	{
-		[Desc("Sequence name to use")]
+		[Desc("Sequence name to use.")]
 		[SequenceReference] public readonly string Sequence = "make";
 
 		[GrantedConditionReference]
 		[Desc("The condition to grant to self while the make animation is playing.")]
 		public readonly string Condition = null;
 
+		[Desc("Apply to sprite bodies with these names.")]
+		public readonly string[] BodyNames = { "body" };
+
 		public object Create(ActorInitializer init) { return new WithMakeAnimation(init, this); }
 	}
 
-	public class WithMakeAnimation : INotifyCreated
+	public class WithMakeAnimation : INotifyCreated, INotifyDeployTriggered
 	{
 		readonly WithMakeAnimationInfo info;
-		readonly WithSpriteBody wsb;
+		readonly WithSpriteBody[] wsbs;
 
 		ConditionManager conditionManager;
 		int token = ConditionManager.InvalidConditionToken;
@@ -40,7 +44,7 @@ namespace OpenRA.Mods.Common.Traits.Render
 		{
 			this.info = info;
 			var self = init.Self;
-			wsb = self.Trait<WithSpriteBody>();
+			wsbs = self.TraitsImplementing<WithSpriteBody>().Where(w => info.BodyNames.Contains(w.Info.Name)).ToArray();
 		}
 
 		void INotifyCreated.Created(Actor self)
@@ -56,6 +60,11 @@ namespace OpenRA.Mods.Common.Traits.Render
 			if (conditionManager != null && !string.IsNullOrEmpty(info.Condition) && token == ConditionManager.InvalidConditionToken)
 				token = conditionManager.GrantCondition(self, info.Condition);
 
+			var wsb = wsbs.FirstOrDefault(Exts.IsTraitEnabled);
+
+			if (wsb == null)
+				return;
+
 			wsb.PlayCustomAnimation(self, info.Sequence, () =>
 			{
 				if (token != ConditionManager.InvalidConditionToken)
@@ -71,6 +80,11 @@ namespace OpenRA.Mods.Common.Traits.Render
 			if (conditionManager != null && !string.IsNullOrEmpty(info.Condition) && token == ConditionManager.InvalidConditionToken)
 				token = conditionManager.GrantCondition(self, info.Condition);
 
+			var wsb = wsbs.FirstOrDefault(Exts.IsTraitEnabled);
+
+			if (wsb == null)
+				return;
+
 			wsb.PlayCustomAnimationBackwards(self, info.Sequence, () =>
 			{
 				if (token != ConditionManager.InvalidConditionToken)
@@ -85,11 +99,14 @@ namespace OpenRA.Mods.Common.Traits.Render
 		{
 			Reverse(self, () =>
 			{
+				var wsb = wsbs.FirstOrDefault(Exts.IsTraitEnabled);
+
 				// HACK: The actor remains alive and active for one tick before the followup activity
 				// (sell/transform/etc) runs. This causes visual glitches that we attempt to minimize
 				// by forcing the animation to frame 0 and regranting the make condition.
 				// These workarounds will break the actor if the followup activity doesn't dispose it!
-				wsb.DefaultAnimation.PlayFetchIndex(info.Sequence, () => 0);
+				if (wsb != null)
+					wsb.DefaultAnimation.PlayFetchIndex(info.Sequence, () => 0);
 
 				if (conditionManager != null && !string.IsNullOrEmpty(info.Condition))
 					token = conditionManager.GrantCondition(self, info.Condition);
@@ -97,5 +114,55 @@ namespace OpenRA.Mods.Common.Traits.Render
 				self.QueueActivity(queued, activity);
 			});
 		}
+
+		// TODO: Make this use Forward instead
+		void INotifyDeployTriggered.Deploy(Actor self)
+		{
+			var notified = false;
+
+			foreach (var wsb in wsbs)
+			{
+				if (wsb.IsTraitDisabled)
+					continue;
+
+				var notify = self.TraitsImplementing<INotifyDeployComplete>();
+				wsb.PlayCustomAnimation(self, info.Sequence, () =>
+				{
+					if (notified)
+						return;
+
+					foreach (var n in notify)
+					{
+						n.FinishedDeploy(self);
+						notified = true;
+					}
+				});
+			}
+		}
+
+		// TODO: Make this use Reverse instead
+		void INotifyDeployTriggered.Undeploy(Actor self)
+		{
+			var notified = false;
+
+			foreach (var wsb in wsbs)
+			{
+				if (wsb.IsTraitDisabled)
+					continue;
+
+				var notify = self.TraitsImplementing<INotifyDeployComplete>();
+				wsb.PlayCustomAnimationBackwards(self, info.Sequence, () =>
+				{
+					if (notified)
+						return;
+
+					foreach (var n in notify)
+					{
+						n.FinishedUndeploy(self);
+						notified = true;
+					}
+				});
+			}
+		}
 	}
 }