year2024/
day17.rs

1use std::ops::ControlFlow;
2use utils::prelude::*;
3
4/// Interpreting 3-bit assembly.
5///
6/// Part 2 assumes the input is structured such that the Nth digit in the output depends only on
7/// bits N*3 onwards. This enables working backwards from the right digit, checking 8 possible
8/// 3-bit patterns for each digit.
9#[derive(Clone, Debug)]
10pub struct Day17 {
11    a_start: u64,
12    program: Vec<u8>,
13}
14
15impl Day17 {
16    pub fn new(input: &str, _: InputType) -> Result<Self, InputError> {
17        let (a_start, b_start, c_start, program) = parser::u64()
18            .with_prefix("Register A: ")
19            .then(parser::u64().with_prefix("\nRegister B: "))
20            .then(parser::u64().with_prefix("\nRegister C: "))
21            .then(
22                parser::number_range(0..=7)
23                    .repeat(b',', 2)
24                    .with_prefix("\n\nProgram: "),
25            )
26            .parse_complete(input)?;
27
28        if b_start != 0 || c_start != 0 {
29            return Err(InputError::new(input, 0, "expected B and C to start at 0"));
30        }
31
32        Ok(Self { a_start, program })
33    }
34
35    #[must_use]
36    pub fn part1(&self) -> String {
37        let mut output = String::new();
38
39        self.run(self.a_start, |x| {
40            if !output.is_empty() {
41                output.push(',');
42            }
43            output.push((b'0' + x) as char);
44            ControlFlow::Continue(())
45        });
46
47        output
48    }
49
50    #[inline]
51    fn run(&self, mut a: u64, mut out_fn: impl FnMut(u8) -> ControlFlow<()>) {
52        let mut b = 0;
53        let mut c = 0;
54        let mut pc = 0;
55
56        while pc + 1 < self.program.len() {
57            let opcode = self.program[pc];
58            let operand = self.program[pc + 1] as u64;
59            let combo_operand = || match operand {
60                0..=3 => operand,
61                4 => a,
62                5 => b,
63                6 => c,
64                7 => panic!("combo operand 7 is reserved"),
65                _ => unreachable!(),
66            };
67
68            match opcode {
69                0 => a /= 1 << combo_operand(), // adv
70                1 => b ^= operand,              // bxl
71                2 => b = combo_operand() % 8,   // bst
72                3 => {
73                    // jnz
74                    if a != 0 {
75                        pc = operand as usize;
76                        continue;
77                    }
78                }
79                4 => b ^= c, // bxc
80                5 => {
81                    // out
82                    if out_fn((combo_operand() % 8) as u8) == ControlFlow::Break(()) {
83                        return;
84                    }
85                }
86                6 => b = a / (1 << combo_operand()), // bdv
87                7 => c = a / (1 << combo_operand()), // cdv
88                _ => unreachable!(),
89            }
90
91            pc += 2;
92        }
93    }
94
95    #[must_use]
96    pub fn part2(&self) -> u64 {
97        self.search(0, self.program.len() - 1)
98            .expect("no solution found")
99    }
100
101    fn search(&self, base: u64, pos: usize) -> Option<u64> {
102        if pos == 0 {
103            for x in 0..8 {
104                let a = base | x;
105                if self.check_all_matches(a) {
106                    return Some(a);
107                }
108            }
109            return None;
110        }
111
112        for x in 0..8 {
113            let a = base | (x << (pos * 3));
114            if self.check_pos_matches(a, pos) {
115                if let Some(result) = self.search(a, pos - 1) {
116                    return Some(result);
117                }
118            }
119        }
120        None
121    }
122
123    fn check_pos_matches(&self, a: u64, pos: usize) -> bool {
124        let mut output = None;
125        self.run(a / (1 << (pos * 3)), |out| {
126            output = Some(out);
127            ControlFlow::Break(())
128        });
129        output == Some(self.program[pos])
130    }
131
132    fn check_all_matches(&self, a: u64) -> bool {
133        let mut count = 0;
134        self.run(a, |out| {
135            if count >= self.program.len() || self.program[count] != out {
136                ControlFlow::Break(())
137            } else {
138                count += 1;
139                ControlFlow::Continue(())
140            }
141        });
142        count == self.program.len()
143    }
144}
145
146examples!(Day17 -> (&'static str, u64) [
147    {
148        input: "Register A: 729\n\
149            Register B: 0\n\
150            Register C: 0\n\
151            \n\
152            Program: 0,1,5,4,3,0",
153        part1: "4,6,3,5,6,3,5,2,1,0"
154    },
155    {
156        input: "Register A: 2024\n\
157            Register B: 0\n\
158            Register C: 0\n\
159            \n\
160            Program: 0,3,5,4,3,0",
161        part2: 117440,
162    },
163]);