Day 24: Arithmetic Logic Unit

use crate::Answer;
use std::{
	collections::{BTreeMap as Map, BTreeSet as Set},
	ops::{Index, IndexMut},
};

type Num = i32;
type Output = String;

#[derive(Debug, Clone, Copy)]
enum Operand {
	Number(Num),
	Reg(Register),
}

impl Operand {
	fn from_str(s: &str) -> Option<Self> {
		use Operand::*;
		Some(match s.parse().ok() {
			Some(val) => Number(val),
			None => Reg(Register::from_str(s)?),
		})
	}
}

#[derive(Debug, Clone, Copy)]
enum MathOp {
	Add,
	Mul,
	Div,
	Mod,
	Eql,
}

impl MathOp {
	fn from_str(s: &str) -> Option<Self> {
		use MathOp::*;
		Some(match s {
			"add" => Add,
			"mul" => Mul,
			"div" => Div,
			"mod" => Mod,
			"eql" => Eql,
			_ => return None,
		})
	}
}

#[derive(Debug, Clone)]
struct MathInstr {
	operation: MathOp,
	register: Register,
	operand: Operand,
}

#[derive(Debug)]
struct InstrBlock {
	in_reg: Register,
	instrs: Vec<MathInstr>,
}

#[derive(Debug, Clone, Copy)]
enum Register {
	W = 0,
	X,
	Y,
	Z,
}

impl Register {
	fn from_str(s: &str) -> Option<Self> {
		use Register::*;
		Some(match s {
			"w" => W,
			"x" => X,
			"y" => Y,
			"z" => Z,
			_ => return None,
		})
	}
}

#[derive(Debug)]
struct Alu {
	registers: [Num; 4],
}

impl Index<Register> for Alu {
	type Output = Num;
	fn index(&self, register: Register) -> &Self::Output {
		&self.registers[register as usize]
	}
}

impl IndexMut<Register> for Alu {
	fn index_mut(&mut self, register: Register) -> &mut Self::Output {
		&mut self.registers[register as usize]
	}
}

impl Alu {
	fn new() -> Self {
		Self { registers: [0; 4] }
	}

	fn run_block(&mut self, block: &InstrBlock, input: Num) {
		use MathOp::*;
		use Operand::*;

		self[block.in_reg] = input;

		for &MathInstr {
			operation,
			register,
			operand,
		} in &block.instrs
		{
			let value = match operand {
				Number(n) => n,
				Reg(register) => self[register],
			};

			let r = &mut self[register];
			match operation {
				Add => *r += value,
				Mul => *r *= value,
				Div => *r /= value,
				Mod => *r %= value,
				Eql => *r = if r == &value { 1 } else { 0 },
			}
		}
	}

	fn from_running_block_on(
		block: &InstrBlock,
		input: Num,
		setup: impl FnOnce(&mut Self),
	) -> Self {
		let mut alu = Alu::new();
		setup(&mut alu);
		alu.run_block(block, input);
		alu
	}
}

fn read_input(s: &str) -> Option<Vec<InstrBlock>> {
	let mut blocks = vec![];
	let mut curr_in_reg = None;
	let mut curr_instrs = vec![];

	// Dummy input line at the end that tells the last block it's done
	for line in s.lines().chain(std::iter::once("inp x")) {
		let mut splat = line.split_ascii_whitespace();
		let instr_str = splat.next()?;
		if instr_str == "inp" {
			if let Some(r) = curr_in_reg {
				blocks.push(InstrBlock {
					in_reg: r,
					instrs: curr_instrs.clone(),
				});
			}

			curr_in_reg = Some(Register::from_str(splat.next()?)?);
			curr_instrs.clear();
		} else {
			let operation = MathOp::from_str(instr_str)?;
			let register = Register::from_str(splat.next()?)?;
			let operand = Operand::from_str(splat.next()?)?;
			curr_instrs.push(MathInstr {
				operation,
				register,
				operand,
			});
		}
	}

	Some(blocks)
}

fn ans_for_input(input: &str) -> Answer<Output, Output> {
	let blocks = read_input(input).unwrap();
	let valid_zs = get_valid_zs(&blocks);
	(24, (pt1(&blocks, &valid_zs), pt2(&blocks, &valid_zs))).into()
}

pub fn ans() -> Answer<Output, Output> {
	ans_for_input(include_str!("input.txt"))
}

fn pt1(blocks: impl AsRef<[InstrBlock]>, valid_zs: impl AsRef<[Set<Num>]>) -> Output {
	find_digits(
		blocks,
		valid_zs,
		9,
		|digit| (1..=digit).rev(),
		|digit| digit - 1,
		|digit| digit > 1,
	)
}

fn pt2(blocks: impl AsRef<[InstrBlock]>, valid_zs: impl AsRef<[Set<Num>]>) -> Output {
	find_digits(
		blocks,
		valid_zs,
		1,
		|digit| digit..=9,
		|digit| digit + 1,
		|digit| digit < 9,
	)
}