2020_day08.rs

use std::collections::BTreeSet;

const INPUT: &str = include_str!("../input");

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum Opcode {
    Acc,
    Jmp,
    Nop,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
struct Instruction(Opcode, i32);

struct Cpu {
    acc: i32,
    ip: usize,
}

impl Cpu {
    fn new() -> Self {
        Self { acc: 0, ip: 0 }
    }

    fn step(&mut self, instruction: Instruction) {
        match instruction.0 {
            Opcode::Acc => self.acc = self.acc.wrapping_add(instruction.1),
            Opcode::Jmp => self.ip = self.ip.wrapping_add(instruction.1.wrapping_sub(1) as usize),
            Opcode::Nop => (),
        }

        self.ip = self.ip.wrapping_add(1);
    }

    fn has_loop(&mut self, instructions: &[Instruction]) -> bool {
        let mut set = BTreeSet::new();

        while !set.contains(&self.ip) {
            set.insert(self.ip);
            if let Some(i) = instructions.get(self.ip) {
                self.step(*i);
            } else {
                return false;
            }
        }

        true
    }

    fn detect_loop(&mut self, instructions: &[Instruction]) {
        let mut set = BTreeSet::new();

        while !set.contains(&self.ip) {
            set.insert(self.ip);
            self.step(instructions[self.ip]);
        }
    }

    fn fix_loop(original: &[Instruction]) -> Cpu {
        let mut last_modified = 0;

        loop {
            let mut cpu = Cpu::new();
            let mut modified = original.to_vec();
            let idx = modified[last_modified..]
                .iter()
                .position(|x| matches!(x.0, Opcode::Jmp | Opcode::Nop))
                .unwrap();
            match &mut modified[last_modified + idx].0 {
                v @ Opcode::Jmp => *v = Opcode::Nop,
                v @ Opcode::Nop => *v = Opcode::Jmp,
                _ => panic!("Noe Opcode"),
            }

            if !cpu.has_loop(&modified) {
                break cpu;
            }

            last_modified += idx + 1;
        }
    }
}

fn parse(s: &str) -> Vec<Instruction> {
    s.lines()
        .map(|line| {
            let mut splitted = line.split(' ');
            let ins = splitted.next().unwrap();
            let arg = splitted.next().unwrap();
            assert!(splitted.next().is_none());

            let parsed = match ins {
                "acc" => Opcode::Acc,
                "jmp" => Opcode::Jmp,
                "nop" => Opcode::Nop,
                _ => panic!("Unknown instruction"),
            };

            Instruction(parsed, arg.parse().unwrap())
        })
        .collect()
}

fn main() {
    let instructions = parse(INPUT);
    let mut cpu = Cpu::new();
    cpu.detect_loop(&instructions);

    println!("A: {}", cpu.acc);
    println!("B: {}", Cpu::fix_loop(&instructions).acc);
}

#[test]
fn test1() {
    let input = "nop +0
acc +1
jmp +4
acc +3
jmp -3
acc -99
acc +1
jmp -4
acc +6";

    let instructions = parse(input);
    let mut cpu = Cpu::new();
    cpu.detect_loop(&instructions);
    assert_eq!(cpu.acc, 5);
}

#[test]
fn test2() {
    let input = "nop +0
acc +1
jmp +4
acc +3
jmp -3
acc -99
acc +1
jmp -4
acc +6";

    let instructions = parse(input);
    let fixed_cpu = Cpu::fix_loop(&instructions);
    assert_eq!(fixed_cpu.acc, 8);
}