aoc2019/shared/src/intcode.rs

use failure::Error;

#[repr(u8)]
#[derive(PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Copy, Debug)]
pub enum ParameterMode {
	Position = 0,
	Immediate = 1,
}

impl ParameterMode {
	pub fn decode(code: CELL) -> Result<(Self, CELL), Error> {
		let rem = code / 10;
		let mode = match code % 10 {
			0 => Self::Position,
			1 => Self::Immediate,
			m => failure::bail!("Unknown parameter mode: {}", m),
		};
		Ok((mode, rem))
	}
}

#[repr(u8)]
#[derive(PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Copy, Debug)]
pub enum Instruction {
	Add = 1,
	Mul = 2,
	In = 3,
	Out = 4,
	JumpIfTrue = 5,
	JumpIfFalse = 6,
	LessThan = 7,
	Equal = 8,
	Halt = 99,
}

impl Instruction {
	pub fn decode(opcode: CELL) -> Result<(Self, CELL), Error> {
		let rem = opcode / 100;
		let instr = match opcode % 100 {
			1 => Self::Add,
			2 => Self::Mul,
			3 => Self::In,
			4 => Self::Out,
			5 => Self::JumpIfTrue,
			6 => Self::JumpIfFalse,
			7 => Self::LessThan,
			8 => Self::Equal,
			99 => Self::Halt,
			n => failure::bail!("Unknown instruction code {}", n),
		};
		Ok((instr, rem))
	}
}

#[derive(PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Copy, Debug)]
pub struct OpCode {
	pub instruction: Instruction,
	pub modes: [ParameterMode; 3],
}

impl OpCode {
	pub fn decode(opcode: CELL) -> Result<OpCode, Error> {
		let (instruction, code) = Instruction::decode(opcode)?;
		let (mode0, code) = ParameterMode::decode(code)?;
		let (mode1, code) = ParameterMode::decode(code)?;
		let (mode2, code) = ParameterMode::decode(code)?;
		failure::ensure!(code == 0);
		let modes = [mode0, mode1, mode2];

		Ok(OpCode {
			instruction,
			modes,
		})
	}
}

pub type CELL = i64;

#[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 {
	/// instruction pointer
	pub ip: usize,
	/// next value to use for an "IN" (3) instruction (which will clear it)
	pub next_input: Option<CELL>,
	/// memory of machine (both code and data)
	pub data: Vec<CELL>,
}

impl IntCode {
	#[inline(always)]
	fn fetch(&mut self, modes: &mut &[ParameterMode]) -> CELL {
		let mode = modes[0];
		*modes = &modes[1..];
		let imm = self.data[self.ip];
		self.ip += 1;
		match mode {
			ParameterMode::Position => self.data[imm as usize],
			ParameterMode::Immediate => imm,
		}
	}

	#[inline(always)]
	fn store(&mut self, modes: &mut &[ParameterMode], value: CELL) {
		let mode = modes[0];
		*modes = &modes[1..];
		let addr = self.ip;
		self.ip += 1;
		match mode {
			ParameterMode::Immediate => self.data[addr] = value,
			ParameterMode::Position => {
				let addr = self.data[addr] as usize;
				self.data[addr] = value;
			},
		}
	}

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

	/// run simple machine without I/O to completion
	pub fn simulate(&mut self) {
		assert_eq!(self.ip, 0);
		assert!(self.next_input.is_none());
		self.simulate_step().expect_finished();
	}

	/// run a machine until:
	/// - it needs input, but `next_input` is `None`
	/// - it outputs a value
	/// - it finishes
	pub fn simulate_step(&mut self) -> SimulateStep {
		loop {
			let opcode = OpCode::decode(self.data[self.ip]).unwrap();
			let modes = &mut &opcode.modes[..];
			self.ip += 1;

			match opcode.instruction {
				Instruction::Add => self.binop(|a, b| a + b, modes),
				Instruction::Mul => self.binop(|a, b| a * b, modes),
				Instruction::LessThan => self.binop(|a, b| if a < b { 1 } else { 0 }, modes),
				Instruction::Equal => self.binop(|a, b| if a == b { 1 } else { 0 }, modes),
				Instruction::In => {
					if let Some(v) = self.next_input.take() {
						self.store(modes, v);
					} else {
						self.ip -= 1; // go back to opcode
						return SimulateStep::WaitForInput;
					}
				},
				Instruction::Out => {
					return SimulateStep::Output(self.fetch(modes));
				},
				Instruction::Halt => {
					self.ip -= 1; // make re-running this step end up here again
					return SimulateStep::Finished;
				},
				Instruction::JumpIfTrue => {
					if self.fetch(modes) != 0 {
						self.ip = self.fetch(modes) as usize;
					} else {
						self.ip += 1; // skip parameter for target
					}
				},
				Instruction::JumpIfFalse => {
					if self.fetch(modes) == 0 {
						self.ip = self.fetch(modes) as usize;
					} else {
						self.ip += 1; // skip parameter for target
					}
				},
			}
		}
	}
}

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<_, _>>()?,
		})
	}
}