aoc2020/src/bin/day11.rs

const INPUT: &str = include_str!("../../data/day11");

#[derive(Clone, Copy, PartialEq, Eq, Debug)]
enum Cell {
	Floor,
	Seat { occupied: bool },
}

impl Cell {
	fn parse(b: u8) -> Cell {
		match b {
			b'.' => Self::Floor,
			b'L' => Self::Seat { occupied: false },
			b'#' => Self::Seat { occupied: true }, // shouldn't happen in input...
			_ => panic!("Invalid cell: {}", b),
		}
	}

	fn as_char(self) -> char {
		match self {
			Self::Floor => '.',
			Self::Seat { occupied: false } => 'L',
			Self::Seat { occupied: true } => '#',
		}
	}

	fn occupied(self) -> bool {
		match self {
			Self::Floor => false,
			Self::Seat { occupied: false } => false,
			Self::Seat { occupied: true } => true,
		}
	}
}

#[derive(Clone, PartialEq, Eq, Debug)]
struct Area {
	width: usize,
	height: usize,
	cells: Vec<Cell>,
}

impl Area {
	fn parse(data: &str) -> Self {
		let input_width = data.lines().next().unwrap().len();
		let width = input_width + 2;
		let mut height = 0;
		let mut cells = Vec::new();
		// use stub cells to surround everything to avoid having to check indices for overflow
		const STUB_CELL: Cell = Cell::Seat { occupied: false };
		// stub first row with floor cells
		height += 1;
		cells.resize(cells.len() + width, STUB_CELL);
		for line in data.lines() {
			assert_eq!(input_width, line.len());
			height += 1;
			cells.push(STUB_CELL);
			cells.extend(line.bytes().map(Cell::parse));
			cells.push(STUB_CELL);
		}
		height += 1;
		cells.resize(cells.len() + width, STUB_CELL);
		assert_eq!(width * height, cells.len());
		Self { width, height, cells }
	}

	fn occupied(&self) -> usize {
		self.cells.iter().filter(|c| c.occupied()).count()
	}

	fn step(&self, target: &mut Self) {
		assert_eq!(self.width, target.width);
		assert_eq!(self.height, target.height);
		let w = self.width as isize;
		let adjacent_offsets = [
			-w - 1, -w, -w + 1,
			-1, /* 0, */ 1,
			w - 1, w, w + 1,
		];
		for y in 1..self.height-1 {
			let row_offset = y * self.width;
			for x in 1..self.width-1 {
				let cell_offset = row_offset + x;
				let co = cell_offset as isize;
				let c = match self.cells[cell_offset] {
					Cell::Floor => Cell::Floor,
					Cell::Seat { occupied: false } => {
						if adjacent_offsets.iter().all(|&o| !self.cells[(co + o) as usize].occupied()) {
							Cell::Seat { occupied: true }
						} else {
							Cell::Seat { occupied: false }
						}
					},
					Cell::Seat { occupied: true } => {
						if adjacent_offsets.iter().filter(|&o| self.cells[(co + o) as usize].occupied()).count() < 4 {
							Cell::Seat { occupied: true }
						} else {
							Cell::Seat { occupied: false }
						}
					},
				};
				target.cells[cell_offset] = c;
			}
		}
	}

	fn search_seat(&self, (mut y, mut x): (usize, usize), (diry, dirx): (isize, isize)) -> bool {
		loop {
			y = (y as isize + diry) as usize;
			x = (x as isize + dirx) as usize;
			match self[(y, x)] {
				Cell::Floor => (),
				Cell::Seat { occupied } => return occupied,
			}
		}
	}

	fn step2(&self, target: &mut Self) {
		assert_eq!(self.width, target.width);
		assert_eq!(self.height, target.height);
		const LOOK_DIRECTIONS: [(isize, isize); 8] = [
			(-1, -1), (-1, 0), (-1, 1),
			(0, -1), /* (0, 0), */ (0, 1),
			(1, -1), (1, 0), (1, 1),
		];
		for y in 1..self.height-1 {
			let row_offset = y * self.width;
			for x in 1..self.width-1 {
				let cell_offset = row_offset + x;
				let c = match self.cells[cell_offset] {
					Cell::Floor => Cell::Floor,
					Cell::Seat { occupied: false } => {
						if LOOK_DIRECTIONS.iter().all(|&ld| !self.search_seat((y, x), ld)) {
							Cell::Seat { occupied: true }
						} else {
							Cell::Seat { occupied: false }
						}
					},
					Cell::Seat { occupied: true } => {
						if LOOK_DIRECTIONS.iter().filter(|&&ld| self.search_seat((y, x), ld)).take(5).count() < 5 {
							Cell::Seat { occupied: true }
						} else {
							Cell::Seat { occupied: false }
						}
					},
				};
				target.cells[cell_offset] = c;
			}
		}
	}
}

impl std::ops::Index<(usize, usize)> for Area {
	type Output = Cell;

	fn index(&self, (y, x): (usize, usize)) -> &Self::Output {
		&self.cells[y * self.width + x]
	}
}

impl std::fmt::Display for Area {
	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
		use std::fmt::Write;
		for y in 1..self.height-1 {
			let row_offset = y * self.width;
			f.write_char(' ')?;
			for x in 1..self.width-1 {
				f.write_char(self.cells[row_offset+x].as_char())?;
			}
			f.write_str("\n")?;
		}
		Ok(())
	}
}

fn main() {
	let input_area = Box::new(Area::parse(INPUT));
	{
		let mut area = input_area.clone();
		let mut next_area = area.clone();
		loop {
			area.step(&mut next_area);
			if area == next_area { break; }
			std::mem::swap(&mut area, &mut next_area);
		}
		// print!("{}", area);
		println!("Occupied part 1: {}", area.occupied());
	}
	{
		let mut area = input_area.clone();
		let mut next_area = area.clone();
		loop {
			area.step2(&mut next_area);
			if area == next_area { break; }
			std::mem::swap(&mut area, &mut next_area);
		}
		print!("{}", area);
		println!("Occupied part 2: {}", area.occupied());
	}
}