Better performance for finding things to shoot at

3 files changed, 44 insertions, 27 deletions

diff --git a/src/game.rs b/src/game.rs
index 3a7cd92..9903250 100644
--- a/src/game.rs
+++ b/src/game.rs
@@ -268,31 +268,48 @@ impl GameBoard {
         self.outcome
     }
 
-    pub fn find_target(&self, center: Point2d<i8>, dir: Direction, weapon_range: u8) -> Option<&Worm> {
-        let diff = dir.as_vec();
-
-        let range = if dir.is_diagonal() {
-            ((weapon_range as f32 + 1.) / 2f32.sqrt()).floor() as i8
-        } else {
-            weapon_range as i8
-        };
-
-        let mut target_worm: Option<&Worm> = None;
-        for distance in 1..=range {
-            let target = center + diff * distance;
-            match self.map.at(target) {
-                Some(false) => {
-                    target_worm = self.players.iter()
-                        .flat_map(|p| p.worms.iter())
-                        .find(|w| w.position == target);
-                    if target_worm.is_some() {
-                        break;
+    pub fn find_targets(&self, center: Point2d<i8>, weapon_range: u8) -> Vec<Direction> {
+        let range = weapon_range as i8;
+        let dir_range = ((weapon_range as f32 + 1.) / 2f32.sqrt()).floor() as i8;
+
+        let mut directions: Vec<Direction> = self.players.iter()
+            .flat_map(|p| p.worms.iter())
+            .filter(|w| w.position != center)
+            .filter_map(|w| {
+                let diff = w.position - center;
+                if diff.x == 0 && diff.y.abs() <= range {
+                    if diff.y > 0 {
+                        Some((Direction::South, diff.y))
+                    } else {
+                        Some((Direction::North, -diff.y))
                     }
-                },
-                _ => break
-            }
-        }
-        target_worm
+                } else if diff.y == 0 && diff.x.abs() <= range {
+                    if diff.x > 0 {
+                        Some((Direction::East, diff.x))
+                    } else {
+                        Some((Direction::West, -diff.x))
+                    }
+                } else if diff.x.abs() == diff.y.abs() && diff.x.abs() <= dir_range {
+                    match (diff.x > 0, diff.y > 0) {
+                        (true, true) => Some((Direction::SouthEast, diff.x)),
+                        (false, true) => Some((Direction::SouthWest, -diff.x)),
+                        (true, false) => Some((Direction::NorthEast, diff.x)),
+                        (false, false) => Some((Direction::NorthWest, -diff.x)),
+                    }
+                } else {
+                    None
+                }
+            })
+            .filter(|(dir, range)| {
+                let diff = dir.as_vec();
+                (1..*range).any(|distance| self.map.at(center + diff * distance) != Some(false))
+            })
+            .map(|(dir, _range)| dir)
+            .collect();
+
+        directions.sort();
+        directions.dedup();
+        directions
     }
 }
 
diff --git a/src/geometry/direction.rs b/src/geometry/direction.rs
index a7d9861..737bd00 100644
--- a/src/geometry/direction.rs
+++ b/src/geometry/direction.rs
@@ -1,7 +1,7 @@
 use std::fmt;
 use crate::geometry::vec::Vec2d;
 
-#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
+#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
 pub enum Direction {
     North,
     NorthEast,
diff --git a/src/strategy.rs b/src/strategy.rs
index d6f92a6..3506e75 100644
--- a/src/strategy.rs
+++ b/src/strategy.rs
@@ -217,8 +217,8 @@ fn valid_moves(state: &GameBoard, player_index: usize) -> Vec<Command> {
             _ => None
         })
         .collect::<Vec<_>>();
-    let mut shoots = Direction::all().iter()
-        .filter(|dir| state.find_target(worm.position, **dir, worm.weapon_range).is_some())
+    let mut shoots = state.find_targets(worm.position, worm.weapon_range)
+        .iter()
         .map(|d| Command::Shoot(*d))
         .collect::<Vec<_>>();
     moves.append(&mut shoots);