Touched up selection functionality.

OpenRA.Game/Widgets/WorldInteractionControllerWidget.cs

@@ -40,15 +40,17 @@ public override void Draw()
 		{
 			if (!IsDragging)
 			{
-				foreach (var u in SelectActorsInBoxWithDeadzone(World, lastMousePosition, lastMousePosition, _ => true))
+				// Render actors under the mouse pointer
+				foreach (var u in SelectActorsInBoxWithDeadzone(World, lastMousePosition, lastMousePosition))
 					worldRenderer.DrawRollover(u);
 
 				return;
 			}
 
+			// Render actors in the dragbox
 			var selbox = SelectionBox;
 			Game.Renderer.WorldLineRenderer.DrawRect(selbox.Value.First.ToFloat2(), selbox.Value.Second.ToFloat2(), Color.White);
-			foreach (var u in SelectActorsInBoxWithDeadzone(World, selbox.Value.First, selbox.Value.Second, _ => true))
+			foreach (var u in SelectActorsInBoxWithDeadzone(World, selbox.Value.First, selbox.Value.Second))
 				worldRenderer.DrawRollover(u);
 		}
 
@@ -113,15 +115,15 @@ public override bool HandleMouseInput(MouseInput mi)
 						{
 							// Select actors on the screen that have the same selection class as the already selected actor
 							var sc = unit.Trait<Selectable>().Class;
-							var newSelection = SelectActorsInBox(World, worldRenderer.Viewport.TopLeft, worldRenderer.Viewport.BottomRight,
-								a => a.Owner == unit.Owner && a.Trait<Selectable>().Class == sc);
+							var newSelection = SelectActorsOnScreen(World, worldRenderer, new HashSet<string> { sc }, unit.Owner);
 
 							World.Selection.Combine(World, newSelection, true, false);
 						}
 					}
 					else if (dragStart.HasValue)
 					{
-						var newSelection = SelectActorsInBoxWithDeadzone(World, dragStart.Value, xy, _ => true);
+						// Select actors in the dragbox
+						var newSelection = SelectActorsInBoxWithDeadzone(World, dragStart.Value, xy);
 						World.Selection.Combine(World, newSelection, mi.Modifiers.HasModifier(Modifiers.Shift), dragStart == xy);
 					}
 				}
@@ -232,8 +234,8 @@ public override bool HandleKeyPress(KeyInput e)
 					World.SetPauseState(!World.Paused);
 				else if (key == Game.Settings.Keys.SelectAllUnitsKey)
 				{
-					var ownUnitsOnScreen = SelectActorsInBox(World, worldRenderer.Viewport.TopLeft, worldRenderer.Viewport.BottomRight,
-						a => a.Owner == player);
+					// Select actors on the screen which belong to the current player
+					var ownUnitsOnScreen = SelectActorsOnScreen(World, worldRenderer, null, player);
 					World.Selection.Combine(World, ownUnitsOnScreen, false, false);
 				}
 				else if (key == Game.Settings.Keys.SelectUnitsByTypeKey)
@@ -245,16 +247,15 @@ public override bool HandleKeyPress(KeyInput e)
 						.ToHashSet();
 
 					// Select actors on the screen that have the same selection class as one of the already selected actors
-					Func<Actor, bool> cond = a => a.Owner == player && selectedClasses.Contains(a.Trait<Selectable>().Class);
-					var newSelection = SelectActorsInBox(World, worldRenderer.Viewport.TopLeft, worldRenderer.Viewport.BottomRight, cond).ToList();
+					var newSelection = SelectActorsOnScreen(World, worldRenderer, selectedClasses, player).ToList();
 
 					// Check if selecting actors on the screen has selected new units
 					if (newSelection.Count() > World.Selection.Actors.Count())
 						Game.Debug("Selected across screen");
 					else
 					{
 						// Select actors in the world that have the same selection class as one of the already selected actors
-						newSelection = World.ActorMap.ActorsInWorld().Where(x => x.HasTrait<Selectable>() && x.Trait<Selectable>().Info.Selectable && cond(x)).ToList();
+						newSelection = SelectActorsInWorld(World, selectedClasses, player).ToList();
 						Game.Debug("Selected across map");
 					}
 
@@ -269,18 +270,41 @@ public override bool HandleKeyPress(KeyInput e)
 			return false;
 		}
 
-		static IEnumerable<Actor> SelectActorsInBoxWithDeadzone(World world, int2 a, int2 b, Func<Actor, bool> cond)
+		static IEnumerable<Actor> SelectActorsOnScreen(World world, WorldRenderer wr, IEnumerable<string> sc, Player player)
 		{
-			if (a == b || (a - b).Length > Game.Settings.Game.SelectionDeadzone)
-				return SelectActorsInBox(world, a, b, cond);
-			else
-				return SelectActorsInBox(world, b, b, cond);
+			return SelectActorsByPlayerByClass(world.ScreenMap.ActorsInBox(wr.Viewport.TopLeft, wr.Viewport.BottomRight), sc, player);
 		}
 
-		static IEnumerable<Actor> SelectActorsInBox(World world, int2 a, int2 b, Func<Actor, bool> cond)
+		static IEnumerable<Actor> SelectActorsInWorld(World world, IEnumerable<string> sc, Player player)
 		{
+			return SelectActorsByPlayerByClass(world.ActorMap.ActorsInWorld(), sc, player);
+		}
+
+		static IEnumerable<Actor> SelectActorsByPlayerByClass(IEnumerable<Actor> sel, IEnumerable<string> sc, Player player)
+		{
+			return sel.Where(a =>
+			{
+				if (a.Owner != player)
+					return false;
+				var s = a.TraitOrDefault<Selectable>();
+
+				// sc == null means that units, that meet all other criteria, get selected
+				return s != null && s.Info.Selectable && (sc == null || sc.Contains(s.Class));
+			});
+		}
+
+		static IEnumerable<Actor> SelectActorsInBoxWithDeadzone(World world, int2 a, int2 b)
+		{
+			// For dragboxes that are too small, shrink the dragbox to a single point (point b)
+			if ((a - b).Length <= Game.Settings.Game.SelectionDeadzone)
+				a = b;
+
 			return world.ScreenMap.ActorsInBox(a, b)
-				.Where(x => x.HasTrait<Selectable>() && x.Trait<Selectable>().Info.Selectable && (x.Owner.IsAlliedWith(world.RenderPlayer) || !world.FogObscures(x)) && cond(x))
+				.Where(x =>
+				{
+					var s = x.TraitOrDefault<Selectable>();
+					return s != null && s.Info.Selectable && !world.FogObscures(x);
+				})
 				.GroupBy(x => x.GetSelectionPriority())
 				.OrderByDescending(g => g.Key)
 				.Select(g => g.AsEnumerable())