aoc2019/shared/src/intcode.rs

const OP_ADD: CELL = 1;
const OP_MUL: CELL = 2;
const OP_INPUT: CELL = 3;
const OP_OUTPUT: CELL = 4;
const OP_JUMP_IF_TRUE: CELL = 5;
const OP_JUMP_IF_FALSE: CELL = 6;
const OP_LESS_THAN: CELL = 7;
const OP_EQUAL: CELL = 8;
const OP_HALT: CELL = 99;

pub type CELL = i64;
type Mode = u8;

#[derive(PartialEq, Eq, Clone, Debug)]
pub enum SimulateStep {
	Finished,
	WaitForInput,
	Output(CELL),
}

impl SimulateStep {
	pub fn expect_output(self) -> CELL {
		if let Self::Output(v) = self {
			return v;
		}
		panic!("Expected Output, got {:?}", self);
	}

	pub fn expect_finished(self) {
		if let Self::Finished = self { return; }
		panic!("Expected Finished, got {:?}", self);
	}

	pub fn expect_wait_for_input(self) {
		if let Self::WaitForInput = self { return; }
		panic!("Expected WaitForInput, got {:?}", self);
	}
}

#[derive(PartialEq, Eq, Clone, Debug)]
pub struct IntCode {
	pub ip: usize,
	pub next_input: Option<CELL>,
	pub data: Vec<CELL>,
}

impl IntCode {
	#[inline(always)]
	fn fetch(&mut self, mode: Mode) -> CELL {
		let imm = self.data[self.ip];
		self.ip += 1;
		match mode {
			1 => imm,
			_ => self.data[imm as usize],
		}
	}
	#[inline(always)]
	fn store(&mut self, mode: Mode, value: CELL) {
		let addr = self.ip;
		self.ip += 1;
		match mode {
			1 => self.data[addr] = value,
			_ => {
				let addr = self.data[addr] as usize;
				self.data[addr] = value;
			},
		}
	}

	fn binop<F>(&mut self, op: F, modes: [Mode; 3])
	where
		F: FnOnce(CELL, CELL) -> CELL,
	{
		let src1 = self.fetch(modes[0]);
		let src2 = self.fetch(modes[1]);
		self.store(modes[2], op(src1, src2));
	}

	pub fn simulate(&mut self) {
		assert_eq!(self.ip, 0);
		assert!(self.next_input.is_none());
		self.simulate_step().expect_finished();
	}

	pub fn simulate_step(&mut self) -> SimulateStep {
		loop {
			let modes = self.data[self.ip];
			self.ip += 1;
			let op = modes % 100;
			let modes = modes / 100;
			let mode1 = (modes % 10) as Mode;
			let modes = modes / 10;
			let mode2 = (modes % 10) as Mode;
			let modes = modes / 10;
			let mode3 = (modes % 10) as Mode;
			let modes = [mode1, mode2, mode3];

			match op {
				OP_ADD => self.binop(|a, b| a + b, modes),
				OP_MUL => self.binop(|a, b| a * b, modes),
				OP_LESS_THAN => self.binop(|a, b| if a < b { 1 } else { 0 }, modes),
				OP_EQUAL => self.binop(|a, b| if a == b { 1 } else { 0 }, modes),
				OP_INPUT => {
					if let Some(v) = self.next_input.take() {
						self.store(mode1, v);
					} else {
						self.ip -= 1; // go back to opcode
						return SimulateStep::WaitForInput;
					}
				},
				OP_OUTPUT => {
					return SimulateStep::Output(self.fetch(mode1));
				},
				OP_HALT => {
					self.ip -= 1; // make re-running this step end up here again
					return SimulateStep::Finished;
				},
				OP_JUMP_IF_TRUE => {
					if self.fetch(mode1) != 0 {
						self.ip = self.fetch(mode2) as usize;
					} else {
						self.ip += 1; // skip target
					}
				},
				OP_JUMP_IF_FALSE => {
					if self.fetch(mode1) == 0 {
						self.ip = self.fetch(mode2) as usize;
					} else {
						self.ip += 1; // skip target
					}
				},
				_ => panic!("invalid opcode: {}", op),
			}
		}
	}
}

impl std::str::FromStr for IntCode {
	type Err = <CELL as std::str::FromStr>::Err;

	fn from_str(s: &str) -> Result<Self, Self::Err> {
		Ok(IntCode {
			ip: 0,
			next_input: None,
			data: s.split(',').map(|s| s.trim().parse::<CELL>()).collect::<Result<_, _>>()?,
		})
	}
}