day16: back to elegant recursion now that we track beams

problem 1 solution was busting the stack, so used an iterative approach,
but now that we track lit beams we can bail out early enough not to bust
the stack. performance is the same but it's much nicer.

16/src/main.rs

@@ -1,4 +1,3 @@
-use std::collections::HashSet;
 use std::fmt::Display;
 use std::fs::File;
 use std::io::{BufRead, BufReader, Lines};
@@ -28,12 +27,72 @@ fn main() {
 
 // PARSE
 
+enum Interaction {
+    One(FromDirection),
+    Two((FromDirection, FromDirection)),
+}
+
 #[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
 enum FromDirection {
-    Left,
-    Above,
-    Right,
-    Below,
+    Left = 0,
+    Above = 1,
+    Right = 2,
+    Below = 3,
+}
+
+impl FromDirection {
+    fn mask(&self) -> u8 {
+        1 << *self as u8
+    }
+    // return the new pos for a ray that will be from the direction.
+    // a ray 'from' left 'goes' right
+    fn goes_pos(&self, pos: (i64, i64)) -> (i64, i64) {
+        match self {
+            Self::Left => (pos.0 + 1, pos.1),
+            Self::Right => (pos.0 - 1, pos.1),
+            Self::Above => (pos.0, pos.1 + 1),
+            Self::Below => (pos.0, pos.1 - 1),
+        }
+    }
+    fn reflect_ne(&self) -> Self {
+        match self {
+            Self::Left => Self::Below,
+            Self::Right => Self::Above,
+            Self::Above => Self::Right,
+            Self::Below => Self::Left,
+        }
+    }
+    fn opposite(&self) -> FromDirection {
+        match self {
+            Self::Left => Self::Right,
+            Self::Right => Self::Left,
+            Self::Above => Self::Below,
+            Self::Below => Self::Above,
+        }
+    }
+    fn reflect_se(&self) -> FromDirection {
+        self.reflect_ne().opposite()
+    }
+    fn interact(&self, tile: char) -> Interaction {
+        match tile {
+            '.' => Interaction::One(*self),
+            '/' => Interaction::One(self.reflect_ne()),
+            '\\' => Interaction::One(self.reflect_se()),
+            '|' => match self {
+                FromDirection::Above | FromDirection::Below => Interaction::One(*self),
+                FromDirection::Left | FromDirection::Right => {
+                    Interaction::Two((FromDirection::Above, FromDirection::Below))
+                }
+            },
+            '-' => match self {
+                FromDirection::Left | FromDirection::Right => Interaction::One(*self),
+                FromDirection::Above | FromDirection::Below => {
+                    Interaction::Two((FromDirection::Left, FromDirection::Right))
+                }
+            },
+            c => unimplemented!("invalid tile {}", c),
+        }
+    }
 }
 
 impl Display for FromDirection {
@@ -52,113 +111,77 @@ struct Contraption {
 }
 
 struct VisitState {
-    energized: Vec<Vec<bool>>,
-    visited_rays: HashSet<((i64, i64), FromDirection)>,
+    visited_from: Vec<Vec<u8>>,
 }
 
 impl VisitState {
-    fn score(&self) -> u64 {
-        self.energized.iter().flatten().filter(|c| **c).count() as u64
+    fn visit(&mut self, pos: (i64, i64), dir: FromDirection) -> bool {
+        let pos_state = &mut self.visited_from[pos.1 as usize][pos.0 as usize];
+        if *pos_state & dir.mask() > 0 {
+            false
+        } else {
+            *pos_state |= dir.mask();
+            true
         }
+    }
+    fn score(&self) -> u64 {
+        self.visited_from
+            .iter()
+            .flatten()
+            .filter(|c| **c != 0)
+            .count() as u64
+    }
+    #[allow(dead_code)]
     fn dump(&self) {
         println!("Score {}:", self.score());
-        for line in &self.energized {
+        for line in &self.visited_from {
             println!(
                 "  {}",
-                String::from_iter(line.iter().map(|b| if *b { '#' } else { '.' }))
+                String::from_iter(line.iter().map(|b| if *b == 0 { '.' } else { '#' }))
             );
         }
     }
 }
 
-
 impl Contraption {
-    fn cast_ray<'a>(
+    fn height(&self) -> i64 {
+        self.tiles.len() as i64
+    }
+    fn width(&self) -> i64 {
+        self.tiles[0].len() as i64
+    }
+    fn cast_ray<'a>(&'a mut self, pos: (i64, i64), dir: FromDirection) -> VisitState {
+        let mut state = self.empty_state();
+        self.cast_ray_inner(&mut state, pos, dir);
+        state
+    }
+    fn cast_ray_inner<'a>(
         &'a mut self,
         state: &'a mut VisitState,
         pos: (i64, i64),
         dir: FromDirection,
     ) {
-        let mut new_rays = self.cast_ray_inner(state, pos, dir);
-
-        loop {
-            new_rays = new_rays
-                .iter()
-                .flat_map(|(pos, dir)| self.cast_ray_inner(state, *pos, *dir))
-                .collect();
-            if new_rays.len() == 0 {
-                break;
+        if pos.0 >= 0
+            && pos.1 >= 0
+            && pos.0 < self.width()
+            && pos.1 < self.height()
+            && state.visit(pos, dir)
+        {
+            match dir.interact(self.tiles[pos.1 as usize][pos.0 as usize]) {
+                Interaction::One(dir) => self.cast_ray_inner(state, dir.goes_pos(pos), dir),
+                Interaction::Two((dir1, dir2)) => {
+                    self.cast_ray_inner(state, dir1.goes_pos(pos), dir1);
+                    self.cast_ray_inner(state, dir2.goes_pos(pos), dir2);
                 }
+            };
         }
     }
-    fn cast_ray_inner<'a>(
-        &'a mut self,
-        state: &'a mut VisitState,
-        mut pos: (i64, i64),
-        mut dir: FromDirection,
-    ) -> Vec<((i64, i64), FromDirection)> {
-        let width = state.energized[0].len();
-        let height = state.energized.len();
-        let mut new_rays = Vec::new();
-        while pos.0 >= 0 && pos.1 >= 0 && pos.0 < width as i64 && pos.1 < height as i64 {
-            // visit pos
-            state.energized[pos.1 as usize][pos.0 as usize] = true;
-            if !state.visited_rays.insert((pos, dir)) {
-                break;
-            }
-            (pos, dir) = match self.tiles[pos.1 as usize][pos.0 as usize] {
-                '.' => match dir {
-                    FromDirection::Left => ((pos.0 + 1, pos.1), dir),
-                    FromDirection::Right => ((pos.0 - 1, pos.1), dir),
-                    FromDirection::Above => ((pos.0, pos.1 + 1), dir),
-                    FromDirection::Below => ((pos.0, pos.1 - 1), dir),
-                },
-                '/' => match dir {
-                    // from left, go up, from below
-                    FromDirection::Left => ((pos.0, pos.1 - 1), FromDirection::Below),
-                    // from up, go left, from the right
-                    FromDirection::Above => ((pos.0 - 1, pos.1), FromDirection::Right),
-                    // from right, go down, from above
-                    FromDirection::Right => ((pos.0, pos.1 + 1), FromDirection::Above),
-                    // from below, go right, from left
-                    FromDirection::Below => ((pos.0 + 1, pos.1), FromDirection::Left),
-                },
-                '\\' => match dir {
-                    FromDirection::Left => ((pos.0, pos.1 + 1), FromDirection::Above),
-                    FromDirection::Above => ((pos.0 + 1, pos.1), FromDirection::Left),
-                    FromDirection::Right => ((pos.0, pos.1 - 1), FromDirection::Below),
-                    FromDirection::Below => ((pos.0 - 1, pos.1), FromDirection::Right),
-                },
-                '-' => match dir {
-                    FromDirection::Left => ((pos.0 + 1, pos.1), dir),
-                    FromDirection::Right => ((pos.0 - 1, pos.1), dir),
-                    FromDirection::Above | FromDirection::Below => {
-                        new_rays.push(((pos.0 + 1, pos.1), FromDirection::Left));
-                        ((pos.0 - 1, pos.1), FromDirection::Right)
-                    }
-                },
-                '|' => match dir {
-                    FromDirection::Above => ((pos.0, pos.1 + 1), dir),
-                    FromDirection::Below => ((pos.0, pos.1 - 1), dir),
-                    FromDirection::Left | FromDirection::Right => {
-                        new_rays.push(((pos.0, pos.1 + 1), FromDirection::Above));
-                        ((pos.0, pos.1 - 1), FromDirection::Below)
-                    }
-                },
-                c => unimplemented!("invalid character {}", c),
-            }
-        }
-        new_rays
-    }
     fn empty_state(&self) -> VisitState {
-        let mut energized = Vec::new();
-        for _ in 0..self.tiles.len() {
-            energized.push(Vec::from_iter(repeat(false).take(self.tiles[0].len())));
-        }
-        VisitState {
-            energized,
-            visited_rays: HashSet::new(),
+        let mut visited_from = Vec::new();
+        for _ in 0..self.height() {
+            visited_from.push(Vec::from_iter(repeat(0).take(self.width() as usize)));
         }
+        VisitState { visited_from }
     }
 }
 
@@ -176,48 +199,39 @@ impl<T: BufRead> From<Lines<T>> for Contraption {
 
 fn problem1<T: BufRead>(input: Lines<T>) -> u64 {
     let mut contraption = Contraption::from(input);
-    let mut state = contraption.empty_state();
 
-    contraption.cast_ray(&mut state, (0,0), FromDirection::Left);
-
-    state.energized.iter().flatten().filter(|c| **c).count() as u64
+    contraption.cast_ray((0, 0), FromDirection::Left).score()
 }
 
 // PROBLEM 2 solution
 fn problem2<T: BufRead>(input: Lines<T>) -> u64 {
     let mut contraption = Contraption::from(input);
-    let mut max_tiles = 0u64;
 
-    for y in 0..contraption.tiles.len() as i64 {
-        let mut left_state = contraption.empty_state();
-        contraption.cast_ray(&mut left_state, (0, y), FromDirection::Left);
-        let mut right_state = contraption.empty_state();
-        contraption.cast_ray(
-            &mut right_state,
-            (contraption.tiles[0].len() as i64 - 1, y),
-            FromDirection::Right,
-        );
-        max_tiles = std::cmp::max(
+    let rows_max = (0..contraption.height()).fold(0, |max_tiles, y| {
+        std::cmp::max(
             max_tiles,
-            std::cmp::max(left_state.score(), right_state.score()),
-        );
-    }
-    for x in 0..contraption.tiles[0].len() as i64 {
-        let mut top_state = contraption.empty_state();
-        contraption.cast_ray(&mut top_state, (x, 0), FromDirection::Above);
-        let mut bottom_state = contraption.empty_state();
-        contraption.cast_ray(
-            &mut bottom_state,
-            (x, contraption.tiles.len() as i64 - 1),
-            FromDirection::Below,
-        );
-        max_tiles = std::cmp::max(
-            max_tiles,
-            std::cmp::max(top_state.score(), bottom_state.score()),
-        );
-    }
+            std::cmp::max(
+                contraption.cast_ray((0, y), FromDirection::Left).score(),
+                contraption
+                    .cast_ray((contraption.width() - 1, y), FromDirection::Right)
+                    .score(),
+            ),
+        )
+    });
 
-    max_tiles
+    let cols_max = (0..contraption.width()).fold(0, |max_tiles, x| {
+        std::cmp::max(
+            max_tiles,
+            std::cmp::max(
+                contraption.cast_ray((x, 0), FromDirection::Above).score(),
+                contraption
+                    .cast_ray((x, contraption.height() - 1), FromDirection::Below)
+                    .score(),
+            ),
+        )
+    });
+
+    std::cmp::max(rows_max, cols_max)
 }
 
 #[cfg(test)]