AoC21/src/day16.rs

use std::str::Chars;

use itertools::Itertools;

pub fn day16() {
    assert_eq!(
        parse_hex_str("D2FE28"),
        Packet {
            version: 6,
            type_id: 4,
            length: 21,
            payload: Payload::Literal(2021),
        },
    );
    assert_eq!(
        parse_hex_str("38006F45291200"),
        Packet {
            version: 1,
            type_id: 6,
            length: 49,
            payload: Payload::Operator(vec![
                Packet {
                    version: 6,
                    type_id: 4,
                    length: 11,
                    payload: Payload::Literal(10),
                },
                Packet {
                    version: 2,
                    type_id: 4,
                    length: 16,
                    payload: Payload::Literal(20),
                },
            ]),
        },
    );
    dbg!(parse_hex_str("EE00D40C823060"));

    let sums = &[
        "8A004A801A8002F478",
        "620080001611562C8802118E34",
        "C0015000016115A2E0802F182340",
        "A0016C880162017C3686B18A3D4780",
    ]
    .map(|s| parse_hex_str(s).sum_versions());
    assert_eq!(sums, &[16, 12, 23, 31]);

    let input = include_str!("../input/day16.txt");
    let packet = parse_hex_str(input);
    println!("part 1: {}", packet.sum_versions());

    let values = &[
        "C200B40A82",
        "04005AC33890",
        "880086C3E88112",
        "CE00C43D881120",
        "D8005AC2A8F0",
        "F600BC2D8F",
        "9C005AC2F8F0",
        "9C0141080250320F1802104A08",
    ]
    .map(|s| parse_hex_str(s).eval());
    assert_eq!(values, &[3, 54, 7, 9, 1, 0, 0, 1]);

    println!("part 2: {}", packet.eval());
}

fn parse_hex_str(hex_str: &str) -> Packet {
    let mut bits = HexToBits::new(hex_str.chars());
    parse_packet(&mut bits, 0)
}

fn parse_packet<'a>(bits: &mut HexToBits<'a>, depth: u32) -> Packet {
    let mut packet_len = 6;
    let version = bits.read(3);
    let type_id = bits.read(3);
    let payload: Payload;
    // println!("type id {} at depth {}", type_id, depth);
    if type_id == 4 {
        let mut val = 0;
        loop {
            let is_last = !bits.next().unwrap();
            val <<= 4;
            val |= bits.read(4);
            packet_len += 5;
            if is_last {
                break;
            }
        }
        payload = Payload::Literal(val);
    } else {
        packet_len += 1;
        let length_type = bits.next().unwrap();
        let length = if length_type {
            packet_len += 11;
            bits.read(11)
        } else {
            packet_len += 15;
            bits.read(15)
        };
        // println!("length type: {}, length: {}", length_type, length);
        let mut sub_packets = vec![];
        let mut sub_packets_len = 0;
        loop {
            if length_type && sub_packets.len() >= length {
                break;
            } else if !length_type && sub_packets_len >= length {
                break;
            }
            let packet = parse_packet(bits, depth + 1);
            sub_packets_len += packet.length;
            // println!(
            //     "read packet of length {} at depth {}",
            //     packet.length,
            //     depth + 1
            // );
            sub_packets.push(packet);
        }
        payload = Payload::Operator(sub_packets);
        packet_len += sub_packets_len;
    }
    Packet {
        version,
        type_id,
        length: packet_len,
        payload,
    }
}

struct HexToBits<'a> {
    hex_chars: Chars<'a>,
    current: u32,
    current_bits_remaining: u32,
    read_so_far: usize,
}

impl<'a> HexToBits<'a> {
    fn new(chars: Chars<'a>) -> Self {
        Self {
            hex_chars: chars,
            current: 0,
            current_bits_remaining: 0,
            read_so_far: 0,
        }
    }

    fn read(&mut self, len: u32) -> usize {
        let mut res = 0;
        for _ in 0..len {
            res <<= 1;
            res |= if self.next().unwrap() { 1 } else { 0 };
        }
        res
    }
}

impl<'a> Iterator for HexToBits<'a> {
    type Item = bool;

    fn next(&mut self) -> Option<bool> {
        if self.current_bits_remaining > 0 {
            let res = Some((self.current >> self.current_bits_remaining - 1) & 1 == 1);
            self.current_bits_remaining -= 1;
            self.read_so_far += 1;
            res
        } else if let Some(c) = self.hex_chars.next() {
            self.current = c.to_digit(16).unwrap();
            self.current_bits_remaining = 3;
            self.read_so_far += 1;
            Some((self.current >> 3) & 1 == 1)
        } else {
            None
        }
    }
}

#[derive(Debug, PartialEq)]
struct Packet {
    version: usize,
    type_id: usize,
    length: usize,
    payload: Payload,
}

impl Packet {
    fn sum_versions(&self) -> usize {
        match self.payload {
            Payload::Literal(_) => self.version,
            Payload::Operator(ref sub) => {
                self.version + sub.iter().map(|p| p.sum_versions()).sum::<usize>()
            }
        }
    }

    fn eval(&self) -> usize {
        match self.payload {
            Payload::Literal(n) => n,
            Payload::Operator(ref sub) => {
                let sub_vals = sub.iter().map(|p| p.eval());
                match self.type_id {
                    0 => sub_vals.sum(),
                    1 => sub_vals.product(),
                    2 => sub_vals.min().unwrap(),
                    3 => sub_vals.max().unwrap(),
                    5 => {
                        let (a, b) = sub_vals.collect_tuple().unwrap();
                        if a > b {
                            1
                        } else {
                            0
                        }
                    }
                    6 => {
                        let (a, b) = sub_vals.collect_tuple().unwrap();
                        if a < b {
                            1
                        } else {
                            0
                        }
                    }
                    7 => {
                        let (a, b) = sub_vals.collect_tuple().unwrap();
                        if a == b {
                            1
                        } else {
                            0
                        }
                    }
                    _ => panic!(),
                }
            }
        }
    }
}

#[derive(Debug, PartialEq)]
enum Payload {
    Literal(usize),
    Operator(Vec<Packet>),
}