Remove airstrike/paradrop beacon when the whole squad is shot down

OpenRA.Mods.Common/Traits/SupportPowers/AirstrikePower.cs

@@ -44,8 +44,13 @@ namespace OpenRA.Mods.Common.Traits
 
 	public class AirstrikePower : SupportPower
 	{
+		readonly AirstrikePowerInfo info;
+
 		public AirstrikePower(Actor self, AirstrikePowerInfo info)
-			: base(self, info) { }
+			: base(self, info)
+		{
+			this.info = info;
+		}
 
 		public override void Activate(Actor self, Order order, SupportPowerManager manager)
 		{
@@ -56,8 +61,6 @@ namespace OpenRA.Mods.Common.Traits
 
 		public void SendAirstrike(Actor self, WPos target, bool randomize = true, int attackFacing = 0)
 		{
-			var info = Info as AirstrikePowerInfo;
-
 			if (randomize)
 				attackFacing = 256 * self.World.SharedRandom.Next(info.QuantizedFacings) / info.QuantizedFacings;
 
@@ -87,14 +90,7 @@ namespace OpenRA.Mods.Common.Traits
 					});
 				}
 
-				if (beacon != null)
+				RemoveBeacon(beacon);
-				{
-					self.World.AddFrameEndTask(w =>
-					{
-						w.Remove(beacon);
-						beacon = null;
-					});
-				}
 
 				aircraftInRange[a] = true;
 			};
@@ -105,23 +101,18 @@ namespace OpenRA.Mods.Common.Traits
 
 				// Remove the camera when the final plane leaves the target area
 				if (!aircraftInRange.Any(kv => kv.Value))
+					RemoveCamera(camera);
+			};
+
+			Action<Actor> onRemovedFromWorld = a =>
+			{
+				// Checking for attack range is not relevant here because
+				// aircraft may be shot down before entering. Thus we remove
+				// the camera and beacon only if the whole squad is dead.
+				if (aircraftInRange.All(kv => kv.Key.IsDead))
 				{
-					if (camera != null)
+					RemoveCamera(camera);
-					{
+					RemoveBeacon(beacon);
-						camera.QueueActivity(new Wait(info.CameraRemoveDelay));
-						camera.QueueActivity(new RemoveSelf());
-					}
-
-					camera = null;
-
-					if (beacon != null)
-					{
-						self.World.AddFrameEndTask(w =>
-						{
-							w.Remove(beacon);
-							beacon = null;
-						});
-					}
 				}
 			};
 
@@ -152,7 +143,7 @@ namespace OpenRA.Mods.Common.Traits
 					attack.SetTarget(w, target + targetOffset);
 					attack.OnEnteredAttackRange += onEnterRange;
 					attack.OnExitedAttackRange += onExitRange;
-					attack.OnRemovedFromWorld += onExitRange;
+					attack.OnRemovedFromWorld += onRemovedFromWorld;
 
 					a.QueueActivity(new Fly(a, Target.FromPos(target + spawnOffset)));
 					a.QueueActivity(new Fly(a, Target.FromPos(finishEdge + spawnOffset)));
@@ -183,5 +174,27 @@ namespace OpenRA.Mods.Common.Traits
 				}
 			});
 		}
+
+		void RemoveCamera(Actor camera)
+		{
+			if (camera == null)
+				return;
+
+			camera.QueueActivity(new Wait(info.CameraRemoveDelay));
+			camera.QueueActivity(new RemoveSelf());
+			camera = null;
+		}
+
+		void RemoveBeacon(Beacon beacon)
+		{
+			if (beacon == null)
+				return;
+
+			Self.World.AddFrameEndTask(w =>
+			{
+				w.Remove(beacon);
+				beacon = null;
+			});
+		}
 	}
 }

OpenRA.Mods.Common/Traits/SupportPowers/ParatroopersPower.cs

@@ -58,7 +58,13 @@ namespace OpenRA.Mods.Common.Traits
 
 	public class ParatroopersPower : SupportPower
 	{
-		public ParatroopersPower(Actor self, ParatroopersPowerInfo info) : base(self, info) { }
+		readonly ParatroopersPowerInfo info;
+
+		public ParatroopersPower(Actor self, ParatroopersPowerInfo info)
+			: base(self, info)
+		{
+			this.info = info;
+		}
 
 		public override void Activate(Actor self, Order order, SupportPowerManager manager)
 		{
@@ -107,14 +113,7 @@ namespace OpenRA.Mods.Common.Traits
 					});
 				}
 
-				if (beacon != null)
+				RemoveBeacon(beacon);
-				{
-					self.World.AddFrameEndTask(w =>
-					{
-						w.Remove(beacon);
-						beacon = null;
-					});
-				}
 
 				if (!aircraftInRange.Any(kv => kv.Value))
 					Game.Sound.PlayNotification(self.World.Map.Rules, self.Owner, "Speech",
@@ -129,23 +128,18 @@ namespace OpenRA.Mods.Common.Traits
 
 				// Remove the camera when the final plane leaves the target area
 				if (!aircraftInRange.Any(kv => kv.Value))
+					RemoveCamera(camera);
+			};
+
+			Action<Actor> onRemovedFromWorld = a =>
+			{
+				// Checking for attack range is not relevant here because
+				// aircraft may be shot down before entering. Thus we remove
+				// the camera and beacon only if the whole squad is dead.
+				if (aircraftInRange.All(kv => kv.Key.IsDead))
 				{
-					if (camera != null)
+					RemoveCamera(camera);
-					{
+					RemoveBeacon(beacon);
-						camera.QueueActivity(new Wait(info.CameraRemoveDelay));
-						camera.QueueActivity(new RemoveSelf());
-					}
-
-					camera = null;
-
-					if (beacon != null)
-					{
-						self.World.AddFrameEndTask(w =>
-						{
-							w.Remove(beacon);
-							beacon = null;
-						});
-					}
 				}
 			};
 
@@ -187,7 +181,7 @@ namespace OpenRA.Mods.Common.Traits
 					drop.SetLZ(w.Map.CellContaining(target + targetOffset), !info.AllowImpassableCells);
 					drop.OnEnteredDropRange += onEnterRange;
 					drop.OnExitedDropRange += onExitRange;
-					drop.OnRemovedFromWorld += onExitRange;
+					drop.OnRemovedFromWorld += onRemovedFromWorld;
 
 					var cargo = a.Trait<Cargo>();
 					var passengers = units.Skip(added).Take(passengersPerPlane);
@@ -227,5 +221,27 @@ namespace OpenRA.Mods.Common.Traits
 
 			return units.ToArray();
 		}
+
+		void RemoveCamera(Actor camera)
+		{
+			if (camera == null)
+				return;
+
+			camera.QueueActivity(new Wait(info.CameraRemoveDelay));
+			camera.QueueActivity(new RemoveSelf());
+			camera = null;
+		}
+
+		void RemoveBeacon(Beacon beacon)
+		{
+			if (beacon == null)
+				return;
+
+			Self.World.AddFrameEndTask(w =>
+			{
+				w.Remove(beacon);
+				beacon = null;
+			});
+		}
 	}
 }