Fractal descent of bugs

1 files changed, 119 insertions, 19 deletions

diff --git a/src/bin/day_24.rs b/src/bin/day_24.rs
index 56c4b94..ebc6a1b 100644
--- a/src/bin/day_24.rs
+++ b/src/bin/day_24.rs
@@ -1,5 +1,4 @@
 use rpds::RedBlackTreeSet;
-use std::fmt;
 use std::io;
 use std::io::prelude::*;
 use std::iter;
@@ -12,7 +11,15 @@ use structopt::StructOpt;
 /// Simulates the life and death of Eris bugs
 ///
 /// See https://adventofcode.com/2019/day/24 for details.
-struct Opt {}
+struct Opt {
+    /// How many iterations of the game of life should be run
+    /// If not provided, runs until a state appears twice.
+    #[structopt(short = "n")]
+    depth: Option<usize>,
+    /// Interprets the map as being one part of an infinite fractal map
+    #[structopt(short = "f")]
+    fractal: bool,
+}
 
 fn main() {
     let stdin = io::stdin();
@@ -22,12 +29,16 @@ fn main() {
         .lock()
         .lines()
         .map(|x| exit_on_failed_assertion(x, "Error reading input"))
-        .collect();
+        .collect::<State>()
+        .with_fractal_mode(opt.fractal);
+
+    let final_state = match opt.depth {
+        Some(depth) => initial_state.n_steps_forward(depth),
+        None => initial_state.first_repeated_state(),
+    };
 
-    println!(
-        "{}",
-        initial_state.first_repeated_state().biodiversity_rating()
-    );
+    println!("Bugs: {}", final_state.alive.size());
+    println!("Biodiversity: {}", final_state.biodiversity_rating());
 }
 
 fn exit_on_failed_assertion<A, E: std::error::Error>(data: Result<A, E>, message: &str) -> A {
@@ -43,6 +54,7 @@ fn exit_on_failed_assertion<A, E: std::error::Error>(data: Result<A, E>, message
 #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
 struct State {
     alive: RedBlackTreeSet<Coordinate>,
+    fractal_mode: bool,
 }
 
 impl FromIterator<String> for State {
@@ -63,11 +75,25 @@ impl FromIterator<String> for State {
                         .collect::<Vec<Coordinate>>()
                 })
                 .collect(),
+            fractal_mode: false,
         }
     }
 }
 
 impl State {
+    fn with_fractal_mode(&self, fractal_mode: bool) -> State {
+        State {
+            fractal_mode,
+            ..self.clone()
+        }
+    }
+
+    fn n_steps_forward(&self, n: usize) -> Self {
+        iter::successors(Some(self.clone()), |state| Some(state.next()))
+            .nth(n)
+            .unwrap()
+    }
+
     fn first_repeated_state(&self) -> State {
         iter::successors(
             Some((RedBlackTreeSet::new(), self.clone())),
@@ -84,6 +110,7 @@ impl State {
                 .still_alive_next_round()
                 .chain(self.comes_alive_next_round())
                 .collect(),
+            ..self.clone()
         }
     }
 
@@ -104,7 +131,7 @@ impl State {
     fn comes_alive_next_round<'a>(&'a self) -> impl Iterator<Item = Coordinate> + 'a {
         self.alive
             .iter()
-            .flat_map(|coord| Self::neighbours(*coord))
+            .flat_map(move |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
@@ -112,22 +139,95 @@ impl State {
     }
 
     fn count_alive_neighbours(&self, coordinate: Coordinate) -> usize {
-        Self::neighbours(coordinate)
+        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()
-            .map(|(dx, dy)| Coordinate {
-                x: coordinate.x + dx,
-                y: coordinate.y + dy,
-                depth: coordinate.depth,
-            })
-            .filter(|coord| coord.x >= 0 && coord.x < 5 && coord.y >= 0 && coord.y < 5)
-            .collect()
+    fn neighbours(&self, coordinate: Coordinate) -> Vec<Coordinate> {
+        if self.fractal_mode {
+            [(-1, 0), (1, 0), (0, -1), (0, 1)]
+                .into_iter()
+                .map(|(dx, dy)| Coordinate {
+                    x: coordinate.x + dx,
+                    y: coordinate.y + dy,
+                    depth: coordinate.depth,
+                })
+                .flat_map(move |coord| match (coord.x, coord.y) {
+                    (x, _y) if x < 0 => vec![Coordinate {
+                        x: 1,
+                        y: 2,
+                        depth: coord.depth - 1,
+                    }]
+                    .into_iter(),
+                    (_x, y) if y < 0 => vec![Coordinate {
+                        x: 2,
+                        y: 1,
+                        depth: coord.depth - 1,
+                    }]
+                    .into_iter(),
+                    (x, _y) if x >= 5 => vec![Coordinate {
+                        x: 3,
+                        y: 2,
+                        depth: coord.depth - 1,
+                    }]
+                    .into_iter(),
+                    (_x, y) if y >= 5 => vec![Coordinate {
+                        x: 2,
+                        y: 3,
+                        depth: coord.depth - 1,
+                    }]
+                    .into_iter(),
+                    (2, 2) => match (coordinate.x, coordinate.y) {
+                        (1, 2) => (0..5)
+                            .map(|y| Coordinate {
+                                x: 0,
+                                y,
+                                depth: coord.depth + 1,
+                            })
+                            .collect::<Vec<_>>()
+                            .into_iter(),
+                        (2, 1) => (0..5)
+                            .map(|x| Coordinate {
+                                x,
+                                y: 0,
+                                depth: coord.depth + 1,
+                            })
+                            .collect::<Vec<_>>()
+                            .into_iter(),
+                        (3, 2) => (0..5)
+                            .map(|y| Coordinate {
+                                x: 4,
+                                y,
+                                depth: coord.depth + 1,
+                            })
+                            .collect::<Vec<_>>()
+                            .into_iter(),
+                        (2, 3) => (0..5)
+                            .map(|x| Coordinate {
+                                x,
+                                y: 4,
+                                depth: coord.depth + 1,
+                            })
+                            .collect::<Vec<_>>()
+                            .into_iter(),
+                        (_, _) => vec![].into_iter(),
+                    },
+                    _ => vec![coord.clone()].into_iter(),
+                })
+                .collect()
+        } else {
+            [(-1, 0), (1, 0), (0, -1), (0, 1)]
+                .into_iter()
+                .map(|(dx, dy)| Coordinate {
+                    x: coordinate.x + dx,
+                    y: coordinate.y + dy,
+                    depth: coordinate.depth,
+                })
+                .filter(|coord| coord.x >= 0 && coord.x < 5 && coord.y >= 0 && coord.y < 5)
+                .collect()
+        }
     }
 }