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