AOC 2019 Day 2 (First solution only)

use std::io::{self, BufRead};

#[derive(Debug)]
enum EmulationError {
    UnknownOpcode(usize),
    MemoryAccess(usize),
    Overflow,
}

fn produce_program(s: &str) -> Vec<usize> {
    s.split(',').flat_map(|v| v.parse()).collect()
}

fn execute_program(program: &mut [usize]) -> Result<&mut [usize], EmulationError> {
    for i in (0..program.len()).step_by(4) {
        match program[i] {
            // direct indexing is fine because we are always in bound due to the for loop
            99 => break,
            1 => binary_op(program, (i + 1, i + 2, i + 3), |x, y| {
                x.checked_add(y).ok_or(EmulationError::Overflow)
            })?,
            2 => binary_op(program, (i + 1, i + 2, i + 3), |x, y| {
                x.checked_mul(y).ok_or(EmulationError::Overflow)
            })?,
            x => return Err(EmulationError::UnknownOpcode(x)),
        }
    }
    Ok(program)
}

fn read_mem(program: &mut [usize], pos: usize) -> Result<usize, EmulationError> {
    program
        .get(pos)
        .copied()
        .ok_or(EmulationError::MemoryAccess(pos))
}

fn write_mem(program: &mut [usize], pos: usize, val: usize) -> Result<(), EmulationError> {
    *program
        .get_mut(pos)
        .ok_or(EmulationError::MemoryAccess(pos))? = val;
    Ok(())
}

fn binary_op<F>(
    program: &mut [usize],
    (i1, i2, o): (usize, usize, usize),
    mut f: F,
) -> Result<(), EmulationError>
where
    F: FnMut(usize, usize) -> Result<usize, EmulationError>,
{
    let (i1, i2, o) = (
        read_mem(program, i1)?,
        read_mem(program, i2)?,
        read_mem(program, o)?,
    );
    let val = f(read_mem(program, i1)?, read_mem(program, i2)?)?;
    write_mem(program, o, val)
}

fn main() {
    assert_eq!(
        execute_program(&mut [1, 0, 0, 0, 99]).unwrap(),
        &[2, 0, 0, 0, 99]
    );

    assert_eq!(
        execute_program(&mut [2, 3, 0, 3, 99]).unwrap(),
        &[2, 3, 0, 6, 99]
    );

    assert_eq!(
        execute_program(&mut [2, 4, 4, 5, 99, 0]).unwrap(),
        &[2, 4, 4, 5, 99, 9801]
    );

    assert_eq!(
        execute_program(&mut [1, 1, 1, 4, 99, 5, 6, 0, 99]).unwrap(),
        &[30, 1, 1, 4, 2, 5, 6, 0, 99]
    );

    let i = io::stdin();
    let mut stdin = i.lock();
    let mut input = String::new();
    stdin.read_line(&mut input).expect("Error reading stdin.");
    let mut data = produce_program(&input);
    // before running the program, replace position 1 with the value 12 and replace position 2 with the value 2.
    data[1] = 12;
    data[2] = 2;

    println!(
        "The first part has answer: {}",
        execute_program(&mut data.clone()).unwrap()[0]
    );
}