Pruned moves better

1 files changed, 65 insertions, 12 deletions

diff --git a/src/game.rs b/src/game.rs
index 11108e3..32efef8 100644
--- a/src/game.rs
+++ b/src/game.rs
@@ -614,14 +614,24 @@ impl GameBoard {
             .collect()
     }
 
+    // TODO: These two dominate in terms of time taken. Optimize.
     fn pruned_valid_bomb_commands(&self, player_index: usize) -> Vec<Command> {
         self.selects_iter(player_index)
             .filter(|(_, worm)| worm.bombs > 0)
             .flat_map(|(select, worm)| {
                 let mut result = Vec::with_capacity((BOMB_RANGE * 2 + 1).pow(2) as usize - 12);
+                let own_worm_positions: ArrayVec<[Point2d<i8>; 3]> = self.players[player_index]
+                    .worms
+                    .iter()
+                    .map(|w| w.position)
+                    .collect();
+                let opponent_worm_positions: ArrayVec<[Point2d<i8>; 3]> = self.players
+                    [GameBoard::opponent(player_index)]
+                .worms
+                .iter()
+                .map(|w| w.position)
+                .collect();
 
-                // TODO: Don't push the ones where you're hurt by the bomb
-                // TODO: Only push the ones that hurt an opponent worm
                 for y in worm.position.y - BOMB_RANGE..=worm.position.y + BOMB_RANGE {
                     for x in worm.position.x - BOMB_RANGE..=worm.position.x + BOMB_RANGE {
                         let target = Point2d::new(x, y);
@@ -629,10 +639,27 @@ impl GameBoard {
                             && (worm.position - target).magnitude_squared()
                                 < (BOMB_RANGE + 1).pow(2)
                         {
-                            result.push(Command {
-                                worm: select,
-                                action: Action::Bomb(target),
-                            });
+                            let own_affected_worms = BOMB_DAMAGES
+                                .iter()
+                                .map(|(damage_offset, _)| target + *damage_offset)
+                                .filter(|explosion_position| {
+                                    own_worm_positions.contains(explosion_position)
+                                })
+                                .count();
+                            let opponent_affected_worms = BOMB_DAMAGES
+                                .iter()
+                                .map(|(damage_offset, _)| target + *damage_offset)
+                                .filter(|explosion_position| {
+                                    opponent_worm_positions.contains(explosion_position)
+                                })
+                                .count();
+
+                            if own_affected_worms == 0 && opponent_affected_worms > 0 {
+                                result.push(Command {
+                                    worm: select,
+                                    action: Action::Bomb(target),
+                                });
+                            }
                         }
                     }
                 }
@@ -647,9 +674,18 @@ impl GameBoard {
             .filter(|(_, worm)| worm.snowballs > 0)
             .flat_map(|(select, worm)| {
                 let mut result = Vec::with_capacity((SNOWBALL_RANGE * 2 + 1).pow(2) as usize - 12);
+                let own_worm_positions: ArrayVec<[Point2d<i8>; 3]> = self.players[player_index]
+                    .worms
+                    .iter()
+                    .map(|w| w.position)
+                    .collect();
+                let opponent_worm_positions: ArrayVec<[Point2d<i8>; 3]> = self.players
+                    [GameBoard::opponent(player_index)]
+                .worms
+                .iter()
+                .map(|w| w.position)
+                .collect();
 
-                // TODO: Don't push the ones where you're freezing yourself
-                // TODO: Only push the ones that freeze an opponent worm
                 for y in worm.position.y - SNOWBALL_RANGE..=worm.position.y + SNOWBALL_RANGE {
                     for x in worm.position.x - SNOWBALL_RANGE..=worm.position.x + SNOWBALL_RANGE {
                         let target = Point2d::new(x, y);
@@ -657,10 +693,27 @@ impl GameBoard {
                             && (worm.position - target).magnitude_squared()
                                 < (SNOWBALL_RANGE + 1).pow(2)
                         {
-                            result.push(Command {
-                                worm: select,
-                                action: Action::Snowball(target),
-                            });
+                            let own_affected_worms = SNOWBALL_FREEZES
+                                .iter()
+                                .map(|offset| target + *offset)
+                                .filter(|explosion_position| {
+                                    own_worm_positions.contains(explosion_position)
+                                })
+                                .count();
+                            let opponent_affected_worms = SNOWBALL_FREEZES
+                                .iter()
+                                .map(|offset| target + *offset)
+                                .filter(|explosion_position| {
+                                    opponent_worm_positions.contains(explosion_position)
+                                })
+                                .count();
+
+                            if own_affected_worms == 0 && opponent_affected_worms > 0 {
+                                result.push(Command {
+                                    worm: select,
+                                    action: Action::Snowball(target),
+                                });
+                            }
                         }
                     }
                 }