AoC 2019 day 11

main.rs

use std::fs;
use std::collections::VecDeque;
use std::collections::HashMap;

#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
struct Point {
    x: isize,
    y: isize
}

#[derive(Copy, Clone)]
enum Direction {
    Up,
    Down,
    Left,
    Right
}

#[derive(Copy, Clone, Debug)]
enum Color {
    Black = 0,
    White = 1,
}

impl From<usize> for Color {
    fn from(n: usize) -> Color {
        match n {
            0 => Color::Black,
            1 => Color::White,
            _ => panic!("Unknown color!")
        }
    }
}

#[derive(Copy, Clone)]
enum RobotState {
    WaitingForColor,
    WaitingForDirection,
}


fn main() {
    let opcodes: Vec<isize> = fs::read_to_string("data.txt").unwrap()
            .split(",")
            .map(|x| x.trim().parse().unwrap())
            .collect();

    let mut painted: HashMap<Point, Color> = HashMap::new();

    let mut robot_position = Point { x: 0, y: 0 };
    let mut robot_state = RobotState::WaitingForColor;
    let mut robot_direction = Direction::Up;

    let mut vm = IntCodeVirtualMachine::new();
    vm.instructions = opcodes.clone();

    painted.insert(robot_position, Color::White);

    loop {
        match vm.run() {
            ProgramState::Halted => break,
            ProgramState::NeedInput => vm.push_input(*painted.get(&robot_position).unwrap_or(&Color::Black) as isize),
            ProgramState::Output(out) => {
                match robot_state {
                    RobotState::WaitingForColor => {
                        painted.insert(robot_position, (out as usize).into());
                        robot_state = RobotState::WaitingForDirection;
                    },
                    RobotState::WaitingForDirection => {
                        robot_direction = match out {
                            0 => {
                                match robot_direction {
                                    Direction::Up => Direction::Left,
                                    Direction::Down => Direction::Right,
                                    Direction::Left => Direction::Down,
                                    Direction::Right => Direction::Up,
                                }
                            },
                            1 => {
                                match robot_direction {
                                    Direction::Up => Direction::Right,
                                    Direction::Down => Direction::Left,
                                    Direction::Left => Direction::Up,
                                    Direction::Right => Direction::Down,
                                }
                            }
                            _ => panic!("Program passed invalid direction")
                        };

                        robot_position = Point { 
                            x: robot_position.x + match robot_direction {
                                Direction::Right => 1,
                                Direction::Left => -1,
                                _ => 0,
                            },
                            y: robot_position.y + match robot_direction {
                                Direction::Up => 1,
                                Direction::Down => -1,
                                _ => 0,
                            }
                        };

                        robot_state = RobotState::WaitingForColor;
                    }
                }
            }
        }
    }

    let least_x = painted.iter().fold(std::isize::MAX, |least, val| if val.0.x < least { val.0.x } else { least });
    let least_y = painted.iter().fold(std::isize::MAX, |least, val| if val.0.y < least { val.0.y } else { least });

    let points: HashMap<Point, Color> = painted.iter().map(|x| (Point { x: x.0.x + least_x.abs(), y: x.0.y + least_y.abs() }, x.1.clone())).collect();

    let max_x = points.iter().fold(std::isize::MIN, |most, val| if val.0.x > most { val.0.x } else { most });
    let max_y = points.iter().fold(std::isize::MIN, |most, val| if val.0.y > most { val.0.y } else { most });

    for x in 0..=max_x {
        for y in 0..=max_y {
            match points.get(&Point { x: x, y: y }).unwrap_or(&Color::Black) {
                Color::Black => print!(" "),
                Color::White => print!("█")
            }
        }
        println!();
    }
}

struct Opcode {
    instruction: isize,
    modes: [OpcodeMode; 3],
}

impl Opcode {
    fn parse(number: isize) -> Self {
        Opcode {
            instruction: get_digit(number, 2) * 10 + get_digit(number, 1),
            modes: [
                OpcodeMode::get_mode(get_digit(number, 3)),
                OpcodeMode::get_mode(get_digit(number, 4)),
                OpcodeMode::get_mode(get_digit(number, 5)),
            ],
        }
    }
}

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

impl OpcodeMode {
    fn get_mode(number: isize) -> Self {
        match number {
            0 => OpcodeMode::Position,
            1 => OpcodeMode::Immediate,
            2 => OpcodeMode::Relative,
            _ => panic!("Invalid opcode mode!"),
        }
    }
}

fn get_digit(number: isize, place: isize) -> isize {
    (number / 10isize.pow((place - 1) as u32)) % 10
}

struct IntCodeVirtualMachine {
    instructions: Vec<isize>,
    instruction_ptr: usize,
    current_opcode: Option<Opcode>,
    current_operands: [isize; 3],
    inputs: VecDeque<isize>,
    jumped: bool,
    relative_base: isize,
    memory: HashMap<usize, isize>,
}

enum ProgramState {
    Halted,
    NeedInput,
    Output(isize)
}

impl IntCodeVirtualMachine {
    fn new() -> Self {
        IntCodeVirtualMachine {
            instructions: vec![],
            memory: HashMap::new(),
            instruction_ptr: 0,
            current_opcode: None,
            current_operands: [0; 3],
            inputs: VecDeque::new(),
            jumped: false,
            relative_base: 0,
        }
    }

    fn reset(&mut self) {
        self.memory.clear();
        self.instructions.clear();
        self.instruction_ptr = 0;
        self.current_opcode = None;
        self.current_operands = [0; 3];
        self.jumped = false;
        self.relative_base = 0;
        self.inputs.clear();
    }

    fn write(&mut self, location: usize, number: isize) {
        if location >= self.instructions.len() {
            self.memory.insert(location - self.instructions.len(), number);
        } else {
            self.instructions[location] = number;
        }
    }

    fn read(&self, location: usize) -> isize {
        if location >= self.instructions.len() {
            *self.memory.get(&(location - self.instructions.len())).unwrap_or(&0)
        } else {
            self.instructions[location]
        }
    }

    fn push_input(&mut self, input: isize) {
        self.inputs.push_back(input);
    }

    fn run_completely(&mut self) -> Vec<isize> {
        let mut output = vec![];

        loop {
            match self.run() {
                ProgramState::Halted => break,
                ProgramState::NeedInput => panic!("Not enough input"),
                ProgramState::Output(x) => output.push(x),
            }
        }

        output
    }

    fn run(&mut self) -> ProgramState {
        if self.read(self.instruction_ptr) == 99 {
            return ProgramState::Halted;
        }

        self.current_opcode = Some(self.get_current_opcode());
        self.set_current_operands();

        let mut output: Option<isize> = None;

        while {
            if let Some(out) = output {
                return ProgramState::Output(out);
            }
            
            match self.current_opcode.as_ref().unwrap().instruction {
                1 => self.write(self.current_operands[2] as usize, self.current_operands[0] + self.current_operands[1]),
                2 => self.write(self.current_operands[2] as usize, self.current_operands[0] * self.current_operands[1]),
                3 => { 
                    let input = self.inputs.pop_front();
                    self.write(self.current_operands[0] as usize, match input {
                        Some(x) => x,
                        None => return ProgramState::NeedInput
                    }
                )},
                4 => output = Some(self.current_operands[0]),
                5 => if self.current_operands[0] != 0 { self.instruction_ptr = self.current_operands[1] as usize; self.jumped = true; },
                6 => if self.current_operands[0] == 0 { self.instruction_ptr = self.current_operands[1] as usize; self.jumped = true; },
                7 => self.write(self.current_operands[2] as usize, if self.current_operands[0] < self.current_operands[1] { 1 } else { 0 }),
                8 => self.write(self.current_operands[2] as usize, if self.current_operands[0] == self.current_operands[1] { 1 } else { 0 }),
                9 => self.relative_base += self.current_operands[0],
                _ => panic!("Invalid opcode!"),
            };

            self.next_instruction()
        } {}

        if let Some(out) = output {
            return ProgramState::Output(out);
        }

        ProgramState::Halted
    }

    fn next_instruction(&mut self) -> bool {
        if !self.jumped {
            self.instruction_ptr += match self.current_opcode.as_ref().unwrap().instruction {
                1 | 2 | 7 | 8 => 4,
                3 | 4 | 9 => 2,
                5 | 6 => 3,
                _ => panic!("Unknown opcode read!"),
            };
        }

        self.jumped = false;

        let next_opcode = self.get_current_opcode();

        if next_opcode.instruction == 99 {
            self.current_opcode = None;
            self.current_operands = [0, 0, 0];

            return false;
        }

        self.current_opcode = Some(next_opcode);
        self.set_current_operands();

        true
    }

    fn get_current_opcode(&self) -> Opcode {
        Opcode::parse(self.read(self.instruction_ptr))
    }

    fn set_current_operands(&mut self) {
        let current = self.current_opcode.as_ref().unwrap();

        match current.instruction {
            1 | 2 | 5 | 6 | 7 | 8 => {
                self.current_operands[0] = self.resolve_operand(self.read(self.instruction_ptr + 1), current.modes[0], false);
                self.current_operands[1] = self.resolve_operand(self.read(self.instruction_ptr + 2), current.modes[1], false);
                
                if current.instruction != 5 && current.instruction != 6 {
                    self.current_operands[2] = self.resolve_operand(self.read(self.instruction_ptr + 3), current.modes[2], true);
                }
            },
            3 => self.current_operands[0] = self.resolve_operand(self.read(self.instruction_ptr + 1), current.modes[0], true),
            4 | 9 => self.current_operands[0] = self.resolve_operand(self.read(self.instruction_ptr + 1), current.modes[0], false),
            _ => panic!("Unsupported opcode!")
        }
    }

    fn resolve_operand(&self, operand: isize, mode: OpcodeMode, write: bool) -> isize {
        match mode {
            OpcodeMode::Immediate => if !write { operand } else { panic!("Program attempted to resolve write operand with immediate mode!") },
            OpcodeMode::Position => if !write { self.read(operand as usize) } else { operand },
            OpcodeMode::Relative => if !write { self.read((operand + self.relative_base) as usize) } else { operand + self.relative_base },
        }
    }
}