From 66fa27de0bc548ecb0b3b8febe50ab1ae923cffa Mon Sep 17 00:00:00 2001
From: Antoine Martin <antoine@alarsyo.net>
Date: Sat, 4 Dec 2021 15:38:33 +0100
Subject: [PATCH] 2021: day04: use Vec::retain

---
 aoc2021/src/day04.rs | 39 +++++++++++++++++++--------------------
 1 file changed, 19 insertions(+), 20 deletions(-)

diff --git a/aoc2021/src/day04.rs b/aoc2021/src/day04.rs
index 2a39388..15a1250 100644
--- a/aoc2021/src/day04.rs
+++ b/aoc2021/src/day04.rs
@@ -36,7 +36,8 @@ fn part1(input: &str) -> Result<u64> {
 
     'draw_loop: for draw in draws {
         for grid in &mut grids {
-            if grid.mark(draw) {
+            grid.mark(draw);
+            if grid.is_winning() {
                 wgrid = Some(grid.clone());
                 wdraw = Some(draw);
                 break 'draw_loop;
@@ -74,21 +75,15 @@ fn part2(input: &str) -> Result<u64> {
             .context("no draws available, didn't find last grid")?;
 
         // TODO: replace with drain_filter when stabilized
-        let mut i = 0;
-        while i < grids.len() {
-            let grid = &mut grids[i];
-            if grid.mark(draw) {
-                grids.remove(i);
-            } else {
-                i += 1;
-            }
-        }
+        grids.iter_mut().for_each(|grid| grid.mark(draw));
+        grids.retain(|grid| !grid.is_winning());
     }
 
     let last_grid = &mut grids[0];
 
     for draw in draws {
-        if last_grid.mark(draw) {
+        last_grid.mark(draw);
+        if last_grid.is_winning() {
             return Ok(draw as u64 * last_grid.unmarked_numbers().map(|n| n as u64).sum::<u64>());
         }
     }
@@ -102,23 +97,26 @@ struct Grid {
     grid: [bool; GRID_SIZE],
     rows: [u8; GRID_HEIGHT],
     cols: [u8; GRID_WIDTH],
+    winning: bool,
 }
 
 impl Grid {
-    fn mark(&mut self, draw: u8) -> bool {
-        match self.number_to_pos.get(&draw) {
-            Some(&(x, y)) => {
-                *self.access_grid_mut(x, y) = true;
+    fn mark(&mut self, draw: u8) {
+        if let Some(&(x, y)) = self.number_to_pos.get(&draw) {
+            *self.access_grid_mut(x, y) = true;
 
-                self.rows[y] += 1;
-                self.cols[x] += 1;
+            self.rows[y] += 1;
+            self.cols[x] += 1;
 
-                self.rows[y] as usize == GRID_WIDTH || self.cols[x] as usize == GRID_HEIGHT
-            }
-            None => false,
+            self.winning =
+                self.rows[y] as usize == GRID_WIDTH || self.cols[x] as usize == GRID_HEIGHT;
         }
     }
 
+    fn is_winning(&self) -> bool {
+        self.winning
+    }
+
     fn unmarked_numbers(&self) -> impl Iterator<Item = u8> + '_ {
         self.number_to_pos
             .iter()
@@ -157,6 +155,7 @@ impl std::str::FromStr for Grid {
             grid: [false; GRID_SIZE],
             rows: [0; GRID_HEIGHT],
             cols: [0; GRID_WIDTH],
+            winning: false,
         })
     }
 }