Day 17

src/bin/day17.rs

@@ -0,0 +1,248 @@
+use std::time::Instant;
+
+use aoc_2024::{load, print_res};
+use bstr::BString;
+use color_eyre::eyre::{bail, Context, OptionExt};
+use itertools::Itertools;
+
+#[derive(Debug, Clone, Copy)]
+pub struct Computer {
+    a: u64,
+    b: u64,
+    c: u64,
+
+    pc: usize,
+}
+
+type Parsed = (Computer, Vec<u8>);
+
+#[inline(never)]
+pub fn parsing(input: &BString) -> color_eyre::Result<Parsed> {
+    let Some((regs, program)) = std::str::from_utf8(input)?.split_once("\n\n") else {
+        bail!("Malformed input")
+    };
+
+    let mut regs = regs.lines();
+    let mut extract_reg = |name: char| -> color_eyre::Result<u64> {
+        regs.next()
+            .ok_or_eyre("Missing register")?
+            .strip_prefix("Register ")
+            .ok_or_eyre("Malformed prefix")?
+            .strip_prefix(name)
+            .ok_or_eyre("Malformed name")?
+            .strip_prefix(": ")
+            .ok_or_eyre("Malformed suffix")?
+            .parse()
+            .map_err(Into::into)
+    };
+
+    let program = program
+        .strip_prefix("Program: ")
+        .ok_or_eyre("Malformed program")?
+        .trim();
+
+    Ok((
+        Computer {
+            a: extract_reg('A').with_context(|| "Parsing A")?,
+            b: extract_reg('B').with_context(|| "Parsing B")?,
+            c: extract_reg('C').with_context(|| "Parsing C")?,
+
+            pc: 0,
+        },
+        program
+            .split(',')
+            .map(|n| {
+                color_eyre::eyre::ensure!(
+                    n.len() == 1 && (b'0'..=b'7').contains(&n.as_bytes()[0]),
+                    "Instruction: {n}"
+                );
+                Ok(n.as_bytes()[0] - b'0')
+            })
+            .collect::<Result<_, _>>()?,
+    ))
+}
+
+impl Computer {
+    fn combo(&self, v: u8) -> u64 {
+        match v {
+            0..=3 => v as u64,
+            4 => self.a,
+            5 => self.b,
+            6 => self.c,
+            invalid => unreachable!("Invalid combo: {invalid}"),
+        }
+    }
+
+    pub fn run(&mut self, code: &[u8]) -> Vec<u8> {
+        let mut output = Vec::new();
+        while self.pc < code.len() {
+            let operand = code[self.pc + 1];
+            match code[self.pc] {
+                // adv
+                0 => {
+                    self.a >>= self.combo(operand);
+                }
+                // bxl
+                1 => {
+                    self.b ^= operand as u64;
+                }
+                // bst
+                2 => {
+                    self.b = self.combo(operand) % 8;
+                }
+                // jnz
+                3 => {
+                    if self.a != 0 {
+                        self.pc = operand as usize;
+                        continue;
+                    }
+                }
+                // bxc
+                4 => {
+                    self.b ^= self.c;
+                }
+                // out
+                5 => {
+                    output.push((self.combo(operand) % 8) as u8);
+                }
+                // bdv
+                6 => {
+                    self.b = self.a >> self.combo(operand);
+                }
+                // cdv
+                7 => {
+                    self.c = self.a >> self.combo(operand);
+                }
+                invalid => unreachable!("Invalid opcode: {invalid}"),
+            }
+
+            self.pc += 2;
+        }
+
+        output
+    }
+}
+
+#[inline(never)]
+pub fn part1((mut computer, code): Parsed) {
+    print_res!(
+        "Result: {}",
+        computer.run(&code).iter().map(|s| s.to_string()).join(",")
+    );
+}
+
+fn display_code(code: &[u8]) {
+    fn combo_display(v: u8) {
+        match v {
+            0..=3 => print!("{v}"),
+            4 => print!("A"),
+            5 => print!("B"),
+            6 => print!("C"),
+            _ => print!("<INVALID>"),
+        }
+    }
+
+    println!("Program:");
+    for (i, chunk) in code.chunks_exact(2).enumerate() {
+        print!("  [{:2}] ", i * 2);
+        match chunk[0] {
+            0 => {
+                print!("A >>= ");
+                combo_display(chunk[1]);
+            }
+            1 => {
+                print!("B ^= {}", chunk[1]);
+            }
+            2 => {
+                print!("B = ");
+                combo_display(chunk[1]);
+                print!(" % 8");
+            }
+            3 => {
+                print!("A != 0 goto {}", chunk[1]);
+            }
+            4 => {
+                print!("B ^= C");
+            }
+            5 => {
+                print!("out(");
+                combo_display(chunk[1]);
+                print!(")");
+            }
+            6 => {
+                print!("B = A >> ");
+                combo_display(chunk[1]);
+            }
+            7 => {
+                print!("C = A >> ");
+                combo_display(chunk[1]);
+            }
+            _ => {
+                print!("<INVALID>")
+            }
+        }
+        println!()
+    }
+}
+
+#[inline(never)]
+pub fn part2((computer, code): Parsed) {
+    display_code(&code);
+
+    let mut possibles = vec![0u64];
+
+    for i in 0..code.len() {
+        let mut new = Vec::new();
+
+        for a in possibles {
+            for x in 0..=7 {
+                let tentative = a | x << (i * 3);
+                let mut computer = computer;
+                computer.a = tentative;
+
+                let result = computer.run(&code);
+                if i > 3 && (result.len() < i - 3 || result[0..(i - 3)] != code[0..(i - 3)]) {
+                    continue;
+                }
+
+                new.push(tentative);
+            }
+        }
+
+        possibles = new;
+    }
+
+    for tentative in possibles {
+        let mut computer = computer;
+        computer.a = tentative;
+
+        let result = computer.run(&code);
+        if result == code {
+            print_res!("Value: {tentative}");
+            return;
+        }
+    }
+
+    panic!("No solution found");
+}
+
+pub fn main() -> color_eyre::Result<()> {
+    let context = load()?;
+
+    let start = Instant::now();
+    let parsed = parsing(&context.input)?;
+    let elapsed = humantime::format_duration(start.elapsed());
+
+    let start = Instant::now();
+    if context.part == 1 {
+        part1(parsed);
+    } else {
+        part2(parsed);
+    }
+    let elapsed_part = humantime::format_duration(start.elapsed());
+
+    println!("  Parsing: {elapsed}");
+    println!("  Solving: {elapsed_part}");
+
+    Ok(())
+}