Refactor buggy game of life to use Set instead of Vector

This sets me up for the infinite play space of part 2

1 files changed, 58 insertions, 55 deletions

diff --git a/src/bin/day_24.rs b/src/bin/day_24.rs
index ef10ae8..56c4b94 100644
--- a/src/bin/day_24.rs
+++ b/src/bin/day_24.rs
@@ -1,4 +1,3 @@
-use rpds::vector::Vector;
 use rpds::RedBlackTreeSet;
 use std::fmt;
 use std::io;
@@ -43,45 +42,36 @@ fn exit_on_failed_assertion<A, E: std::error::Error>(data: Result<A, E>, message
 
 #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
 struct State {
-    grid: Vector<Vector<bool>>,
+    alive: RedBlackTreeSet<Coordinate>,
 }
 
 impl FromIterator<String> for State {
     fn from_iter<T: IntoIterator<Item = String>>(iter: T) -> Self {
         State {
-            grid: iter
+            alive: iter
                 .into_iter()
-                .map(|line| line.chars().map(|c| c == '#').collect::<Vector<_>>())
+                .enumerate()
+                .flat_map(move |(y, line)| {
+                    line.chars()
+                        .enumerate()
+                        .filter(move |(_x, c)| *c == '#')
+                        .map(move |(x, _c)| Coordinate {
+                            x: x as isize,
+                            y: y as isize,
+                            depth: 0,
+                        })
+                        .collect::<Vec<Coordinate>>()
+                })
                 .collect(),
         }
     }
 }
 
-impl fmt::Display for State {
-    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        self.grid
-            .iter()
-            .map(|row| {
-                writeln!(
-                    f,
-                    "{}",
-                    row.iter()
-                        .map(|alive| if *alive { '#' } else { '.' })
-                        .collect::<String>()
-                )
-            })
-            .collect()
-    }
-}
-
 impl State {
     fn first_repeated_state(&self) -> State {
         iter::successors(
             Some((RedBlackTreeSet::new(), self.clone())),
-            |(seen, state)| {
-                eprintln!("{}", state);
-                Some((seen.insert(state.clone()), state.next()))
-            },
+            |(seen, state)| Some((seen.insert(state.clone()), state.next())),
         )
         .find(|(seen, state)| seen.contains(state))
         .unwrap()
@@ -90,47 +80,60 @@ impl State {
 
     fn next(&self) -> State {
         State {
-            grid: self
-                .grid
-                .iter()
-                .enumerate()
-                .map(|(y, row)| {
-                    row.iter()
-                        .enumerate()
-                        .map(|(x, alive)| {
-                            if *alive {
-                                self.count_alive_neighbours(x, y) == 1
-                            } else {
-                                self.count_alive_neighbours(x, y) == 1
-                                    || self.count_alive_neighbours(x, y) == 2
-                            }
-                        })
-                        .collect()
-                })
+            alive: self
+                .still_alive_next_round()
+                .chain(self.comes_alive_next_round())
                 .collect(),
         }
     }
 
     fn biodiversity_rating(&self) -> usize {
-        self.grid
+        self.alive
             .iter()
-            .flat_map(|row| row.iter())
-            .enumerate()
-            .map(|(i, state)| if *state { 2_usize.pow(i as u32) } else { 0 })
+            .map(|coord| 2_usize.pow((coord.y * 5 + coord.x) as u32))
             .sum()
     }
 
-    fn count_alive_neighbours(&self, x: usize, y: usize) -> usize {
+    fn still_alive_next_round<'a>(&'a self) -> impl Iterator<Item = Coordinate> + 'a {
+        self.alive
+            .iter()
+            .filter(move |coord| self.count_alive_neighbours(**coord) == 1)
+            .cloned()
+    }
+
+    fn comes_alive_next_round<'a>(&'a self) -> impl Iterator<Item = Coordinate> + 'a {
+        self.alive
+            .iter()
+            .flat_map(|coord| Self::neighbours(*coord))
+            .filter(move |coord| !self.alive.contains(coord))
+            .filter(move |coord| {
+                self.count_alive_neighbours(*coord) == 1 || self.count_alive_neighbours(*coord) == 2
+            })
+    }
+
+    fn count_alive_neighbours(&self, coordinate: Coordinate) -> usize {
+        Self::neighbours(coordinate)
+            .into_iter()
+            .filter(|coord| self.alive.contains(coord))
+            .count()
+    }
+
+    fn neighbours(coordinate: Coordinate) -> Vec<Coordinate> {
         [(-1, 0), (1, 0), (0, -1), (0, 1)]
             .into_iter()
-            .filter_map(|(dx, dy)| {
-                if (x > 0 || *dx >= 0) && (y > 0 || *dy >= 0) {
-                    Some(((x as i32 + dx) as usize, (y as i32 + dy) as usize))
-                } else {
-                    None
-                }
+            .map(|(dx, dy)| Coordinate {
+                x: coordinate.x + dx,
+                y: coordinate.y + dy,
+                depth: coordinate.depth,
             })
-            .filter(|(x, y)| self.grid.get(*y).and_then(|row| row.get(*x)).cloned() == Some(true))
-            .count()
+            .filter(|coord| coord.x >= 0 && coord.x < 5 && coord.y >= 0 && coord.y < 5)
+            .collect()
     }
 }
+
+#[derive(Debug, Clone, Copy, PartialOrd, Ord, PartialEq, Eq)]
+struct Coordinate {
+    x: isize,
+    y: isize,
+    depth: isize,
+}