aoc-2024 / src / day6.rs

use std::collections::HashSet;

pub fn day6_part1(input: &str) -> String {
    let map = Map::parse(input);
    map.step_forever().len().to_string()
}

pub fn day6_part2(input: &str) -> String {
    let map = Map::parse(input);
    map.new_maps()
        .map(|map| map.loops())
        .filter(|&map| map)
        .count()
        .to_string()
}

#[derive(Clone, Copy, PartialEq, PartialOrd, Eq, Ord, Hash, Debug)]
struct Point {
    row: usize,
    col: usize,
}

impl Point {
    fn step(self, dir: CardinalDirection) -> Self {
        match dir {
            CardinalDirection::Up => self.up(),
            CardinalDirection::Left => self.left(),
            CardinalDirection::Right => self.right(),
            CardinalDirection::Down => self.down(),
        }
    }
    fn up(self) -> Self {
        Point {
            row: self.row - 1,
            ..self
        }
    }
    fn down(self) -> Self {
        Point {
            row: self.row + 1,
            ..self
        }
    }
    fn left(self) -> Self {
        Point {
            col: self.col - 1,
            ..self
        }
    }
    fn right(self) -> Self {
        Point {
            col: self.col + 1,
            ..self
        }
    }
}

#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Debug)]
enum CardinalDirection {
    Up,
    Left,
    Right,
    Down,
}

impl CardinalDirection {
    fn rotate_90degrees_clockwise(self) -> Self {
        match self {
            Self::Up => Self::Right,
            Self::Left => Self::Up,
            Self::Right => Self::Down,
            Self::Down => Self::Left,
        }
    }
}

impl From<char> for CardinalDirection {
    fn from(c: char) -> Self {
        match c {
            '^' => Self::Up,
            'v' | 'V' => Self::Down,
            '<' => Self::Left,
            '>' => Self::Right,
            _ => unreachable!(),
        }
    }
}

#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Debug)]
struct Guard {
    location: Point,
    direction: CardinalDirection,
}

impl Guard {
    fn step(self) -> Self {
        Self {
            location: self.location.step(self.direction),
            ..self
        }
    }
    fn step_obstacle(self) -> Self {
        let new_dir = self.direction.rotate_90degrees_clockwise();
        Self {
            direction: new_dir,
            ..self
        }
    }
}

struct Map {
    width: usize,
    height: usize,
    guard: Guard,
    obstructions: HashSet<Point>,
    visited: HashSet<Guard>,
    stuck_in_a_loop: bool,
}

impl Map {
    fn new_maps(self) -> impl Iterator<Item = Self> {
        (0..self.width)
            .flat_map(|col| (0..self.height).map(move |row| Point { row, col }))
            .filter(|point| &self.guard.location != point && !self.obstructions.contains(point))
            .collect::<Vec<_>>()
            .into_iter()
            .map(move |point| Self {
                obstructions: {
                    let mut obstructions = self.obstructions.clone();
                    obstructions.insert(point);
                    obstructions
                },
                visited: self.visited.clone(),
                ..self
            })
    }

    fn loops(mut self) -> bool {
        while self.step() && !self.stuck_in_a_loop {}
        self.stuck_in_a_loop
    }

    fn step_forever(mut self) -> HashSet<Point> {
        while self.step() {}
        self.visited.iter().map(|g| g.location).collect()
    }

    /// returns whether a step was taken
    fn step(&mut self) -> bool {
        self.stuck_in_a_loop |= !self.visited.insert(self.guard);
        if self.is_guard_about_to_leave() {
            return false;
        }

        if self.is_guard_facing_obstruction() {
            self.guard = self.guard.step_obstacle();
        } else {
            self.guard = self.guard.step()
        }

        true
    }

    fn is_guard_facing_obstruction(&self) -> bool {
        if self.is_guard_about_to_leave() {
            return false;
        }
        let point = self.guard.step().location;
        self.obstructions.contains(&point)
    }

    fn is_guard_about_to_leave(&self) -> bool {
        match self.guard.direction {
            CardinalDirection::Up => self.guard.location.row == 0,
            CardinalDirection::Down => self.guard.location.row == self.height - 1,
            CardinalDirection::Left => self.guard.location.col == 0,
            CardinalDirection::Right => self.guard.location.col == self.width - 1,
        }
    }

    fn parse(input: &str) -> Self {
        let lines = input.lines().collect::<Vec<_>>();
        let height = lines.len();
        let width = lines[0].chars().count();
        let guard = lines
            .iter()
            .enumerate()
            .find_map(|(row, line)| {
                line.char_indices()
                    .find(|&(_, ch)| ch == '^' || ch == '>' || ch == '<' || ch == 'V' || ch == 'v')
                    .map(|(col, _)| (row, col))
            })
            .unwrap();
        let obstructions = lines
            .iter()
            .enumerate()
            .flat_map(|(row, line)| {
                line.char_indices()
                    .filter(|&(_, ch)| ch == '#')
                    .map(move |(col, _): (usize, char)| Point { row, col })
            })
            .collect();

        Self {
            width,
            height,
            guard: Guard {
                location: Point {
                    row: guard.0,
                    col: guard.1,
                },
                direction: lines[guard.0].chars().nth(guard.1).unwrap().into(),
            },
            obstructions,
            visited: HashSet::new(),
            stuck_in_a_loop: false,
        }
    }
}