Time for another Advent of Code 2020 problem!

That one sounds like it's going to be fun. Our input is pretty much assembly, like this:

nop +0
acc +1
jmp +4
acc +3
jmp -3
acc -99
acc +1
jmp -4
acc +6

So, the first thing we're going to do is write down some types.

There's more than one way to approach this problem, but let's go with this:

Rust code
#[derive(Debug, Clone, Copy)]
enum InstructionKind {
    Nop,
    Acc,
    Jmp,
}

#[derive(Debug, Clone, Copy)]
struct Instruction {
    kind: InstructionKind,
    operand: isize,
}

type Program = Vec<Instruction>;

Then let's write a very quick parser - we won't even reach for peg this time.

Rust code
fn parse_program(input: &str) -> Program {
    input
        .lines()
        .map(|l| {
            let mut tokens = l.split(' ');
            Instruction {
                kind: match tokens.next() {
                    Some(tok) => match tok {
                        "nop" => InstructionKind::Nop,
                        "acc" => InstructionKind::Acc,
                        "jmp" => InstructionKind::Jmp,
                        _ => panic!("unknown instruction kind {}", tok),
                    },
                    None => panic!("for line {}, expected instruction kind", l),
                },
                operand: match tokens.next() {
                    Some(tok) => tok.parse().unwrap(),
                    None => panic!("for line {}, expected operand", l),
                },
            }
        })
        .collect()
}
Rust code
fn main() {
    let program = parse_program(include_str!("input.txt"));
    dbg!(program);
}
Shell session
$ cargo run --quiet
[src/main.rs:42] program = [
    Instruction {
        kind: Nop,
        operand: 0,
    },
    Instruction {
        kind: Acc,
        operand: 1,
    },
    Instruction {
        kind: Jmp,
        operand: 4,
    },
    Instruction {
        kind: Acc,
        operand: 3,
    },
    Instruction {
        kind: Jmp,
        operand: -3,
    },
    Instruction {
        kind: Acc,
        operand: -99,
    },
    Instruction {
        kind: Acc,
        operand: 1,
    },
    Instruction {
        kind: Jmp,
        operand: -4,
    },
    Instruction {
        kind: Acc,
        operand: 6,
    },
]

Next step - we're emulating a machine, so what does our machine state look like?

We have an accumulator, and a program counter - they both start at 0, so we can derive the Default trait easily.

Rust code
#[derive(Debug, Clone, Copy, Default)]
struct State {
    /// Program counter
    pc: usize,
    /// Accumulator
    acc: isize,
}

From there, we can implement a .next() method on State, that takes a program, and returns the next State!

Our State is small, and Copy, so we can take it by value and return it by value. I don't feel like doing error handling today so we'll just panic! if we encounter an invalid instruction, like jumping before 0 or after the end of the program.

Rust code
use std::convert::TryInto;

impl State {
    fn next(self, program: &Program) -> Self {
        let ins = program[self.pc];
        match ins.kind {
            InstructionKind::Nop => Self {
                pc: self.pc + 1,
                ..self
            },
            InstructionKind::Acc => Self {
                pc: self.pc + 1,
                acc: self.acc + ins.operand,
            },
            InstructionKind::Jmp => Self {
                pc: (self.pc as isize + ins.operand).try_into().unwrap(),
                ..self
            },
        }
    }
}

And then we can give it a go!

Rust code
fn main() {
    let program = parse_program(include_str!("input.txt"));

    let mut state: State = Default::default();
    dbg!("initial state", state);

    for _ in 0..5 {
        println!("will execute {:?}", program[state.pc]);
        state = state.next(&program);
        dbg!(state);
    }
}
Shell session
$ cargo run --quiet
[src/main.rs:74] "initial state" = "initial state"
[src/main.rs:74] state = State {
    pc: 0,
    acc: 0,
}
will execute Instruction { kind: Nop, operand: 0 }
[src/main.rs:79] state = State {
    pc: 1,
    acc: 0,
}
will execute Instruction { kind: Acc, operand: 1 }
[src/main.rs:79] state = State {
    pc: 2,
    acc: 1,
}
will execute Instruction { kind: Jmp, operand: 4 }
[src/main.rs:79] state = State {
    pc: 6,
    acc: 1,
}
will execute Instruction { kind: Acc, operand: 1 }
[src/main.rs:79] state = State {
    pc: 7,
    acc: 2,
}
will execute Instruction { kind: Jmp, operand: -4 }
[src/main.rs:79] state = State {
    pc: 3,
    acc: 2,
}

The question we need to answer is: Immediately before any instruction is executed a second time, what value is in the accumulator?

We could do this the imperative way... or we could somehow shoehorn our problem into an Iterator problem, as is customary around here.

Rust code
fn main() {
    let program = parse_program(include_str!("input.txt"));

    let mut state: State = Default::default();

    let iter = std::iter::from_fn(|| {
        state = state.next(&program);
        Some(state)
    });

    dbg!(iter.take(5).collect::<Vec<_>>());
}
Shell session
$ cargo run --quiet
[src/main.rs:81] iter.take(5).collect::<Vec<_>>() = [
    State {
        pc: 1,
        acc: 0,
    },
    State {
        pc: 2,
        acc: 1,
    },
    State {
        pc: 6,
        acc: 1,
    },
    State {
        pc: 7,
        acc: 2,
    },
    State {
        pc: 3,
        acc: 2,
    },
]

There is, however, a nicer way to write this if we use itertools, so let's try it:

Shell session
$ cargo add itertools
      Adding itertools v0.9.0 to dependencies
Rust code
fn main() {
    let program = parse_program(include_str!("input.txt"));

    let iter = itertools::iterate(State::default(), |s| s.next(&program));
    dbg!(iter.take(5).collect::<Vec<_>>());
}

The output now includes the initial state as well:

Shell session
$ cargo run --quiet
[src/main.rs:74] iter.take(5).collect::<Vec<_>>() = [
    State {
        pc: 0,
        acc: 0,
    },
    State {
        pc: 1,
        acc: 0,
    },
    State {
        pc: 2,
        acc: 1,
    },
    State {
        pc: 6,
        acc: 1,
    },
    State {
        pc: 7,
        acc: 2,
    },
]

Now, since we need to determine when we run an instruction for the second time, I'd like to maintain a hashset of all the instructions' positions we have already executed.

Whenever HashSet::insert returns false (it did have this value present), we can stop and return what's in the accumulator.

Rust code
use std::collections::HashSet;

fn main() {
    let program = parse_program(include_str!("input.txt"));

    let mut iter = itertools::iterate(State::default(), |s| s.next(&program));
    let mut set: HashSet<usize> = Default::default();
    let answer = iter.find(|state| !set.insert(state.pc)).unwrap();

    println!(
        "Before executing {} a second time, the accumulator was {}",
        answer.pc, answer.acc
    );
}
Shell session
$ cargo run --quiet
Before executing 1 a second time, the accumulator was 5

Let's try it with the real input now:

Shell session
$ cargo run --quiet
Before executing 296 a second time, the accumulator was 1594

Part 2

In Part 2, we need to fix the program so it actually terminates - by turning one jmp into a nop, or one nop into a jmp.

Part 1 gave us which instruction was about to be executed a second time, so we need to figure out how we ended up there. But with our current setup, we only have the next instruction.

To get both the previous and the next instruction, we can use itertools' tuple_windows method:

Rust code
use itertools::Itertools;

fn main() {
    let program = parse_program(include_str!("input.txt"));

    let iter = itertools::iterate(State::default(), |s| s.next(&program));
    let mut set: HashSet<usize> = Default::default();
    let answer = iter
        .tuple_windows()
        .find(|(_, next)| !set.insert(next.pc))
        .unwrap();

    println!(
        "Before {:?}, we were at state {:?} and executed {:?}",
        answer.1, answer.0, program[answer.0.pc]
    );
}
Shell session
$ cargo run --quiet
Before State { pc: 296, acc: 1594 }, we were at state State { pc: 317, acc: 1594 } and executed Instruction { kind: Jmp, operand: -21 }

So, if my calculations are right... on line 318 (since lines are 1-based in most text editors), we should find...

jmp -21

But is that really what's causing an infinite loop? If we change it to nop -21, then we just loop somewhere else:

Shell session
$ cargo run --quiet
Before State { pc: 345, acc: 1546 }, we were at state State { pc: 322, acc: 1546 } and executed Instruction { kind: Jmp, operand: 23 }

And we're only allowed to change one instruction...

So, there's probably a smart way to solve this. But I'm tired, and there's only 656 instructions in this program...

Rust code
fn main() {
    let program = parse_program(include_str!("input.txt"));

    let num_jmp_and_nop = program
        .iter()
        .filter(|i| matches!(i.kind, InstructionKind::Jmp | InstructionKind::Nop))
        .count();
    dbg!(num_jmp_and_nop);
}
Shell session
$ [src/main.rs:85] num_jmp_and_nop = 291

...291 of which are jmp and nop. So here's an idea - how about we use brute force? How about we generate 291 different versions of our program, and simulate them all in parallel? First one to complete, wins the prize!

First we'll need a way to indicate completion, so we'll change State::next to return an Option instead:

Rust code
impl State {
    fn next(self, program: &Program) -> Option<Self> {
        if !(0..program.len()).contains(&self.pc) {
            return None;
        }

        let ins = program[self.pc];
        Some(match ins.kind {
            InstructionKind::Nop => Self {
                pc: self.pc + 1,
                ..self
            },
            InstructionKind::Acc => Self {
                pc: self.pc + 1,
                acc: self.acc + ins.operand,
            },
            InstructionKind::Jmp => Self {
                pc: (self.pc as isize + ins.operand).try_into().unwrap(),
                ..self
            },
        })
    }
}

And then it's off to the race:

Rust code
fn main() {
    let program = parse_program(include_str!("input.txt"));
    find_variant(&program);
}

fn flip_kind(kind: &mut InstructionKind) {
    *kind = match *kind {
        InstructionKind::Jmp => InstructionKind::Nop,
        InstructionKind::Nop => InstructionKind::Jmp,
        x => x,
    };
}

fn find_variant(program: &Program) {
    let mut variants: Vec<_> = program
        .iter()
        .enumerate()
        .filter_map(|(index, ins)| match ins.kind {
            InstructionKind::Jmp | InstructionKind::Nop => Some(index),
            _ => None,
        })
        .map(|i| {
            let mut variant = program.clone();
            flip_kind(&mut variant[i].kind);
            (i, variant)
        })
        .map(|(index, variant)| {
            itertools::iterate(Some(State::default()), move |state| {
                state
                    .unwrap_or_else(|| panic!("variant {} terminated!", index))
                    .next(&variant)
            })
        })
        .collect();

    loop {
        for v in &mut variants {
            v.next();
        }
    }
}
Shell session
$ cargo run --quiet
thread 'main' panicked at 'variant 364 terminated!', src/main.rs:99:40
note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace

Interesting! Now we can finally answer the question: what is the value of the accumulator when the program terminates?

Rust code
use itertools::Itertools;

fn eval(program: &Program) -> Option<isize> {
    itertools::iterate(Some(State::default()), |state| {
        state.and_then(|state| state.next(program))
    })
    .while_some()
    .last()
    .map(|s| s.acc)
}

fn main() {
    let mut program = parse_program(include_str!("input.txt"));
    flip_kind(&mut program[364].kind);
    dbg!(eval(&program));
}
Shell session
$ cargo run --quiet
[src/main.rs:79] eval(&program) = Some(
    758,
)

🎉🎉🎉!

Until next time, take care.