Ex 11, year 2020 part 1 & part 2

lib.rs

use std::fmt;
use std::str::FromStr;

// [x - 1, y - 1]   [  x  , y - 1]    [x + 1, y - 1]
// [x - 1,   y  ]      (current)      [x + 1,   y  ]
// [x - 1, y + 1]   [  x  , y + 1]    [x + 1, y + 1]
const DIRECTIONS: [(i8, i8); 8] = [
    (-1, -1),
    (0, -1),
    (1, -1),
    (-1, 0),
    (1, 0),
    (-1, 1),
    (0, 1),
    (1, 1),
];

#[derive(Debug)]
enum Tile {
    EmptySeat,
    OccupiedSeat,
    Floor,
}

impl Tile {
    fn from_char(c: char) -> Self {
        match c {
            'L' => Tile::EmptySeat,
            '#' => Tile::OccupiedSeat,
            '.' => Tile::Floor,
            _ => panic!("Unexpected char '{}'", c),
        }
    }

    fn to_char(&self) -> char {
        match self {
            Tile::EmptySeat => 'L',
            Tile::OccupiedSeat => '#',
            Tile::Floor => '.',
        }
    }
}

#[derive(Debug)]
struct Map(Vec<Vec<Tile>>);

impl Map {
    fn is_occupied_at(&self, x: usize, y: usize) -> bool {
        if let Some(line) = self.0.get(y) {
            if let Some(Tile::OccupiedSeat) = line.get(x) {
                return true;
            }
        }
        false
    }

    fn count_occupied(&self) -> u32 {
        let mut cnt = 0;
        for line in self.0.iter() {
            for tile in line.iter() {
                if let Tile::OccupiedSeat = tile {
                    cnt += 1;
                }
            }
        }
        cnt
    }

    fn count_occupied_around(&self, x: usize, y: usize) -> u8 {
        DIRECTIONS
            .iter()
            .map(|(dx, dy)| {
                if (x == 0 && *dx < 0) || (y == 0 && *dy < 0) {
                    0
                } else if self.is_occupied_at((x as i8 + dx) as usize, (y as i8 + dy) as usize) {
                    1
                } else {
                    0
                }
            })
            .sum()
    }

    fn count_occupied_in_sight(&self, x: usize, y: usize) -> u8 {
        DIRECTIONS
            .iter()
            .map(|(dx, dy)| {
                let mut cx = x;
                let mut cy = y;
                loop {
                    if (cx == 0 && *dx < 0) || (cy == 0 && *dy < 0) {
                        return 0;
                    } else {
                        cx = (cx as i8 + dx) as usize;
                        cy = (cy as i8 + dy) as usize;
                        if let Some(row) = self.0.get(cy) {
                            if let Some(tile) = row.get(cx) {
                                match tile {
                                    Tile::EmptySeat => return 0,
                                    Tile::OccupiedSeat => return 1,
                                    _ => {
                                        continue;
                                    }
                                };
                            }
                        }
                        return 0;
                    }
                }
            })
            .sum()
    }

    fn next(&self) -> Map {
        let mut map: Vec<Vec<Tile>> = Vec::with_capacity(self.0.len());
        for (y, line) in self.0.iter().enumerate() {
            let mut new_line: Vec<Tile> = Vec::with_capacity(self.0.get(0).unwrap().len());
            for (x, tile) in line.iter().enumerate() {
                match tile {
                    Tile::EmptySeat => {
                        if self.count_occupied_around(x, y) == 0 {
                            new_line.push(Tile::OccupiedSeat)
                        } else {
                            new_line.push(Tile::EmptySeat)
                        }
                    }
                    Tile::OccupiedSeat => {
                        if self.count_occupied_around(x, y) >= 4 {
                            new_line.push(Tile::EmptySeat)
                        } else {
                            new_line.push(Tile::OccupiedSeat)
                        }
                    }
                    _ => new_line.push(Tile::Floor),
                }
            }
            map.push(new_line);
        }

        Map(map)
    }

    fn next_v2(&self) -> Map {
        let mut map: Vec<Vec<Tile>> = Vec::with_capacity(self.0.len());
        for (y, line) in self.0.iter().enumerate() {
            let mut new_line: Vec<Tile> = Vec::with_capacity(self.0.get(0).unwrap().len());
            for (x, tile) in line.iter().enumerate() {
                match tile {
                    Tile::EmptySeat => {
                        if self.count_occupied_in_sight(x, y) == 0 {
                            new_line.push(Tile::OccupiedSeat)
                        } else {
                            new_line.push(Tile::EmptySeat)
                        }
                    }
                    Tile::OccupiedSeat => {
                        if self.count_occupied_in_sight(x, y) >= 5 {
                            new_line.push(Tile::EmptySeat)
                        } else {
                            new_line.push(Tile::OccupiedSeat)
                        }
                    }
                    _ => new_line.push(Tile::Floor),
                }
            }
            map.push(new_line);
        }

        Map(map)
    }
}

impl FromStr for Map {
    type Err = ();
    fn from_str(s: &str) -> Result<Self, ()> {
        let data = s
            .lines()
            .map(|l| l.chars().map(Tile::from_char).collect::<Vec<Tile>>())
            .collect();

        Ok(Map(data))
    }
}

impl fmt::Display for Map {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "{}",
            self.0
                .iter()
                .map(|line| line.iter().map(Tile::to_char).collect::<String>())
                .collect::<Vec<String>>()
                .join("\n")
        )
    }
}

pub fn part1(input: &str) -> u32 {
    let mut map: Map = input.parse().expect("Cannot parse input");
    let mut prev = map.to_string();
    loop {
        map = map.next();
        let next = map.to_string();
        if prev == next {
            return map.count_occupied();
        }
        prev = next;
    }
}

pub fn part2(input: &str) -> u32 {
    let mut map: Map = input.parse().expect("Cannot parse input");
    let mut prev = map.to_string();
    loop {
        map = map.next_v2();
        let next = map.to_string();
        if prev == next {
            return map.count_occupied();
        }
        prev = next;
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_map_part_1() {
        let input = r#"L.LL.LL.LL
LLLLLLL.LL
L.L.L..L..
LLLL.LL.LL
L.LL.LL.LL
L.LLLLL.LL
..L.L.....
LLLLLLLLLL
L.LLLLLL.L
L.LLLLL.LL"#;

        let exp_step1 = r#"#.##.##.##
#######.##
#.#.#..#..
####.##.##
#.##.##.##
#.#####.##
..#.#.....
##########
#.######.#
#.#####.##"#;

        let exp_step2 = r#"#.LL.L#.##
#LLLLLL.L#
L.L.L..L..
#LLL.LL.L#
#.LL.LL.LL
#.LLLL#.##
..L.L.....
#LLLLLLLL#
#.LLLLLL.L
#.#LLLL.##"#;

        let exp_step3 = r#"#.##.L#.##
#L###LL.L#
L.#.#..#..
#L##.##.L#
#.##.LL.LL
#.###L#.##
..#.#.....
#L######L#
#.LL###L.L
#.#L###.##"#;

        let map: Map = input.parse().expect("Cannot parse input");
        assert_eq!(map.to_string(), input);
        let map = map.next();
        assert_eq!(map.to_string(), exp_step1);
        let map = map.next();
        assert_eq!(map.to_string(), exp_step2);
        let map = map.next();
        assert_eq!(map.to_string(), exp_step3);
    }

    #[test]
    fn test_map_part_2() {
        let input = r#"L.LL.LL.LL
LLLLLLL.LL
L.L.L..L..
LLLL.LL.LL
L.LL.LL.LL
L.LLLLL.LL
..L.L.....
LLLLLLLLLL
L.LLLLLL.L
L.LLLLL.LL"#;

        let exp_step1 = r#"#.##.##.##
#######.##
#.#.#..#..
####.##.##
#.##.##.##
#.#####.##
..#.#.....
##########
#.######.#
#.#####.##"#;

        let exp_step2 = r#"#.LL.LL.L#
#LLLLLL.LL
L.L.L..L..
LLLL.LL.LL
L.LL.LL.LL
L.LLLLL.LL
..L.L.....
LLLLLLLLL#
#.LLLLLL.L
#.LLLLL.L#"#;

        let exp_step3 = r#"#.L#.##.L#
#L#####.LL
L.#.#..#..
##L#.##.##
#.##.#L.##
#.#####.#L
..#.#.....
LLL####LL#
#.L#####.L
#.L####.L#"#;

        let map: Map = input.parse().expect("Cannot parse input");
        assert_eq!(map.to_string(), input);
        let map = map.next_v2();
        assert_eq!(map.to_string(), exp_step1);
        let map = map.next_v2();
        assert_eq!(map.to_string(), exp_step2);
        let map = map.next_v2();
        assert_eq!(map.to_string(), exp_step3);
    }

    #[test]
    fn part_1() {
        let input = include_str!("../input.txt");
        assert_eq!(part1(input), 2261);
    }

    #[test]
    fn part_2() {
        let input = include_str!("../input.txt");
        assert_eq!(part2(input), 2039);
    }
}