aoc2024/day06/src/main.rs

use std::{
    collections::HashSet,
    fmt::{Display, Write},
};

use winnow::{
    PResult, Parser,
    ascii::{newline, till_line_ending},
    combinator::{alt, opt, repeat, separated},
};

fn main() {
    let input = std::fs::read_to_string("input").unwrap();
    println!("{}", pt1(&input));
    println!("{}", pt2(&input));
}

fn pt1(input: &str) -> usize {
    let mut board = Board::new(input);

    board.run()
}

fn pt2(input: &str) -> usize {
    let mut total = 0;

    let mut board = Board::new(input);

    let guard = board.guard;

    let no = [guard.0, guard.1.next(guard.0)];
    let mut trace = HashSet::new();
    while let Some((loc, _)) = board.step() {
        if !no.contains(&loc) {
            trace.insert(loc);
        }
    }

    let board = Board::new(input);
    for loc in trace {
        let mut board = board.clone();
        let mut steps = HashSet::new();
        board.set(loc, Cell::Obstacle);
        steps.insert(board.guard);
        while let Some(guard) = board.step() {
            if !steps.insert(guard) {
                total += 1;
                break;
            }
        }
    }

    total
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Hash)]
enum Dir {
    #[default]
    N,
    S,
    E,
    W,
}

impl Dir {
    fn next(&self, loc: Loc) -> Loc {
        let Loc(r, c) = loc;
        let (nr, nc) = match self {
            Dir::N => (r.wrapping_sub(1), c),
            Dir::S => (r + 1, c),
            Dir::E => (r, c + 1),
            Dir::W => (r, c.wrapping_sub(1)),
        };
        Loc(nr, nc)
    }

    fn rot(&self) -> Self {
        use Dir::*;
        match self {
            N => E,
            E => S,
            S => W,
            W => N,
        }
    }
}

#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, Hash)]
struct Loc(usize, usize);

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum Cell {
    Obstacle,
    Empty,
    Guard(Dir),
    Visited,
}

impl Display for Cell {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        use Cell::*;
        use Dir::*;
        let c = match self {
            Obstacle => '#',
            Empty => '.',
            Visited => 'X',
            Guard(d) => match d {
                N => '^',
                E => '>',
                S => 'v',
                W => '<',
            },
        };
        f.write_char(c)
    }
}

#[derive(Debug, Default, Clone, PartialEq, Eq)]
struct Board {
    cells: Vec<Vec<Cell>>,
    guard: (Loc, Dir),
}

impl Display for Board {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str(
            &self
                .cells
                .iter()
                //.map(|row| row.iter().map(|c| format!("{c}")).collect::<String>())
                .map(|row| {
                    row.iter().fold(String::new(), |mut out, c| {
                        let _ = write!(out, "{c}");
                        out
                    })
                })
                .collect::<Vec<_>>()
                .join("\n"),
        )
    }
}

impl Board {
    fn new(input: &str) -> Self {
        parse(input).unwrap()
    }

    fn set(&mut self, loc: Loc, cell: Cell) -> bool {
        if let Some(row) = self.cells.get_mut(loc.0) {
            if let Some(c) = row.get_mut(loc.1) {
                *c = cell;
                return true;
            }
        }
        false
    }

    fn step(&mut self) -> Option<(Loc, Dir)> {
        let (loc, dir) = &mut self.guard;
        *self
            .cells
            .get_mut(loc.0)
            .and_then(|r| r.get_mut(loc.1))
            .unwrap() = Cell::Visited;
        let nloc = dir.next(*loc);
        if let Some(next_cell) = self
            .cells
            .get_mut(nloc.0)
            .and_then(|row| row.get_mut(nloc.1))
        {
            match next_cell {
                Cell::Obstacle => {
                    *dir = dir.rot();
                    Some((*loc, *dir))
                }
                Cell::Empty | Cell::Visited => {
                    *next_cell = Cell::Guard(*dir);
                    *loc = nloc;
                    Some((nloc, *dir))
                }
                _ => unreachable!(),
            }
        } else {
            None
        }
    }

    fn run(&mut self) -> usize {
        let mut locs = HashSet::new();
        locs.insert(self.guard.0);
        while let Some(loc) = self.step().map(|l| l.0) {
            locs.insert(loc);
        }
        locs.len()
    }
}

fn parse_cell(input: &mut &str) -> PResult<Cell> {
    alt((parse_guard, parse_floor)).parse_next(input)
}

fn parse_guard(input: &mut &str) -> PResult<Cell> {
    let g = alt(('^', '>', 'v', '<')).parse_next(input)?;
    use Cell::Guard;
    use Dir::*;
    match g {
        '^' => Ok(Guard(N)),
        '>' => Ok(Guard(E)),
        'v' | 'V' => Ok(Guard(S)),
        '<' => Ok(Guard(W)),
        _ => unreachable!(),
    }
}

fn parse_floor(input: &mut &str) -> PResult<Cell> {
    let g = alt(('.', '#')).parse_next(input)?;
    match g {
        '.' => Ok(Cell::Empty),
        '#' => Ok(Cell::Obstacle),
        'X' => Ok(Cell::Visited),
        _ => unreachable!(),
    }
}

fn parse_line(input: &mut &str) -> PResult<Vec<Cell>> {
    let mut line = till_line_ending.parse_next(input)?;
    repeat(1.., parse_cell).parse_next(&mut line)
}

fn parse(input: &str) -> PResult<Board> {
    let mut gdir = Dir::N;
    let mut loc = Loc(0, 0);

    let (cells, _): (Vec<Vec<Cell>>, _) =
        winnow::combinator::seq!(separated(1.., parse_line, newline), opt(newline))
            .parse(input)
            .unwrap();

    for (row, line) in cells.iter().enumerate() {
        for (col, cell) in line.iter().enumerate() {
            if let Cell::Guard(d) = cell {
                gdir = *d;
                loc = Loc(row, col);
            }
        }
    }

    Ok(Board {
        cells,
        guard: (loc, gdir),
    })
}

#[cfg(test)]
mod test {
    use super::*;
    static INPUT: &str = "....#.....
.........#
..........
..#.......
.......#..
..........
.#..^.....
........#.
#.........
......#...";

    #[test]
    fn p1() {
        let v = pt1(INPUT);
        assert_eq!(v, 41)
    }

    #[test]
    fn p2() {
        assert_eq!(6, pt2(INPUT));
    }
}