Rusty intcode

intcode.rs

use thiserror::Error;
use std::convert::{TryFrom, TryInto};
use std::collections::{VecDeque, HashMap};
use std::path::Path;

//==============================================================================

/// Represents an Intcode code file.
#[derive(Debug, Clone)]
pub struct Intcode {
    code: HashMap<usize, i64>,
}

impl Intcode {
    pub fn from_path<P: AsRef<Path>>(path: P) -> Self {
        std::fs::read_to_string(path).expect("error reading file")
            .trim()
            .split(",")
            .map(|s| s.parse::<i64>().expect("error parsing intcode"))
            .collect()
    }

    pub fn run<I: IntoIterator<Item=i64>>(&self, input: I) -> Runner<I::IntoIter> {
        Runner::new(self.clone(), input.into_iter())
    }
}

impl std::iter::FromIterator<i64> for Intcode {
    fn from_iter<I: IntoIterator<Item=i64>>(iter: I) -> Self {
        Intcode {
            code: iter.into_iter().enumerate().collect(),
        }
    }
}

impl std::ops::Index<usize> for Intcode {
    type Output = i64;

    fn index(&self, i: usize) -> &Self::Output {
        match self.code.get(&i) {
            Some(p) => p,
            None => &0,
        }
    }
}

impl std::ops::IndexMut<usize> for Intcode {
    fn index_mut(&mut self, i: usize) -> &mut Self::Output {
        self.code.entry(i).or_insert(0)
    }
}

//==============================================================================

/// Represents a running Intcode interpreter.
#[derive(Debug, Clone)]
pub struct Runner<I> {
    cur: usize,
    rel_base: usize,
    code: Intcode,
    pub input: I,
}

#[derive(Debug, Copy, Clone, PartialOrd, Ord, PartialEq, Eq, Hash)]
pub enum Status {
    /// The Intcode interpreter has paused after executing an output instruction. (opcode 4)
    Output(i64),
    /// The Intcode interpreter has paused at an input instruction (opcode 3) because it
    /// has completely drained its input iterator.  It will continue to produce this status
    /// until more elements are added.
    ///
    /// Not sure how to meet these demands?  Try using an [`ExtendableIter`] for the input.
    WaitingForInput,
    /// The Intcode interpreter has stopped after encountering opcode 99.
    Stopped,
}

impl Status {
    pub fn unwrap_or_else<F>(&self, func: F) -> i64
    where
        F: FnOnce(Status) -> i64,
    {
        match *self {
            Status::Output(x) => x,
            status => func(status),
        }
    }
}

impl<I: Iterator<Item=i64>> Runner<I> {
    pub fn new(code: Intcode, input: I) -> Self {
        let cur = 0;
        let rel_base = 0;
        Runner { cur, rel_base, code, input }
    }

    /// See the current state of the intcode.
    ///
    /// Intcode is self-modifying, so it may be different from the initial code.
    pub fn code(&self) -> &Intcode {
        &self.code
    }

    /// Advance the interpreter until one of the following occurs:
    ///
    /// * An output instruction (opcode 4).
    /// * An input instruction (opcode 3) when the input iterator is empty.
    /// * A stop instruction (opcode 99).
    pub fn step(&mut self) -> Status {
        self.step_impl()
    }

    /// Step repeatedly, printing all ascii characters to STDOUT.
    pub fn print_until_non_ascii(&mut self) -> Status {
        loop {
            match self.step() {
                Status::Output(x) if 0 < x && x < 128 => print!("{}", x as u8 as char),
                status => return status,
            }
        }
    }
}

//==============================================================================
// Intcode runner implementation details

macro_rules! deriving_try_from {
    (
        #[try_from_error($Error:path)]
        $(#[$($meta:tt)+])+
        $vis:vis enum $Enum:ident {
            $(
                $Variant:ident = $value:expr,
            )+
        }
    ) => {
        $(#[$($meta)+])+
        $vis enum $Enum {
            $( $Variant = $value, )+
        }

        impl TryFrom<i64> for $Enum {
            type Error = $Error;

            fn try_from(x: i64) -> Result<Self, Self::Error> {
                match x {
                    $( $value => Ok($Enum::$Variant), )+
                    value => Err($Error(value)),
                }
            }
        }
    };
}

#[derive(Error, Debug)]
#[error("bad opcode: {0}")]
struct OpcodeError(i64);

#[derive(Error, Debug)]
#[error("bad parameter mode: {0}")]
struct ParameterModeError(i64);

deriving_try_from! {
    #[try_from_error(OpcodeError)]
    #[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord)]
    enum Opcode {
        Add = 1,
        Mul = 2,
        Read = 3,
        Write = 4,
        JumpIfTrue = 5,
        JumpIfFalse = 6,
        LessThan = 7,
        EqualTo = 8,
        AdjustRel = 9,
        Stop = 99,
    }
}

deriving_try_from! {
    #[try_from_error(ParameterModeError)]
    #[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord)]
    enum ParameterMode {
        Absolute = 0,
        Immediate = 1,
        Relative = 2,
    }
}

impl<I: Iterator<Item=i64>> Runner<I> {
    fn step_impl(&mut self) -> Status {
        loop {
            let last_pos = self.cur;

            let command = self.advance();
            assert!(command > 0, "negative opcode!? {}", command);
            let opcode = Opcode::try_from(command % 100).unwrap_or_else(|e| panic!("{}", e));
            let parameter_modes = [
                ParameterMode::try_from((command / 100) % 10).unwrap_or_else(|e| panic!("{}", e)),
                ParameterMode::try_from((command / 1000) % 10).unwrap_or_else(|e| panic!("{}", e)),
                ParameterMode::try_from((command / 10000) % 10).unwrap_or_else(|e| panic!("{}", e)),
            ];

            fn add_offset(a: usize, b: i64) -> usize {
                (a as i64 + b).try_into().unwrap_or_else(|e| panic!("{}", e))
            }
            fn resolve_read(value: i64, code: &Intcode, rel_base: usize, mode: ParameterMode) -> i64 {
                match mode {
                    ParameterMode::Absolute => code[add_offset(0, value)],
                    ParameterMode::Relative => code[add_offset(rel_base, value)],
                    ParameterMode::Immediate => value,
                }
            }
            fn resolve_write(value: i64, code: &mut Intcode, rel_base: usize, mode: ParameterMode) -> &mut i64 {
                match mode {
                    ParameterMode::Absolute => &mut code[add_offset(0, value)],
                    ParameterMode::Relative => &mut code[add_offset(rel_base, value)],
                    ParameterMode::Immediate => panic!("attempt to write to immediate operand!"),
                }
            }

            match opcode {
                | Opcode::Add
                | Opcode::Mul
                | Opcode::EqualTo
                | Opcode::LessThan
                => {
                    let binop: fn(i64, i64) -> i64 = match opcode {
                        Opcode::Add => |a, b| a + b,
                        Opcode::Mul => |a, b| a * b,
                        Opcode::EqualTo => |a, b| (a == b) as i64,
                        Opcode::LessThan => |a, b| (a < b) as i64,
                        _ => unreachable!(),
                    };
                    let a = resolve_read(self.advance(), &self.code, self.rel_base, parameter_modes[0]);
                    let b = resolve_read(self.advance(), &self.code, self.rel_base, parameter_modes[1]);
                    let c = resolve_write(self.advance(), &mut self.code, self.rel_base, parameter_modes[2]);
                    *c = binop(a, b);
                },

                | Opcode::Read
                => {
                    let dest = resolve_write(self.advance(), &mut self.code, self.rel_base, parameter_modes[0]);
                    match self.input.next() {
                        None => {
                            self.cur = last_pos; // repeat this instruction next time we step
                            return Status::WaitingForInput;
                        },
                        Some(value) => *dest = value,
                    }
                },

                | Opcode::Write
                => {
                    let value = resolve_read(self.advance(), &self.code, self.rel_base, parameter_modes[0]);
                    return Status::Output(value);
                },

                | Opcode::JumpIfFalse
                | Opcode::JumpIfTrue
                => {
                    let value = resolve_read(self.advance(), &self.code, self.rel_base, parameter_modes[0]);
                    let dest = resolve_read(self.advance(), &self.code, self.rel_base, parameter_modes[1]);
                    let is_true = value != 0;
                    if is_true == (opcode == Opcode::JumpIfTrue) {
                        self.cur = dest.try_into().unwrap_or_else(|e| panic!("{}", e));
                    }
                },

                | Opcode::AdjustRel
                => {
                    let amount = resolve_read(self.advance(), &self.code, self.rel_base, parameter_modes[0]);
                    self.rel_base = add_offset(self.rel_base, amount);
                },

                | Opcode::Stop
                => {
                    self.cur = last_pos; // make sure we stay at this instruction
                    return Status::Stopped;
                }
            }
        }
    }
    
    fn advance(&mut self) -> i64 {
        let out = self.code[self.cur];
        self.cur += 1;
        out
    }
}

//==============================================================================
// Utility iterator types for users to use as inputs

/// An iterator over a sequence of elements, onto which more elements can be added at any time.
///
/// You can even add elements after this returns `None`, and it will begin to produce `Some(_)`
/// again! (this behavior is in fact crucial at times for communicating with Intcode)
#[derive(Debug, Clone)]
pub struct ExtendableIter<T> {
    data: VecDeque<T>,
}

impl<T> ExtendableIter<T> {
    pub fn new() -> Self {
        Self { data: VecDeque::new() }
    }
}

impl<T> Default for ExtendableIter<T> {
    fn default() -> Self { Self::new() }
}

impl<T> ExtendableIter<T> {
    pub fn push(&mut self, x: T) {
        self.data.push_back(x);
    }
}

impl<T> Extend<T> for ExtendableIter<T> {
    fn extend<I: IntoIterator<Item=T>>(&mut self, iter: I) {
        self.data.extend(iter)
    }
}

impl<T> std::iter::FromIterator<T> for ExtendableIter<T> {
    fn from_iter<I: IntoIterator<Item=T>>(iter: I) -> Self {
        let mut out = Self::new();
        out.extend(iter);
        out
    }
}

impl<T> Iterator for ExtendableIter<T> {
    type Item = T;

    fn next(&mut self) -> Option<T> { self.data.pop_front() }
    fn size_hint(&self) -> (usize, Option<usize>) { (0, None) }
}

impl ExtendableIter<i64> {
    /// Helper for UTF-8 and ASCII input.
    pub fn extend_from_utf8(&mut self, s: &str) {
        self.extend(s.bytes().map(|b| b as i64))
    }
}

/// An iterator that keeps producing a single value, which can be changed at any time.
#[derive(Debug, Clone)]
pub struct ConstantIter<T> {
    value: T,
}

impl<T> ConstantIter<T> {
    pub fn new(value: T) -> Self {
        ConstantIter { value }
    }

    pub fn set_value(&mut self, value: T) {
        self.value = value;
    }
}

impl<T: Clone> Iterator for ConstantIter<T> {
    type Item = T;

    fn next(&mut self) -> Option<T> {
        Some(self.value.clone())
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        (usize::max_value(), None)
    }
}