Compiler Explorer

Source code

//! Logic for parsing and computing the "packet" language in day 16 of Advent of Code 2021
#![no_std]
use core::cmp;
use core::ops::{Add, Mul};

pub fn begin(input: &str) -> usize {
    let bits = input
        .trim()
        .bytes()
        .map(|b| Nibble::from_hex_ascii(b).unwrap())
        .flat_map(Nibble::into_bits);

solve(bits).unwrap()
}

fn solve(bits: impl Iterator<Item = bool>) -> Result<usize, ComputeError> {
 compute(&mut bits.counted())
}

/// Returns either the solution to the given packet or an error.
///
/// # Arguments
/// * `bits` - a mutable reference to an [Iterator] over the bits of the packet.
fn compute(bits: &mut CountIter) -> Result<usize, ComputeError>
where
 I: Iterator<Item = bool>,
{
 let (_version, operation) = get_header(bits)?;

use Operation as Op;
 match operation {
 Op::Sum => reduce(Add::add, bits),
 Op::Product => reduce(Mul::mul, bits),
 Op::Minimum => reduce(cmp::min, bits),
 Op::Maximim => reduce(cmp::max, bits),
 Op::Literal => literal(bits),
 Op::Greater => compare(|a, b| a > b, bits),
 Op::Less => compare(|a, b| a < b, bits),
 Op::Equal => compare(|a, b| a == b, bits),
 }
}

fn literal(bits: &mut CountIter) -> Result<usize, ComputeError>
where
 I: Iterator<Item = bool>,
{
 Ok(usize::from_bits(&mut LiteralBits::new(bits)))
}

#[inline]
fn reduce(f: fn(usize, usize) -> usize, bits: &mut CountIter) -> Result<usize, ComputeError>
where
 I: Iterator<Item = bool>,
{
 let length_type_id = get_length_type(bits)?;
 match length_type_id {
 false => reduce_t0(f, bits),
 true => reduce_t1(f, bits),
 }
}

fn reduce_t0(
 f: fn(usize, usize) -> usize,
 bits: &mut CountIter,
) -> Result<usize, ComputeError>
where
 I: Iterator<Item = bool>,
{
 let num_bits = get_length_t0(bits)?;

let final_bits_read = bits.iter_count() + num_bits as usize;

let mut accum = compute(bits)?;

while bits.iter_count() != final_bits_read {
        let subpacket = compute(bits)?;
        accum = f(accum, subpacket);
    }

Ok(accum)
}

fn reduce_t1(
 f: fn(usize, usize) -> usize,
 bits: &mut CountIter,
) -> Result<usize, ComputeError>
where
 I: Iterator<Item = bool>,
{
 let num_packets = get_length_t1(bits)?;

let mut accum = compute(bits)?;
    let num_packets = num_packets - 1;

for _ in 0..num_packets {
        let subpacket = compute(bits)?;
        accum = f(accum, subpacket);
    }

Ok(accum)
}

#[inline]
fn compare(f: fn(usize, usize) -> bool, bits: &mut CountIter) -> Result<usize, ComputeError>
where
 I: Iterator<Item = bool>,
{
 let length_type_id = get_length_type(bits)?;
 match length_type_id {
 false => compare_t0(f, bits),
 true => compare_t1(f, bits),
 }
}

fn compare_t0(
 f: fn(usize, usize) -> bool,
 bits: &mut CountIter,
) -> Result<usize, ComputeError>
where
 I: Iterator<Item = bool>,
{
 let num_bits = get_length_t0(bits)?;

let bits_read = bits.iter_count();

let first = compute(bits)?;
    let second = compute(bits)?;

let packet_bits_read = bits.iter_count() - bits_read;

if packet_bits_read != num_bits as usize {
        Err("Comparison operation length field did not match exactly two subpackets")
    } else {
        Ok(f(first, second) as usize)
    }
}

fn compare_t1(
 f: fn(usize, usize) -> bool,
 bits: &mut CountIter,
) -> Result<usize, ComputeError>
where
 I: Iterator<Item = bool>,
{
 let num_packets = get_length_t1(bits)?;
 if num_packets != 2 {
 Err("Comparison operation length field did not match exactly two subpackets")
 } else {
 let first = compute(bits)?;
 let second = compute(bits)?;

Ok(f(first, second) as usize)
    }
}

#[inline]
fn get_header(bits: &mut CountIter) -> Result<(u8, Operation), ComputeError>
where
 I: Iterator<Item = bool>,
{
 const VERSION_FIELD_SIZE: usize = 3;
 const TYPE_ID_FIELD_SIZE: usize = 3;
 let version = try_read_bits(bits, VERSION_FIELD_SIZE)
 .ok_or("Expected packet version, but bit stream ended")?;
 let operation = try_read_bits::<u8, _>(bits, TYPE_ID_FIELD_SIZE)
 .ok_or("Expected packet type ID, but bit stream ended")?
 .try_into()?;

Ok((version, operation))
}

#[inline]
fn get_length_type(bits: &mut CountIter) -> Result<bool, ComputeError>
where
 I: Iterator<Item = bool>,
{
 //const LENGTH_TYPE_FIELD_SIZE: usize = 1;
 bits.next()
 .ok_or("Expected length type ID, but bit stream ended")
}

#[inline]
fn get_length_t0(bits: &mut CountIter) -> Result<u16, ComputeError>
where
 I: Iterator<Item = bool>,
{
 const T0_LEN_FIELD_SIZE: usize = 15;
 try_read_bits(bits, T0_LEN_FIELD_SIZE).ok_or("Expected type 0 length, but bit stream ended")
}

#[inline]
fn get_length_t1(bits: &mut CountIter) -> Result<u16, ComputeError>
where
 I: Iterator<Item = bool>,
{
 const T1_LEN_FIELD_SIZE: usize = 11;
 try_read_bits(bits, T1_LEN_FIELD_SIZE).ok_or("Expected type 1 length, but bit stream ended")
}

#[inline(always)]
fn try_read_bits<T, I>(bits: &mut CountIter, n: usize) -> Option<T>
where
 I: Iterator<Item = bool>,
 T: FromBits,
{
 let bits_read = bits.iter_count();
 let mut number_bits = bits.take(n);
 let number = T::from_bits(&mut number_bits);
 let bits_read = bits.iter_count() - bits_read;
 if bits_read == n {
 Some(number)
 } else {
 None
 }
}

type ComputeError = &'static str;

enum Operation {
    Sum,
    Product,
    Minimum,
    Maximim,
    Literal,
    Greater,
    Less,
    Equal,
}

impl TryFrom<u8> for Operation {
 type Error = ComputeError;

#[inline]
 fn try_from(value: u8) -> Result<Self, Self::Error> {
 match value {
 0 => Ok(Self::Sum),
 1 => Ok(Self::Product),
 2 => Ok(Self::Minimum),
 3 => Ok(Self::Maximim),
 4 => Ok(Self::Literal),
 5 => Ok(Self::Greater),
 6 => Ok(Self::Less),
 7 => Ok(Self::Equal),
 _ => Err("Unrecognized operation type"),
 }
 }
}

struct LiteralBits<'a, I> {
 inner: &'a mut I,
 state: u8,
 last: bool,
}

impl<'a, I> LiteralBits<'a, I> {
 #[inline]
 fn new(inner: &'a mut I) -> Self {
 LiteralBits {
 inner,
 state: 0,
 last: false,
 }
 }
}

impl<'a, I> Iterator for LiteralBits<'a, I>
where
 I: Iterator<Item = bool>,
{
 type Item = bool;

#[inline]
 fn next(&mut self) -> Option<Self::Item> {
 if self.state == 0 {
 if self.last {
 return None;
 } else {
 let group_id = self.inner.next()?;
 if !group_id {
 self.last = true;
 }
 self.state = 4;
 }
 }
 self.state -= 1;
 self.inner.next()
 }
}

trait FromBits: From<bool> {
 fn from_bits(bits: I) -> Self
 where
 I: Iterator<Item = bool>;
}

impl FromBits for u8 {
 #[inline]
 fn from_bits(bits: I) -> Self
 where
 I: Iterator<Item = bool>,
 {
 bits.fold(0, |acc, b| (acc << 1) | b as u8)
 }
}
impl FromBits for u16 {
 #[inline]
 fn from_bits(bits: I) -> Self
 where
 I: Iterator<Item = bool>,
 {
 bits.fold(0, |acc, b| (acc << 1) | b as u16)

}
}
impl FromBits for usize {
 #[inline]
 fn from_bits(bits: I) -> Self
 where
 I: Iterator<Item = bool>,
 {
 bits.fold(0, |acc, b| (acc << 1) | b as usize)
 }

}

use core::iter::FusedIterator;

/// A Nibble is used to represent a 4 bit value. It is half of a byte.
struct Nibble(u8);

impl Nibble {
 /// Returns a nibble from a [u8].
 ///
 /// # Arguments
 /// * `hex` - A u8 value interpreted as an ASCII character
 #[inline]
 const fn from_hex_ascii(hex: u8) -> Result<Self, &'static str> {
 match hex {
 c @ b'0'..=b'9' => Ok(Nibble(c - b'0')),
 c @ b'A'..=b'F' => Ok(Nibble(c - b'A' + 10)),
 _ => Err("Can't parse non-hexadecimal character"),
 }
 }

/// Returns an [Iterator] which iterates over each bit.
    /// Iterates from the most significant bit to the least significant bit
    #[inline]
    fn into_bits(self) -> NibbleBits {
        NibbleBits::new(self)
    }
}

/// A struct used to iterate over the bits of a [Nibble].
/// See [Nibble::into_bits]
struct NibbleBits {
    nibble: Nibble,
    len: u8,
}

impl NibbleBits {
    #[inline]
    fn new(nibble: Nibble) -> Self {
        Self { len: 4, nibble }
    }

#[inline]
 fn shift(&mut self) {
 self.nibble.0 <<= 1;
 self.len -= 1;
 }
}

impl Iterator for NibbleBits {
    type Item = bool;

#[inline]
 fn next(&mut self) -> Option<Self::Item> {
 const MASK: u8 = 0b00010000;
 if self.len > 0 {
 self.shift();
 Some((self.nibble.0 & MASK) != 0)
 } else {
 None
 }
 }

#[inline]
 fn size_hint(&self) -> (usize, Option<usize>) {
 (self.len as usize, Some(self.len as usize))
 }
}

impl ExactSizeIterator for NibbleBits {}

impl FusedIterator for NibbleBits {}

pub struct CountIter {
 inner: I,
 count: usize,
}

impl CountIter {
 #[inline]
 pub fn new(inner: I) -> CountIter {
 CountIter { inner, count: 0 }
 }

#[inline]
    pub fn iter_count(&self) -> usize {
        self.count
    }
}

impl Iterator for CountIter
where
 I: Iterator,
{
 type Item = ::Item;

#[inline]
 fn next(&mut self) -> Option<Self::Item> {
 self.count += 1;
 self.inner.next()
 }

#[inline]
 fn size_hint(&self) -> (usize, Option<usize>) {
 self.inner.size_hint()
 }
}

impl<I: ExactSizeIterator> ExactSizeIterator for CountIter {}
impl<I: FusedIterator> FusedIterator for CountIter {}

pub trait Countable {
 #[inline]
 fn counted(self) -> CountIter<Self>
 where
 Self: Sized,
 {
 CountIter::new(self)
 }
}

impl<I: Iterator> Countable for I {}