Advent of Code 2021 - Day 24: Arithmetic Logic Unit

arithmetic_logic_unit.rs

fn main() {
    let program = Program::from(MONAD);

    let (largest, smallest) = find_max_and_min_model_numbers(&program);

    println!("The largest model number is: {}", largest);
    println!("The smallest model number is: {}", smallest);
}

fn find_max_and_min_model_numbers(program: &Program) -> (i64, i64) {
    let function_maps = generate_function_maps(program);

    (
        find_model_number(Criterion::Largest, &function_maps),
        find_model_number(Criterion::Smallest, &function_maps),
    )
}

fn generate_function_maps(program: &Program) -> Vec<HashMap<i64, HashSet<(i64, i64)>>> {
    let mut function_maps = Vec::new();

    let functions = program.instructions.chunks(18).rev();
    let mut acceptable_inputs = HashSet::new();
    let mut acceptable_outputs = HashSet::from([0]);

    for function in functions {
        let program = Program {
            instructions: Vec::from(function),
        };

        let mut function_map = HashMap::new();
        let mut z = 0;

        'search: loop {
            for digit in [1, 2, 3, 4, 5, 6, 7, 8, 9] {
                let inputs = [digit];

                let mut alu = Alu::new(&inputs);
                alu.z = z;

                alu.execute(&program);

                if acceptable_outputs.contains(&alu.z) {
                    acceptable_inputs.insert(z);
                    function_map
                        .entry(z)
                        .or_insert_with(HashSet::new)
                        .insert((digit, alu.z));
                }
            }

            z += 1;

            if z > 250_000 {
                break 'search;
            }
        }

        function_maps.push(function_map);

        acceptable_outputs.clear();
        swap(&mut acceptable_outputs, &mut acceptable_inputs);
    }

    function_maps.reverse();

    function_maps
}

fn find_model_number(
    criterion: Criterion,
    function_maps: &[HashMap<i64, HashSet<(i64, i64)>>],
) -> i64 {
    let compare_fn = match criterion {
        Criterion::Largest => |a: &&(i64, i64), b: &&(i64, i64)| a.0.cmp(&b.0),
        Criterion::Smallest => |a: &&(i64, i64), b: &&(i64, i64)| b.0.cmp(&a.0),
    };

    let mut model_number = 0;
    let mut z = 0;

    for map in function_maps {
        let (max_digit, new_z) = map[&z].iter().max_by(compare_fn).unwrap();
        model_number *= 10;
        model_number += max_digit;
        z = *new_z;
    }

    model_number
}

enum Criterion {
    Largest,
    Smallest,
}

#[derive(Copy, Clone, Debug, PartialEq)]
enum Variable {
    W,
    X,
    Y,
    Z,
}

impl Variable {
    fn get(&self, alu: &Alu) -> i64 {
        alu.get(*self)
    }
}

impl From<&str> for Variable {
    fn from(value: &str) -> Self {
        match value {
            "w" => Self::W,
            "x" => Self::X,
            "y" => Self::Y,
            "z" => Self::Z,
            _ => panic!("Invalid variable: {}", value),
        }
    }
}

#[derive(Copy, Clone, Debug, PartialEq)]
enum Operand {
    Variable(Variable),
    Constant(i64),
}

impl From<&str> for Operand {
    fn from(value: &str) -> Self {
        if let Ok(number) = value.parse::<i64>() {
            Operand::Constant(number)
        } else {
            Operand::Variable(Variable::from(value))
        }
    }
}

impl Operand {
    fn get(&self, alu: &Alu) -> i64 {
        match self {
            Operand::Variable(v) => alu.get(*v),
            Operand::Constant(c) => *c,
        }
    }
}

#[derive(Copy, Clone, Debug)]
enum Instruction {
    Inp(Operand),
    Add(Operand, Operand),
    Mul(Operand, Operand),
    Div(Operand, Operand),
    Mod(Operand, Operand),
    Eql(Operand, Operand),
}

impl From<&str> for Instruction {
    fn from(value: &str) -> Self {
        let tokens = value.split(' ').collect::<Vec<_>>();

        match tokens[0] {
            "inp" => Instruction::Inp(Operand::from(tokens[1])),
            "add" => Instruction::Add(Operand::from(tokens[1]), Operand::from(tokens[2])),
            "mul" => Instruction::Mul(Operand::from(tokens[1]), Operand::from(tokens[2])),
            "div" => Instruction::Div(Operand::from(tokens[1]), Operand::from(tokens[2])),
            "mod" => Instruction::Mod(Operand::from(tokens[1]), Operand::from(tokens[2])),
            "eql" => Instruction::Eql(Operand::from(tokens[1]), Operand::from(tokens[2])),
            _ => panic!("Invalid instruction: {}", tokens[0]),
        }
    }
}

impl Instruction {
    fn execute(&self, alu: &mut Alu) {
        match self {
            Instruction::Inp(Operand::Variable(a)) => {
                let input = alu.get_input();
                alu.set(*a, input);
            }
            Instruction::Add(Operand::Variable(a), b) => alu.set(*a, a.get(alu) + b.get(alu)),
            Instruction::Mul(Operand::Variable(a), b) => alu.set(*a, a.get(alu) * b.get(alu)),
            Instruction::Div(Operand::Variable(a), b) => alu.set(*a, a.get(alu) / b.get(alu)),
            Instruction::Mod(Operand::Variable(a), b) => alu.set(*a, a.get(alu) % b.get(alu)),
            Instruction::Eql(Operand::Variable(a), b) => {
                alu.set(*a, (a.get(alu) == b.get(alu)) as i64)
            }
            _ => panic!("Invalid operands!"),
        }
    }
}

struct Program {
    instructions: Vec<Instruction>,
}

impl From<&str> for Program {
    fn from(value: &str) -> Self {
        let instructions = value.lines().map(Instruction::from).collect();

        Program { instructions }
    }
}

#[derive(Debug)]
struct Alu<'a> {
    w: i64,
    x: i64,
    y: i64,
    z: i64,
    ic: usize,
    inputs: &'a [i64],
}

impl<'a> Alu<'a> {
    fn new(inputs: &'a [i64]) -> Self {
        Self {
            w: 0,
            x: 0,
            y: 0,
            z: 0,
            ic: 0,
            inputs,
        }
    }

    fn execute(&mut self, program: &Program) {
        for instruction in &program.instructions {
            instruction.execute(self);
        }
    }

    fn get_input(&mut self) -> i64 {
        let value = self.inputs[self.ic];
        self.ic += 1;

        value
    }

    fn get(&self, variable: Variable) -> i64 {
        match variable {
            Variable::W => self.w,
            Variable::X => self.x,
            Variable::Y => self.y,
            Variable::Z => self.z,
        }
    }

    fn set(&mut self, variable: Variable, value: i64) {
        match variable {
            Variable::W => self.w = value,
            Variable::X => self.x = value,
            Variable::Y => self.y = value,
            Variable::Z => self.z = value,
        }
    }
}