Make AttractsWorms upgradable

OpenRA.Mods.D2k/Traits/AttractsWorms.cs

@@ -8,12 +8,13 @@
 */
 #endregion
 
+using OpenRA.Mods.Common.Traits;
 using OpenRA.Traits;
 
 namespace OpenRA.Mods.D2k.Traits
 {
 	[Desc("This actor makes noise, which causes them to be targeted by actors with the Sandworm trait.")]
-	public class AttractsWormsInfo : ITraitInfo
+	public class AttractsWormsInfo : UpgradableTraitInfo, ITraitInfo
 	{
 		[Desc("How much noise this actor produces.")]
 		public readonly int Intensity = 0;
@@ -27,22 +28,23 @@ namespace OpenRA.Mods.D2k.Traits
 		[Desc("Ranges at which each Falloff step is defined. Overrides Spread.")]
 		public WRange[] Range = null;
 
-		public object Create(ActorInitializer init) { return new AttractsWorms(this); }
+		public object Create(ActorInitializer init) { return new AttractsWorms(init, this); }
 	}
 
-	public class AttractsWorms
+	public class AttractsWorms : UpgradableTrait<AttractsWormsInfo>
 	{
-		public readonly AttractsWormsInfo Info;
+		readonly Actor self;
 
-		public AttractsWorms(AttractsWormsInfo info)
+		public AttractsWorms(ActorInitializer init, AttractsWormsInfo info)
+			: base(info)
 		{
-			Info = info;
+			self = init.Self;
 
 			if (info.Range == null)
 				info.Range = Exts.MakeArray(info.Falloff.Length, i => i * info.Spread);
 		}
 
-		public int GetNoisePercentageAtDistance(int distance)
+		int GetNoisePercentageAtDistance(int distance)
 		{
 			var inner = Info.Range[0].Range;
 			for (var i = 1; i < Info.Range.Length; i++)
@@ -56,5 +58,23 @@ namespace OpenRA.Mods.D2k.Traits
 
 			return 0;
 		}
+
+		public WVec AttractionAtPosition(WPos pos)
+		{
+			if (IsTraitDisabled)
+				return WVec.Zero;
+
+			var distance = self.CenterPosition - pos;
+			var length = distance.Length;
+
+			// Actor is too far to hear anything.
+			if (length > Info.Range[Info.Range.Length - 1].Range)
+				return WVec.Zero;
+
+			var direction = 1024 * distance / length;
+			var percentage = GetNoisePercentageAtDistance(length);
+
+			return direction * Info.Intensity * percentage / 100;
+		}
 	}
 }

OpenRA.Mods.D2k/Traits/Sandworm.cs

@@ -9,6 +9,7 @@
 #endregion
 
 using System;
+using System.Linq;
 using OpenRA.Mods.Common.Traits;
 using OpenRA.Traits;
 
@@ -91,42 +92,27 @@ namespace OpenRA.Mods.D2k.Traits
 		{
 			targetCountdown = Info.TargetRescanInterval;
 
-			var actorsInRange = self.World.FindActorsInCircle(self.CenterPosition, Info.MaxSearchRadius);
+			// If close enough, we don't care about other actors.
-			var noiseDirection = WVec.Zero;
+			var target = self.World.FindActorsInCircle(self.CenterPosition, Info.IgnoreNoiseAttackRange).FirstOrDefault(x => x.HasTrait<AttractsWorms>());
-
+			if (target != null)
-			foreach (var actor in actorsInRange)
-			{
-				if (!actor.IsInWorld)
-					continue;
-
-				// TODO: Test if we really want to ignore actors that are on rock
-				if (!mobile.Value.CanEnterCell(actor.Location, null, false))
-					continue;
-
-				var noise = actor.TraitOrDefault<AttractsWorms>();
-				if (noise == null)
-					continue;
-
-				var distance = actor.CenterPosition - self.CenterPosition;
-				var length = distance.Length;
-
-				// Actor is too far to be heard
-				if (noise.Info.Range[noise.Info.Range.Length - 1].Range < length)
-					continue;
-
-				// If close enough, we don't care about other actors
-				if (length <= Info.IgnoreNoiseAttackRange.Range)
 			{
 				self.CancelActivity();
-					attackTrait.Value.ResolveOrder(self, new Order("Attack", actor, true) { TargetActor = actor });
+				attackTrait.Value.ResolveOrder(self, new Order("Attack", target, true) { TargetActor = target });
 				return;
 			}
 
-				var direction = 1024 * distance / length;
+			Func<Actor, bool> isValidTarget = a =>
-				var percentage = noise.GetNoisePercentageAtDistance(length);
+			{
+				if (!a.HasTrait<AttractsWorms>())
+					return false;
 
-				noiseDirection += direction * noise.Info.Intensity * percentage / 100;
+				return mobile.Value.CanEnterCell(a.Location, null, false);
-			}
+			};
+
+			var actorsInRange = self.World.FindActorsInCircle(self.CenterPosition, Info.MaxSearchRadius)
+				.Where(isValidTarget).SelectMany(a => a.TraitsImplementing<AttractsWorms>());
+
+			var noiseDirection = actorsInRange.Aggregate(WVec.Zero, (a, b) => a + b.AttractionAtPosition(self.CenterPosition));
 
 			// No target was found
 			if (noiseDirection == WVec.Zero)