aoc19intcode 0.3.2

//! # A standalone Advent of Code 2019 Intcode VM implementation
//!
//! That's it. IDK what else to tell ya.
//!
//! # Examples
//!
//! ```
//! use aoc19intcode::IntcodeVM;
//! assert_eq!(
//!     IntcodeVM::from_prog(&[2, 4, 4, 5, 99, 0])
//!         .run_prog()
//!         .unwrap()[&5],
//!     9801
//! );
//! ```

use std::{
    cmp::{Ordering, PartialEq},
    collections::HashMap,
    fmt,
    ops::{Add, Index, IndexMut, Mul},
    sync::mpsc::{Receiver, Sender},
};

#[derive(Clone, Copy, PartialEq, Debug)]
enum Mode {
    Position,
    Immediate,
    Relative,
}

type Cell = i64;
type Tape = HashMap<usize, Cell>;

/// A VM that runs Intcode programs.
///
/// # Examples
///
/// ```
/// use aoc19intcode::IntcodeVM;
/// assert_eq!(
///     IntcodeVM::from_prog(&[2, 4, 4, 5, 99, 0])
///         .run_prog()
///         .expect("Tried to run a halted program.")[&5],
///     9801
/// );
/// ```
/// ```
/// use std::{thread, sync::mpsc};
/// use aoc19intcode::IntcodeVM;
/// let input = [2, 4, 4, 5, 99, 0];
/// let (mtx, prx) = mpsc::channel();
/// let (ptx, mrx) = mpsc::channel();
/// let (ftx, frx) = mpsc::channel();
/// mtx.send(2).expect("Failed to send an input to the VM.");
/// thread::spawn(move || {
///     IntcodeVM::with_inp_flag(&input, Some(prx), Some(ptx), Some(ftx))
///         .run_prog()
/// });
/// for _ in frx.recv() {
///     println!("received {}", mrx.recv().expect("VM disconnected."));
/// }
/// ```
#[derive(Debug, Default)]
pub struct IntcodeVM {
    tape: Tape,
    ip: usize,
    rp: isize,
    input: Option<Receiver<Cell>>,
    output: Option<Sender<Cell>>,
    inp_wait_flag: Option<Sender<()>>,
    halted: bool,
}

trait Flags {
    fn flags(&self) -> Result<(Mode, Mode, Mode), IntcodeError>;
}

impl Flags for Cell {
    fn flags(&self) -> Result<(Mode, Mode, Mode), IntcodeError> {
        let flag = |argpos| match self / 10_i64.pow(argpos) % 10 {
            0 => Ok(Mode::Position),
            1 => Ok(Mode::Immediate),
            2 => Ok(Mode::Relative),
            _ => Err(IntcodeError::InvalidInstruction),
        };
        Ok((flag(2)?, flag(3)?, flag(4)?))
    }
}

impl IntcodeVM {
    /// Returns a new VM with a program and no I/O controls.
    pub fn from_prog(prog: &[Cell]) -> Self {
        Self {
            tape: prog.iter().copied().enumerate().collect(),
            ..Default::default()
        }
    }

    /// Returns a new VM with a program and I/O controls.
    pub fn with_io(
        prog: &[Cell],
        input: Option<Receiver<Cell>>,
        output: Option<Sender<Cell>>,
    ) -> Self {
        Self {
            tape: prog.iter().copied().enumerate().collect(),
            input,
            output,
            ..Default::default()
        }
    }

    /// Returns a new VM with a program and I/O controls, includes an input
    /// waiting flag sender that sends a unit (`()`) when the VM is expecting
    /// input.
    pub fn with_inp_flag(
        prog: &[Cell],
        input: Option<Receiver<Cell>>,
        output: Option<Sender<Cell>>,
        inp_wait_flag: Option<Sender<()>>,
    ) -> Self {
        Self {
            tape: prog.iter().copied().enumerate().collect(),
            input,
            output,
            inp_wait_flag,
            ..Default::default()
        }
    }

    /// Runs the VM.
    pub fn run_prog(&mut self) -> Result<Tape, IntcodeError> {
        if self.halted {
            return Err(IntcodeError::RanHaltedVM);
        }
        while !self.halted {
            self.step()?;
        }
        Ok(self.tape.clone())
    }

    // Runs one instruction.
    pub fn step(&mut self) -> Result<Tape, IntcodeError> {
        if self.halted {
            return Err(IntcodeError::RanHaltedVM);
        }
        match self[self.ip] % 100 {
            1 => self.maths(<i64 as Add>::add)?,
            2 => self.maths(<i64 as Mul>::mul)?,
            3 => self.input()?,
            4 => self.outpt()?,
            5 => self.test0(<i64 as PartialEq>::ne)?,
            6 => self.test0(<i64 as PartialEq>::eq)?,
            7 => self.test2(Ordering::Less)?,
            8 => self.test2(Ordering::Equal)?,
            9 => self.relpt()?,
            99 => self.halted = true,
            _ => return Err(IntcodeError::InvalidInstruction),
        }

        Ok(self.tape.clone())
    }

    fn get(&mut self, mode: Mode, pad: usize) -> Cell {
        let in1 = self.ip + pad;
        let in2 = match mode {
            Mode::Immediate => return self[in1],
            Mode::Position => self[in1] as usize,
            Mode::Relative => (self[in1] as isize + self.rp) as usize,
        };
        *self.tape.entry(in2).or_default()
    }

    fn maths(&mut self, op: fn(Cell, Cell) -> Cell) -> Result<(), IntcodeError> {
        let (one, two, three) = self[self.ip].flags()?;
        let arg1 = self.get(one, 1);
        let arg2 = self.get(two, 2);
        let rel = if three == Mode::Relative { self.rp } else { 0 };
        let loc = (self[self.ip + 3] + rel as i64) as usize;
        self[loc] = op(arg1, arg2);
        self.ip += 4;
        Ok(())
    }

    fn input(&mut self) -> Result<(), IntcodeError> {
        let (one, ..) = self[self.ip].flags()?;
        let rel = if one == Mode::Relative { self.rp } else { 0 };
        let loc = (self[self.ip + 1] + rel as i64) as usize;
        if let Some(sender) = &self.inp_wait_flag {
            sender.send(()).expect(
                "Error: failed to send input waiting flag: \
                    receiver disconnected.",
            );
        }
        self[loc] = self
            .input
            .as_ref()
            .expect("Error: found input instruction but no receiver was created.")
            .recv()
            .expect("Error: failed to receive input: sender disconnected.");
        self.ip += 2;
        Ok(())
    }

    fn outpt(&mut self) -> Result<(), IntcodeError> {
        let (one, ..) = self[self.ip].flags()?;
        let val = self.get(one, 1);
        self.output
            .as_ref()
            .expect("Error: found output instruction but no sender was created.")
            .send(val)
            .expect("Error: failed to send output: receiver disconnected.");
        self.ip += 2;
        Ok(())
    }

    // compare a cell against 0
    fn test0(&mut self, checker: fn(&Cell, &Cell) -> bool) -> Result<(), IntcodeError> {
        let (one, two, _) = self[self.ip].flags()?;
        let test = self.get(one, 1);
        if checker(&test, &0) {
            self.ip = self.get(two, 2) as usize;
        } else {
            self.ip += 3;
        }
        Ok(())
    }

    // compare 2 arguments against each other
    fn test2(&mut self, mode: Ordering) -> Result<(), IntcodeError> {
        let (one, two, three) = self[self.ip].flags()?;
        let arg1 = self.get(one, 1);
        let arg2 = self.get(two, 2);
        let rel = if three == Mode::Relative { self.rp } else { 0 };
        let loc = (self[self.ip + 3] + rel as i64) as usize;
        let test = arg1.cmp(&arg2) == mode;
        self[loc] = test as i64;
        self.ip += 4;
        Ok(())
    }

    fn relpt(&mut self) -> Result<(), IntcodeError> {
        let (one, ..) = self[self.ip].flags()?;
        let arg1 = self.get(one, 1);
        self.rp += arg1 as isize;
        self.ip += 2;
        Ok(())
    }
}

impl Index<usize> for IntcodeVM {
    type Output = Cell;

    fn index(&self, ix: usize) -> &Self::Output {
        &self.tape[&ix]
    }
}

impl IndexMut<usize> for IntcodeVM {
    fn index_mut(&mut self, ix: usize) -> &mut Self::Output {
        self.tape.entry(ix).or_default();
        self.tape.get_mut(&ix).unwrap()
    }
}

/// Possible errors.
pub enum IntcodeError {
    RanHaltedVM,
    InvalidInstruction,
}

impl std::error::Error for IntcodeError {}

impl fmt::Display for IntcodeError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let err_msg = match self {
            IntcodeError::RanHaltedVM => "Error: tried to run an already halted intcode VM",
            IntcodeError::InvalidInstruction => "Error: an instruction was malformed",
        };
        f.write_str(err_msg)
    }
}

impl fmt::Debug for IntcodeError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let err_msg = match self {
            IntcodeError::RanHaltedVM => "Error: tried to run an already halted intcode VM",
            IntcodeError::InvalidInstruction => "Error: an instruction was malformed",
        };
        f.write_str(err_msg)
    }
}