path: root/2019/rust/day09/src/main.rs

use std::env;
use std::io::{self, BufRead, BufReader};
use std::fs::File;

type BusWidth = i128;
type IOType = i128;
type IOData = Vec<IOType>;

fn read_program(filename: &str) -> Vec<BusWidth> {
    let f = File::open(filename).expect(filename);
    let f = BufReader::new(f);

    f.lines()
        .filter(|line| match line {
            Ok(s) => {
                let b = s.as_bytes();
                if b.len() == 0 || b.len() > 0 && b[0] == b'#' { false } else { true }
            },
            _ => true,
        })
        .map(|line| match line {
            Ok(line) => {
                let data:Vec<BusWidth> = line.split(',').map(|e| { e.trim_start().trim_end() })
                .filter(|v| match v.parse::<BusWidth>() { Ok(_) => true, _ => {
                            eprintln!("Could not parse: {:?}", v);
                        false
                    },
                })
                .map(|r| r.parse().unwrap())
                .collect();

                data
            },
            Err(_) => panic!("Could not parse {}.", filename),
        }).flatten().collect()
}

#[derive(Clone, Debug)]
struct Memory {
    mem: Vec<BusWidth>,
}

impl Memory {
    fn new() -> Memory {
        Memory { mem: Vec::new(), }
    }

    fn init(&mut self, mem:Vec<BusWidth>) {
        self.mem = mem;
    }

    fn read(&self, address: usize) -> BusWidth {
        match self.mem.get(address) {
            Some(value) => *value,
            None =>        0,
        }
    }

    fn write(&mut self, address: usize, value: BusWidth) {
        if self.len() < address {
            self.mem.resize(address*2, 0);
        }
        self.mem[address] = value;
    }

    fn len(&self) -> usize {
        self.mem.len()
    }
}

#[derive(Debug)]
struct Registers {
    ip: usize,
    base: usize,
}

#[derive(Debug)]
enum Arithm {
    Add,
    Mul,
    LessThan,
    Equals,
}

#[derive(Debug)]
enum IO {
    In,
    Out,
}

#[derive(Debug)]
enum Jmp {
    JmpTrue,
    JmpFalse,
}

#[derive(Debug)]
enum RegOp {
    BaseMod,
}

#[derive(Debug, PartialEq)]
enum ParameterMode {
    Position,
    Immediate,
    Relative,
}

#[derive(Debug, PartialEq)]
enum DataDirection {
    Read,
    Write,
}

#[derive(PartialEq)]
enum IOMethod {
    StdIO,
    Variables,
}

struct Intcode {
    r: Registers,
    io_method: IOMethod,
    input: IOData,
    output: IOData,
}

impl Intcode {
    fn new() -> Intcode {
        Intcode {
            r: Registers { ip: 0, base: 0, },
            io_method: IOMethod::Variables,
            input: vec![],
            output: vec![],
        }
    }

    fn param_to_argument(&self, mode: &ParameterMode, mem: &Memory, offset: usize,
                         dir: &DataDirection) -> usize
    {
        match mode {
            ParameterMode::Position =>  {
                if mem.read(offset) < 0 {
                    panic!("Attempt to address negative memory at IP: {}");
                }
                mem.read(offset) as usize
            },
            ParameterMode::Immediate => {
                if *dir == DataDirection::Write {
                    panic!("Immediate parameter mode at IP: {}", self.r.ip);
                }
                offset
            },
            ParameterMode::Relative => {
                if (mem.read(offset) + self.r.base as BusWidth) < 0 {
                    panic!("Attempt to address negative memory at IP: {}");
                }
                (mem.read(offset) + self.r.base as BusWidth) as usize
            },
        }
    }

    fn arithm(&self, op: Arithm, modes: Vec<ParameterMode>, memory: &mut Memory) {
        let parameter_directions =
            [DataDirection::Read, DataDirection::Read, DataDirection::Write];
        let params:Vec<usize> = (0..parameter_directions.len())
            .zip(parameter_directions.iter())
            .map(|(p, d)| {
                self.param_to_argument(&modes[p], memory, self.r.ip + 1 + p, d)
            }).collect();
        let left   = params[0];
        let right  = params[1];
        let result = params[2];

        match op {
            Arithm::Add =>      memory.write(result, memory.read(left) + memory.read(right)),
            Arithm::Mul =>      memory.write(result, memory.read(left) * memory.read(right)),
            Arithm::LessThan => memory.write(result, if memory.read(left) <  memory.read(right) {
                1 } else { 0 }),
            Arithm::Equals =>   memory.write(result, if memory.read(left) == memory.read(right) {
                1 } else { 0 }),
        }
    }

    fn jmp(&mut self, op: Jmp, modes: Vec<ParameterMode>, memory: &mut Memory)
    {
        let parameter_directions = [DataDirection::Read, DataDirection::Read];
        let params:Vec<usize> = (0..parameter_directions.len())
            .zip(parameter_directions.iter())
            .map(|(p, d)| {
                self.param_to_argument(&modes[p], memory, self.r.ip + 1 + p, d)
            }).collect();
        let value   = params[0];
        let address = params[1];

        self.r.ip = match op {
            Jmp::JmpTrue =>  (if memory.read(value) != 0 {
                memory.read(address) as usize } else { self.r.ip + 3 }),
            Jmp::JmpFalse => (if memory.read(value) == 0 {
                memory.read(address) as usize } else { self.r.ip + 3 }),
        };

    }

    fn regop(&mut self, op: RegOp, modes: Vec<ParameterMode>, memory: &mut Memory)
    {
        let parameter_directions = [DataDirection::Read];
        let params:Vec<usize> = (0..parameter_directions.len())
            .zip(parameter_directions.iter())
            .map(|(p, d)| {
                self.param_to_argument(&modes[p], memory, self.r.ip + 1 + p, d)
            }).collect();
        let value = memory.read(params[0]);

        match op {
            RegOp::BaseMod => {
                if value < 0 {
                    self.r.base -= (-1 * value) as usize;
                } else {
                    self.r.base += value as usize;
                }
            },
        };
    }

    fn io_in(&mut self, _op: IO, modes: Vec<ParameterMode>, memory:&mut Memory)
    {
        let parameter_directions = [DataDirection::Write];
        let params:Vec<usize> = (0..parameter_directions.len())
            .zip(parameter_directions.iter())
            .map(|(p, d)| {
                self.param_to_argument(&modes[p], memory, self.r.ip + 1 + p, d)
            }).collect();
        let address = params[0];
        let stdin = io::stdin();

        match &self.io_method {
            IOMethod::StdIO => {
                let mut value: Option<BusWidth> = None;
                while {
                    stdin.lock().lines().next().map(|s| {
                        match s {
                            Ok(v) => {
                                let parsed_string = v.trim_end().trim_start().parse::<BusWidth>();
                                match parsed_string {
                                    Ok(v) => { value = Some(v); value },
                                    _ => None,
                                };
                                value
                            },
                            _ => { None },
                        }
                    });
                    if value == None {
                        eprintln!("Could not parse input, please retry.");
                        true
                    } else {
                        false
                    }
                } {}
                match value {
                    Some(v) => memory.write(address, v),
                    _ => {},
               };
            },
            IOMethod::Variables => {
                memory.write(address, self.input.remove(0) as BusWidth);
            },
        }
    }

    fn io_out(&mut self, _op: IO, modes: Vec<ParameterMode>, memory:&mut Memory )
    {
        let parameter_directions = [DataDirection::Read];
        let params:Vec<usize> = (0..parameter_directions.len())
            .zip(parameter_directions.iter())
            .map(|(p, d)| {
                self.param_to_argument(&modes[p], memory, self.r.ip + 1 + p, d)
            }).collect();
        let address = params[0];

        let o = memory.read(address);
        println!("{}", o);
        match &self.io_method {
            IOMethod::Variables => self.output.push(o as IOType),
            _ => {},
        }
    }

    fn intcode(&mut self, mut memory:Memory)
    {
        while self.r.ip < memory.len() {
            let op = memory.read(self.r.ip) % 100;
            let modes:Vec<ParameterMode> = (2..=4)
                .map(|e| memory.read(self.r.ip) / 10i32.pow(e) as BusWidth % 10)
                .map(|n| match n {
                    0 => ParameterMode::Position,
                    1 => ParameterMode::Immediate,
                    2 => ParameterMode::Relative,
                    p => panic!("Invalid parameter mode {} at IP: {}", p, self.r.ip),
                }).collect();

            match op {
                1 => { self.arithm(Arithm::Add,      modes, &mut memory); self.r.ip += 4 },
                2 => { self.arithm(Arithm::Mul,      modes, &mut memory); self.r.ip += 4 },
                7 => { self.arithm(Arithm::LessThan, modes, &mut memory); self.r.ip += 4 },
                8 => { self.arithm(Arithm::Equals,   modes, &mut memory); self.r.ip += 4 },
                3 => { self.io_in( IO::In,           modes, &mut memory); self.r.ip += 2 },
                4 => { self.io_out(IO::Out,          modes, &mut memory); self.r.ip += 2 },
                5 => { self.jmp(  Jmp::JmpTrue,      modes, &mut memory); },
                6 => { self.jmp(  Jmp::JmpFalse,     modes, &mut memory); },
                9 => { self.regop(RegOp::BaseMod,    modes, &mut memory); self.r.ip += 2},
                99 => break,
                op => panic!("Invalid operation {} at IP: {}", op, self.r.ip),
            }
        }
    }
}

fn main() {
    let args: Vec<String> = env::args().collect();

    let mut opts = getopts::Options::new();
    opts.reqopt("m", "mode", "Mode of IO operations", "[arg|stdio]");
    opts.reqopt("p", "program", "File to load program from", "FILE");
    let matches = match opts.parse(&args[1..]) {
        Ok(m) => { m }
        Err(f) => { panic!(f.to_string()) }
    };
    let program:String = matches.opt_str("program").unwrap();
    let mode = matches.opt_str("mode").unwrap();
    let mut memory = Memory::new();
    memory.init(read_program(&program));
    let mut cpu = Intcode::new();

    if mode == "arg" {
        cpu.io_method = IOMethod::Variables;
        let input:IOData = matches.free[0].split(",").map(|v| {
            let i = v.trim_end().parse().unwrap();
            i
        }).collect();
        cpu.input = input;
        while cpu.input.len() > 1 {
            cpu.intcode(memory.clone());
        }
    } else if mode == "stdio" {
        cpu.io_method = IOMethod::StdIO;

        cpu.intcode(memory.clone());
    } else {
        panic!("Unknown mode: {}", mode);
    }
}