advent of code 2021:16

aoc16.rs

use itertools::Itertools;
use bitvec::prelude::*;


// BIT UTILITIES

// Creates an iterator of bits from a hexadecimal input file
// Does not need to load the whole file into memory
fn bit_stream_from_hex_input(filename: &'static str) -> impl Iterator<Item=bool> {
    std::fs::File::open(filename)
        .unwrap()
        .bytes()
        .map(|mb| mb.unwrap())
        .map(|byte| byte as char)
        .map(|char| char.to_digit(16))
        .take_while(|mc| mc.is_some())
        .map(|mc| mc.unwrap())
        .tuples()
        .map(|(msn, lsn)| (msn * 16 + lsn) as u8)
        .flat_map(|byte| BitVec::<Msb0, _>::from_element(byte).into_iter())
}

// Reads a Big-endian number from the bit stream with a certain number of binary digits.
fn read_number<B, N, const C: u8>(from_bit_stream: &mut B) -> Option<N> where 
    B: Iterator<Item=bool>,
    N: CanStoreBits<C>,
    N: core::ops::Shl<u8, Output=N>,
    N: core::ops::BitOrAssign<N>,
    N: From<bool>
{
    let mut ret = N::from(false);
    for i in (0..C).rev() {
        let bit = from_bit_stream.next()?;
        ret |= N::from(bit) << i;
    }
    return Some(ret);
}

// Helper trait (and macro) used to determine whether an output integer format
// can represent a given number of bits from a bit stream (C).
trait CanStoreBits<const C: u8> {}

macro_rules! can_store_up_to_bits {
    ($entry:ident $($tokens:tt)*) => {
        can_store_up_to_bits!{0, $entry $($tokens)*}
    };
    ($acc:expr, $entry:ident $token:tt $($tokens:tt)*) => {
        can_store_up_to_bits!{2*$acc, $entry $($tokens)*}
        can_store_up_to_bits!{2*$acc + 1, $entry $($tokens)*}
    };
    ($count:expr, $entry:ident) => {
        impl CanStoreBits<{$count + 1u8}> for $entry {}
    };
}

can_store_up_to_bits!(u8 ###);
can_store_up_to_bits!(u16 ####);
can_store_up_to_bits!(u32 #####);
can_store_up_to_bits!(usize #####);
can_store_up_to_bits!(u64 ######);


// PACKET PARSING

#[derive(Debug)]
struct Packet {
    version: u8,
    type_id: u8,
    bit_length: u16,
    content: PacketContent
}

#[derive(Debug)]
enum PacketContent {
    Literal(usize),
    Children(Vec<Packet>)
}

fn read_packet<B>(from_bit_stream: &mut B) -> Option<Packet> where B: Iterator<Item=bool> {
    let version = read_number::<_, _, 3>(from_bit_stream)?;
    let type_id = read_number::<_, _, 3>(from_bit_stream)?;
    let content_and_length = match type_id {
        4 => read_packet_literal(from_bit_stream)?,
        _ => read_child_packets(from_bit_stream)?
    };
    Some(Packet {
        version,
        type_id, 
        bit_length: content_and_length.1 + 6, 
        content: content_and_length.0
    })
}

fn read_packet_literal<B>(from_bit_stream: &mut B) -> Option<(PacketContent, u16)> 
    where B: Iterator<Item=bool> 
{
    let mut ret = 0;
    let mut bit_length = 5;
    while from_bit_stream.next()? == true {
        bit_length += 5;
        ret <<= 4;
        ret += read_number::<_, usize, 4>(from_bit_stream)?;
    }
    ret <<= 4;
    ret += read_number::<_, usize, 4>(from_bit_stream)?;
    return Some((PacketContent::Literal(ret), bit_length));
}

fn read_child_packets<B>(from_bit_stream: &mut B) -> Option<(PacketContent, u16)> 
    where B: Iterator<Item=bool> 
{
    if from_bit_stream.next()? {
        let count = read_number::<_, _, 11>(from_bit_stream)?;
        let children = read_child_packets_by_count(from_bit_stream, count)?;
        Some((children.0, children.1 + 12))
    } else {
        let count = read_number::<_, _, 15>(from_bit_stream)?;
        let children = read_child_packets_by_length(from_bit_stream, count)?;
        Some((children.0, children.1 + 16))
    }
}

fn read_child_packets_by_length<B>(from_bit_stream: &mut B, bit_length: u16) -> Option<(PacketContent, u16)>
    where B: Iterator<Item=bool> 
{
    let mut accumulated_length = 0u16;
    let mut children: Vec<Packet> = vec![];
    loop {
        let new_child_packet = read_packet(from_bit_stream)?;
        accumulated_length += new_child_packet.bit_length;
        children.push(new_child_packet);
        if bit_length == accumulated_length {
            return Some((PacketContent::Children(children), bit_length));
        } else if accumulated_length > bit_length {
            panic!("Invalid packet children. Found a split at {}, not {}", accumulated_length, bit_length);
        }
    }
}

fn read_child_packets_by_count<B>(from_bit_stream: &mut B, count: u16) -> Option<(PacketContent, u16)> 
    where B: Iterator<Item=bool> 
{
    let mut accumulated_length = 0u16;
    let mut children: Vec<Packet> = vec![];
    for _i in 0..count {
        let new_child_packet = read_packet(from_bit_stream)?;
        accumulated_length += new_child_packet.bit_length;
        children.push(new_child_packet);
    }
    return Some((PacketContent::Children(children), accumulated_length));
}


// PART ONE

fn sum_versions(packet: &Packet) -> usize {
    match &packet.content {
        PacketContent::Literal(_c) => packet.version as usize,
        PacketContent::Children(children) => {
            packet.version as usize + children.iter().map(|p| sum_versions(p)).sum::<usize>()
        }
    }
}


// PART TWO

const SUM: u8 = 0;
const PRODUCT: u8 = 1;
const MINIMUM: u8 = 2;
const MAXIMUM: u8 = 3;
const GREATER_THAN: u8 = 5;
const LESS_THAN: u8 = 6;
const EQUAL_TO: u8 = 7;

impl ToString for Packet {
    fn to_string(&self) -> String {
        match &self.content {
            PacketContent::Literal(lit) => lit.to_string(),
            PacketContent::Children(children) => {
                let child_strings: Vec<String> = children.iter().map(|p| p.to_string()).collect(); 
                match self.type_id {
                    SUM => format!("({})", child_strings.join(" + ")),
                    PRODUCT => format!("({})", child_strings.join(" * ")),
                    MINIMUM => format!("min({})", child_strings.join(", ")),
                    MAXIMUM => format!("max({})", child_strings.join(", ")),
                    GREATER_THAN => format!("({})", child_strings.join(" > ")),
                    LESS_THAN => format!("({})", child_strings.join(" < ")),
                    EQUAL_TO => format!("({})", child_strings.join(" == ")),
                    _ => panic!("Invalid type Id: {}", self.type_id)
                }
            }
        } 
    }
}

fn evaluate(packet: &Packet) -> usize {
    match &packet.content {
        PacketContent::Literal(lit) => *lit,
        PacketContent::Children(children) => {
            let mut evaluated_children = children.iter().map(|packet| evaluate(packet));
            match packet.type_id {
                SUM => sum(evaluated_children),
                PRODUCT => product(evaluated_children),
                MINIMUM => minimum(evaluated_children),
                MAXIMUM => maximum(evaluated_children),
                GREATER_THAN => greater_than(evaluated_children),
                LESS_THAN => less_than(evaluated_children),
                EQUAL_TO => equal_to(evaluated_children),
                _ => panic!("Invalid type Id: {}", packet.type_id)
            }
        }
    }
}

fn sum<I>(mut evaluated_children: I) -> usize where I: Iterator<Item=usize> {
    evaluated_children.sum()
}

fn product<I>(mut evaluated_children: I) -> usize where I: Iterator<Item=usize> {
    evaluated_children.fold(1usize, |acc, child| acc * child)
}

fn minimum<I>(mut evaluated_children: I) -> usize where I: Iterator<Item=usize> {
    evaluated_children.fold(usize::MAX, |acc, child| std::cmp::min(acc, child))
}

fn maximum<I>(mut evaluated_children: I) -> usize where I: Iterator<Item=usize> {
    evaluated_children.fold(0usize, |acc, child| std::cmp::max(acc, child))
}

fn greater_than<I>(mut evaluated_children: I) -> usize where I: Iterator<Item=usize> {
   if evaluated_children.next().unwrap() > evaluated_children.next().unwrap() {1} else {0}
}

fn less_than<I>(mut evaluated_children: I) -> usize where I: Iterator<Item=usize> {
   if evaluated_children.next().unwrap() < evaluated_children.next().unwrap() {1} else {0}
}

fn equal_to<I>(mut evaluated_children: I) -> usize where I: Iterator<Item=usize> {
   if evaluated_children.next().unwrap() == evaluated_children.next().unwrap() {1} else {0}
}


// MAIN

fn main() {
    let top_level_packet = read_packet(&mut bit_stream_from_hex_input("16.txt")).unwrap();
    println!("Part One: {}", sum_versions(&top_level_packet));
    println!("Packet Expression: '{}'", top_level_packet.to_string());
    println!("Part Two: {}", evaluate(&top_level_packet));
}