1 files changed, 55 insertions, 35 deletions

diff --git a/src/knowledge.rs b/src/knowledge.rs
index a9f39e3..4ade5f5 100644
--- a/src/knowledge.rs
+++ b/src/knowledge.rs
@@ -53,15 +53,24 @@ impl Knowledge {
             .map(|weapon| (weapon.clone(), weapon.single_shot_rounds_to_ready(energy, state.map_size)))
             .collect();
             
-        match self.last_action {
-            Action::PlaceShips(_) => {},
+        let points = match self.last_action {
+            Action::PlaceShips(_) => {
+                vec!()
+            },
             Action::Shoot(Weapon::SingleShot, p) => {
-                new_knowledge.opponent_map.update_from_shot(p, &state);
+                vec!(p)
             },
             Action::Shoot(w, p) => {
+                vec!()
                 //TODO
             }
         };
+
+        let misses = points.iter().filter(|p| state.opponent_map.cells[p.x as usize][p.y as usize].missed).cloned().collect();
+        let hits = points.iter().filter(|p| !state.opponent_map.cells[p.x as usize][p.y as usize].missed).cloned().collect();
+        let sunk_ships = new_knowledge.opponent_map.update_sunk_ships(&state);
+        
+        new_knowledge.opponent_map.update_from_shot(hits, misses, sunk_ships);
         
         new_knowledge
     }
@@ -157,60 +166,69 @@ impl OpponentMapKnowledge {
         }
     }
 
-    fn update_from_shot(&mut self, p: Point, state: &State) {
-        let ref shot_cell = state.opponent_map.cells[p.x as usize][p.y as usize];
-        let sunk_ship = self.ships.iter()
+    fn update_sunk_ships(&mut self, state: &State) -> Vec<Ship> {
+        let sunk_ships = self.ships.iter()
             .filter(|&(_, x)| !x.destroyed)
             .filter(|&(s, _)| state.opponent_map.ships.get(s).map(|x| x.destroyed) == Some(true))
             .map(|(s, _)| s.clone())
-            .next(); //only one ship can be sunk at a time
+            .collect();
 
-       
-        sunk_ship
-            .and_then(|ship| self.ships.get_mut(&ship))
-            .map(|ref mut ship_knowledge| ship_knowledge.destroyed = true);
+        for &ship in &sunk_ships {
+            self.ships.get_mut(&ship).map(|ref mut ship_knowledge| ship_knowledge.destroyed = true);
+        }
 
-        if shot_cell.missed {
-            self.cells[p.x as usize][p.y as usize].missed = true;
-            for knowledge in self.ships.values_mut() {
-                knowledge.possible_placements.retain(|x| !x.touches_point(p));
-            }
+        sunk_ships
+    }
+    
+    fn update_from_shot(&mut self, hit_cells: Vec<Point>, missed_cells: Vec<Point>, sunk_ships: Vec<Ship>) {
+        for &missed in &missed_cells {
+            self.cells[missed.x as usize][missed.y as usize].missed = true;
         }
-        else {
-            self.cells[p.x as usize][p.y as usize].hit = true;
-            self.cells[p.x as usize][p.y as usize].known_ship = sunk_ship;
+        for &hit in &hit_cells {
+            self.cells[hit.x as usize][hit.y as usize].hit = true;
         }
 
-        let cells_copy = self.cells.clone();
-        if sunk_ship.is_some() {
-            for knowledge in self.ships.values_mut() {                
-                knowledge.possible_placements.retain(|x| {
-                    (sunk_ship != Some(x.ship) && !x.touches_point(p)) ||
-                        (sunk_ship == Some(x.ship) && x.touches_point(p) && x.all_are_hits(&cells_copy))
+        self.clear_sunk_ship_impossible_placements(&sunk_ships, &hit_cells);
 
-                });
-            }
+        let mut more_changes = true;
+        while more_changes {
+            more_changes = self.derive_ship_positions() || self.clear_impossible_placements();
         }
-
-        self.derive_ship_positions();
     }
 
-    fn derive_ship_positions(&mut self) {
+    fn derive_ship_positions(&mut self) -> bool {
+        let mut any_changes = false;
         for knowledge in self.ships.values() {
             if knowledge.possible_placements.len() == 1 {
                 let ref true_placement = knowledge.possible_placements[0];
                 for p in true_placement.points_on_ship() {
                     self.cells[p.x as usize][p.y as usize].known_ship = Some(true_placement.ship);
                 }
+                any_changes = true;
             }
         }
-        self.clear_impossible_placements();
+        any_changes
     }
 
-    fn clear_impossible_placements(&mut self) {
+    fn clear_impossible_placements(&mut self) -> bool {
+        let mut any_changes = false;
         let ref cells = self.cells;
         for knowledge in self.ships.values_mut() {
+            let before = knowledge.possible_placements.len();
             knowledge.possible_placements.retain(|x| x.all_could_be_hits(&cells));
+            let after = knowledge.possible_placements.len();
+            if before != after {
+                any_changes = true;
+            }
+        }
+        any_changes
+    }
+
+    fn clear_sunk_ship_impossible_placements(&mut self, sunk_ships: &Vec<Ship>, must_touch_any: &Vec<Point>) {
+        let cells_copy = self.cells.clone();
+        
+        for knowledge in self.ships.values_mut() {
+            knowledge.possible_placements.retain(|x| sunk_ships.contains(&x.ship) && x.touches_any_point(&must_touch_any) && x.all_are_hits(&cells_copy));
         }
     }
 
@@ -292,6 +310,9 @@ impl PossibleShipPlacement {
     pub fn touches_point(&self, p: Point) -> bool {
         p.check_for_ship_collision(self.position, self.direction, self.ship.length())
     }
+    pub fn touches_any_point(&self, ps: &Vec<Point>) -> bool {
+        ps.iter().any(|&p| self.touches_point(p))
+    }
 
     pub fn points_on_ship(&self) -> Vec<Point> {
         (0..self.ship.length() as i32).map(|i| {
@@ -308,10 +329,9 @@ impl PossibleShipPlacement {
     fn all_could_be_hits(&self, cells: &Vec<Vec<KnowledgeCell>>) -> bool {
         self.points_on_ship()
             .iter()
-            .fold(true, |acc, p| {
+            .all(|p| {
                 let ref cell = cells[p.x as usize][p.y as usize];
-                acc && !cell.missed &&
-                    cell.known_ship.map(|ship| ship == self.ship).unwrap_or(true)
+                !cell.missed && cell.known_ship.map(|ship| ship == self.ship).unwrap_or(true)
             })
     }