Advent of Code 2019 - Day 05 // Intcode Computer

main.rs

use std::collections::{VecDeque, HashSet};
use std::fmt::{self, Display, Formatter};
use std::io;
use std::process::exit;

// This implementation is needlessly complicated,
// but it comes with a simple debugger which I find kinda neat.
// Just invoke vm.add_breakpoint(0) before vm.run() and use
// simple commands via stdin.

// type alias for the data we store in the memory of the VM
type IntcodeInt = isize;

// We use IOStream as a simple mock of stdin and stdout
pub struct IOStream {
    buffer: VecDeque<IntcodeInt>
}

// We just give it simple FIFO capabilities
impl IOStream {
    pub fn new() -> IOStream {
        IOStream{buffer: VecDeque::new()}
    }

    pub fn push(&mut self, value: IntcodeInt) {
        self.buffer.push_back(value);
    }

    pub fn pop(&mut self) -> Option<IntcodeInt> {
        self.buffer.pop_front()
    }
}

#[repr(u8)]
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
enum Parameter {
    Position(IntcodeInt) = 0,
    Immediate(IntcodeInt) = 1
}

impl Display for Parameter {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        match self {
            Parameter::Position(value) => write!(f, "[0x{:x}]", value),
            Parameter::Immediate(value) => write!(f, "0x{:x}", value)
        }
    }
}

impl Parameter {
    fn fetch(&self, memory: &[IntcodeInt]) -> IntcodeInt {
        match self {
            Parameter::Position(value) => memory[*value as usize],
            Parameter::Immediate(value) => *value
        }
    }

    // A write destination is always a position, but it is used like an immediate
    // (we do not load the address to write to from memory)
    // therefore using fetch is wrong for write destinations
    fn unpack(&self) -> IntcodeInt {
        match self {
            Parameter::Position(value) => *value,
            Parameter::Immediate(value) => *value
        }
    }
}

#[repr(u8)]
enum Instruction {
    Add(Parameter, Parameter, Parameter) = 1,
    Mul(Parameter, Parameter, Parameter) = 2,
    Input(Parameter) = 3,
    Output(Parameter) = 4,
    JumpIfTrue(Parameter, Parameter) = 5,
    JumpIfFalse(Parameter, Parameter) = 6,
    LessThan(Parameter, Parameter, Parameter) = 7,
    Equals(Parameter, Parameter, Parameter) = 8,
    Halt = 99,
}

impl Display for Instruction {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        match self {
            Instruction::Add(src1, src2, dest) => write!(f, "ADD {}, {}, {}", src1, src2, dest),
            Instruction::Mul(src1, src2, dest) => write!(f, "MUL {}, {}, {}", src1, src2, dest),
            Instruction::Input(dest) => write!(f, "INP {}", dest),
            Instruction::Output(src) => write!(f, "OUT {}", src),
            Instruction::JumpIfTrue(cond, target) => write!(f, "JIT {}, {}", cond, target),
            Instruction::JumpIfFalse(cond, target) => write!(f, "JIF {}, {}", cond, target),
            Instruction::LessThan(src1, src2, dest) => write!(f, "LT {}, {}, {}", src1, src2, dest),
            Instruction::Equals(src1, src2, dest) => write!(f, "EQ {}, {}, {}", src1, src2, dest),
            Instruction::Halt => write!(f, "HLT")
        }
    }
}

impl Instruction {
    // I hoped there was a way to get the number of parameters from the enum,
    // but alas, there is not.
    fn get_num_params(opcode: u8) -> usize {
        match opcode {
            1 => 3,
            2 => 3,
            3 => 1,
            4 => 1,
            5 => 2,
            6 => 2,
            7 => 3,
            8 => 3,
            99 => 0,
            _ => 0,
        }
    }
    
    fn to_u8(&self) -> u8 {
        // I hoped this would work automatically if we used #[repr(u8)] on the enum,
        // but alas, it does not.
        match self {
            Instruction::Add(_, _, _) => 1,
            Instruction::Mul(_, _, _) => 2,
            Instruction::Input(_) => 3,
            Instruction::Output(_) => 4,
            Instruction::JumpIfTrue(_, _) => 5,
            Instruction::JumpIfFalse(_, _) => 6,
            Instruction::LessThan(_, _, _) => 7,
            Instruction::Equals(_, _, _) => 8,
            Instruction::Halt => 99
        }
    }
}

pub struct IntcodeVM {
    memory: Vec<IntcodeInt>,
    ip: usize,
    stdin: IOStream,
    stdout: IOStream,
    halted: bool,
    breakpoints: HashSet<usize>
}

impl IntcodeVM {
    pub fn new(state: Vec<IntcodeInt>) -> IntcodeVM {
        IntcodeVM { memory: state, ip: 0,
                    stdin: IOStream::new(),
                    stdout: IOStream::new(),
                    halted: false,
                    breakpoints: HashSet::new() }
    }

    pub fn readmem(&self, address: usize) -> IntcodeInt {
        self.memory[address]
    }

    pub fn input(&mut self, value: IntcodeInt) {
        self.stdin.push(value);
    }

    pub fn input_seq(&mut self, values: &[IntcodeInt]) {
        for value in values {
            self.stdin.push(*value);
        }
    }

    pub fn disassemble(&self, start_address: usize, max_instructions: usize, markfirst: bool) {
        let mut ip = start_address;
        let mut disassembled = 0;
        while disassembled < max_instructions && ip < self.memory.len() {
            if let Ok(instruction) = self.fetch_and_decode(ip) {
                if ip == start_address && markfirst {
                    println!("0x{:06x}: {}  <------ IP", ip, instruction);
                } else {
                    println!("0x{:06x}: {}", ip, instruction);
                }
                ip += Instruction::get_num_params(instruction.to_u8()) + 1;
                disassembled += 1;
            } else {
                println!("0x{:06x}: Invalid opcode", ip);
                return;
            }
        }
    }

    fn read_input(&mut self) -> Option<IntcodeInt> {
        self.stdin.pop()
    }

    pub fn add_breakpoint(&mut self, address: usize) {
        self.breakpoints.insert(address);
    }

    pub fn remove_breakpoint(&mut self, address: usize) {
        self.breakpoints.remove(&address);
    }

    fn debug(&mut self) {
        let mut user_selection = String::new();
        let mut context_len = 5;
        let printhelp = || {
            println!("Debugger Commands:");
            println!("  c, continue: Continue execution");
            println!("  n, next, s, step: Step one instruction");
            println!("  p, print <address> [count]: Print memory contents");
            println!("  b, break <address>: Add breakpoint");
            println!("  d, delete <address>: Remove breakpoint");
            println!("  exit: Stop execution");
            println!("  context <count>: Set number of instructions to disassemble");
            println!("  h, help: Print this help");
        };
        printhelp();
        self.disassemble(self.ip, context_len, true);
        'debug_loop: loop {
            if io::stdin().read_line(&mut user_selection).is_err() {
                break 'debug_loop;
            }
            let selection = user_selection.trim();
            let command = selection.split_whitespace().next().unwrap_or("");
            let num_parse = |num_str: &str| -> Option<usize> {
                // parse decimal or hex
                if num_str.starts_with("0x") {
                    usize::from_str_radix(&num_str[2..], 16).ok()
                } else {
                    usize::from_str_radix(num_str, 10).ok()
                }
            };
            // args are always numbers, but should also parse for hex numbers
            let args = selection.split_whitespace()
                                .skip(1)
                                .map(|arg| num_parse(arg).unwrap_or(0))
                                .collect::<Vec<_>>();

            if ["c", "continue"].contains(&command) {
                break 'debug_loop;
            }

            if ["n", "next", "s", "step"].contains(&command) {
                self.step();
                self.disassemble(self.ip, context_len, true);
                println!("");
            };

            if ["p", "print"].contains(&command) {
                if args.len() == 1 {
                    println!("0x{:06x}: {}", args[0], self.readmem(args[0]));
                } else if args.len() == 2 {
                    println!("0x{:06x}: {:?}", args[0], &self.memory[args[0]..args[0] + args[1]]);
                }
            }

            if ["b", "break"].contains(&command) {
                if args.len() == 1 {
                    self.add_breakpoint(args[0]);
                } else {
                    println!("Invalid arguments");
                }
            }

            if ["d", "delete"].contains(&command) {
                if args.len() == 1 {
                    self.remove_breakpoint(args[0]);
                } else {
                    println!("Invalid arguments");
                }
            }

            if ["exit"].contains(&command) {
                exit(0);
            }

            if ["context"].contains(&command) {
                if args.len() == 1 {
                    context_len = args[0];
                } else {
                    println!("Invalid arguments");
                }
            }

            if ["h", "help"].contains(&command) {
                printhelp();
            }

            user_selection.clear();
        }
    }

    fn output(&mut self, value: IntcodeInt) {
        self.stdout.push(value);
    }

    fn fetch_and_decode(&self, ip: usize) -> Result<Instruction, String> {
        let opcode_mode = self.memory[ip];
        let opcode = (opcode_mode % 100) as u8;
        let mut mode = opcode_mode / 100;
        let num_parameters = Instruction::get_num_params(opcode);
        let mut params = Vec::with_capacity(num_parameters);
        for i in 0..num_parameters {
            let param = self.memory[self.ip + i + 1];
            params.push(match mode % 10 {
                0 => Parameter::Position(param),
                1 => Parameter::Immediate(param),
                _ => panic!("Invalid parameter mode")
            });
            mode /= 10;
        }
        match opcode {
            1 => Ok(Instruction::Add(params[0], params[1], params[2])),
            2 => Ok(Instruction::Mul(params[0], params[1], params[2])),
            3 => Ok(Instruction::Input(params[0])),
            4 => Ok(Instruction::Output(params[0])),
            5 => Ok(Instruction::JumpIfTrue(params[0], params[1])),
            6 => Ok(Instruction::JumpIfFalse(params[0], params[1])),
            7 => Ok(Instruction::LessThan(params[0], params[1], params[2])),
            8 => Ok(Instruction::Equals(params[0], params[1], params[2])),
            99 => Ok(Instruction::Halt),
            _ => Err(format!("Invalid opcode: {}", opcode))
        }
    }

    fn execute(&mut self, instruction: Instruction) {
        match instruction {
            Instruction::Add(src1, src2, dest) => {
                assert!(matches!(dest, Parameter::Position(_)));
                let src1 = src1.fetch(&self.memory);
                let src2 = src2.fetch(&self.memory);
                self.memory[dest.unpack() as usize] = src1 + src2;
                self.ip += 4;
            },
            Instruction::Mul(src1, src2, dest) => {
                assert!(matches!(dest, Parameter::Position(_)));
                let src1 = src1.fetch(&self.memory);
                let src2 = src2.fetch(&self.memory);
                self.memory[dest.unpack() as usize] = src1 * src2;
                self.ip += 4;
            },
            Instruction::Input(dest) => {
                assert!(matches!(dest, Parameter::Position(_)));
                let input = self.read_input().unwrap();
                self.memory[dest.unpack() as usize] = input;
                self.ip += 2;
            },
            Instruction::Output(src) => {
                let src = src.fetch(&self.memory);
                self.output(src);
                self.ip += 2;
            },
            Instruction::JumpIfTrue(cond, target) => {
                let cond = cond.fetch(&self.memory);
                let target = target.fetch(&self.memory);
                if cond != 0 {
                    self.ip = target as usize;
                } else {
                    self.ip += 3;
                }
            },
            Instruction::JumpIfFalse(cond, target) => {
                let cond = cond.fetch(&self.memory);
                let target = target.fetch(&self.memory);
                if cond == 0 {
                    self.ip = target as usize;
                } else {
                    self.ip += 3;
                }
            },
            Instruction::LessThan(src1, src2, dest) => {
                assert!(matches!(dest, Parameter::Position(_)));
                let src1 = src1.fetch(&self.memory);
                let src2 = src2.fetch(&self.memory);
                self.memory[dest.unpack() as usize] = if src1 < src2 { 1 } else { 0 };
                self.ip += 4;
            },
            Instruction::Equals(src1, src2, dest) => {
                assert!(matches!(dest, Parameter::Position(_)));
                let src1 = src1.fetch(&self.memory);
                let src2 = src2.fetch(&self.memory);
                self.memory[dest.unpack() as usize] = if src1 == src2 { 1 } else { 0 };
                self.ip += 4;
            },
            Instruction::Halt => {
                self.ip += 1;
                self.halted = true;
            }
        }
    }


    pub fn step(&mut self) {
        if let Ok(instruction) = self.fetch_and_decode(self.ip) {
            self.execute(instruction);
        } else {
            eprintln!("Invalid opcode at 0x{:06x}", self.ip);
            self.halted = true;
        }
    }

    pub fn run(&mut self) {
        self.halted = false;
        while !self.halted && self.ip < self.memory.len() {
            if self.breakpoints.contains(&self.ip) {
                self.debug();
            }
            self.step();
        }
    }
}

fn part1(lines: &Vec<&str>) -> Option<i64> {
    assert_eq!(1, lines.len());
    let input_seq = lines[0].split(",")
                            .map(|n| n.parse::<IntcodeInt>().unwrap())
                            .collect::<Vec<_>>();
    let mut vm = IntcodeVM::new(input_seq);
    vm.input(1);
    // vm.add_breakpoint(0);
    vm.run();
    while let Some(output) = vm.stdout.pop() {
        if output != 0 {
            return Some(output as i64);
        }
    }
    None
}

fn part2(lines: &Vec<&str>) -> Option<i64> {
    assert_eq!(1, lines.len());
    let input_seq = lines[0].split(",")
                            .map(|n| n.parse::<IntcodeInt>().unwrap())
                            .collect::<Vec<_>>();
    let mut vm = IntcodeVM::new(input_seq);
    vm.input(5);
    vm.run();
    while let Some(output) = vm.stdout.pop() {
        if output != 0 {
            return Some(output as i64);
        }
    }
    None
}

fn main() {
    use std::fs;
    use std::env;
    use std::time::Instant;
    let args: Vec<String> =  env::args().collect();
    let infile = args.get(1).unwrap_or_else(|| {
        println!("Usage: {} <puzzle input>", args[0]);
        std::process::exit(1);
    });

    let contents = fs::read_to_string(infile)
        .expect("Could not read in file");

    let lines: Vec<&str> = contents.lines().collect();

    // execute part 1 and part 2, print their results if they exist
    // later parts may follow, so we loop over the part functions
    let parts = [part1, part2];
    for (index, part) in parts.iter().enumerate() {
        let partstart = Instant::now();
        let result = part(&lines);
        match result {
            Some(result) => println!("Part {}: {}\t({:?})", index+1, result, partstart.elapsed()),
            None => println!("Part {}: No result", index+1),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_intcodevm() {
        let state = vec![1,1,1,4,99,5,6,0,99];
        let mut vm = IntcodeVM::new(state);
        vm.run();
        assert_eq!([30,1,1,4,2,5,6,0,99].as_slice(), vm.memory.as_slice());
    }

    #[test]
    fn test_jmp_input_output() {
        let state = vec![3,21,1008,21,8,20,1005,20,22,107,8,21,20,1006,20,31,
                         1106,0,36,98,0,0,1002,21,125,20,4,20,1105,1,46,104,
                         999,1105,1,46,1101,1000,1,20,4,20,1105,1,46,98,99];
        let mut vm = IntcodeVM::new(state.clone());
        vm.input(7);
        vm.run();
        assert_eq!(vm.stdout.pop(), Some(999));
        let mut vm = IntcodeVM::new(state);
        vm.input(9);
        vm.run();
        assert_eq!(vm.stdout.pop(), Some(1001));
    }
}

This is very cool! It inspired me to make a tui app using ratatui. Great work!