arcly-stream 0.1.3

//! A big-endian bit reader with Exp-Golomb decoding, for codec headers.

/// Reads bits MSB-first from a byte slice, with H.264/H.265 Exp-Golomb helpers.
pub struct BitReader<'a> {
    data: &'a [u8],
    /// Absolute bit position from the start of `data`.
    pos: usize,
}

// Not every codec feature uses every reader method; suppress per-feature
// dead-code warnings on this internal helper.
#[allow(dead_code)]
impl<'a> BitReader<'a> {
    /// Wrap `data`, positioned at bit 0.
    pub fn new(data: &'a [u8]) -> Self {
        Self { data, pos: 0 }
    }

    /// Bits remaining.
    pub fn remaining(&self) -> usize {
        (self.data.len() * 8).saturating_sub(self.pos)
    }

    /// Skip `n` bits, returning `None` if that would run past the end. Useful for
    /// stepping over fixed-size syntax blocks (e.g. VVC `general_constraint_info`).
    pub fn skip_bits(&mut self, n: usize) -> Option<()> {
        let new = self.pos.checked_add(n)?;
        if new > self.data.len() * 8 {
            return None;
        }
        self.pos = new;
        Some(())
    }

    /// Advance to the next byte boundary (the `byte_aligned()` / alignment-bit
    /// idiom in NAL syntax). A no-op when already aligned.
    pub fn align_to_byte(&mut self) {
        let rem = self.pos % 8;
        if rem != 0 {
            self.pos += 8 - rem;
        }
    }

    /// Read a single bit, or `None` past the end.
    pub fn read_bit(&mut self) -> Option<u32> {
        let byte = self.data.get(self.pos / 8)?;
        let shift = 7 - (self.pos % 8);
        self.pos += 1;
        Some(((byte >> shift) & 1) as u32)
    }

    /// Read `n` (≤ 32) bits as an unsigned big-endian integer.
    pub fn read_bits(&mut self, n: u32) -> Option<u32> {
        let mut v = 0u32;
        for _ in 0..n {
            v = (v << 1) | self.read_bit()?;
        }
        Some(v)
    }

    /// Decode an unsigned Exp-Golomb code (`ue(v)`).
    pub fn read_ue(&mut self) -> Option<u32> {
        let mut zeros = 0u32;
        while self.read_bit()? == 0 {
            zeros += 1;
            if zeros > 31 {
                return None; // malformed / runaway
            }
        }
        if zeros == 0 {
            return Some(0);
        }
        let rest = self.read_bits(zeros)?;
        Some((1u32 << zeros) - 1 + rest)
    }

    /// Decode a signed Exp-Golomb code (`se(v)`).
    pub fn read_se(&mut self) -> Option<i32> {
        let k = self.read_ue()?;
        let val = k.div_ceil(2) as i32;
        Some(if k % 2 == 1 { val } else { -val })
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn reads_bits_msb_first() {
        // 0b1011_0010
        let mut r = BitReader::new(&[0b1011_0010]);
        assert_eq!(r.read_bit(), Some(1));
        assert_eq!(r.read_bits(3), Some(0b011));
        assert_eq!(r.read_bits(4), Some(0b0010));
        assert_eq!(r.read_bit(), None);
    }

    #[test]
    fn exp_golomb_ue_matches_spec_table() {
        // Canonical ue(v) codes: 1=>0, 010=>1, 011=>2, 00100=>3, 00111=>6.
        // Bitstream: 1 010 011 00100 00111  → pack MSB-first.
        // bits: 1 0 1 0 0 1 1 0 0 1 0 0 0 0 1 1 1  (17 bits)
        let bytes = [0b1010_0110, 0b0100_0011, 0b1000_0000];
        let mut r = BitReader::new(&bytes);
        assert_eq!(r.read_ue(), Some(0));
        assert_eq!(r.read_ue(), Some(1));
        assert_eq!(r.read_ue(), Some(2));
        assert_eq!(r.read_ue(), Some(3));
        assert_eq!(r.read_ue(), Some(6));
    }

    #[test]
    fn exp_golomb_se_zigzag() {
        // ue 0->se 0, ue 1->se 1, ue 2->se -1, ue 3->se 2, ue 4->se -2.
        // codes: 1(0) 010(1) 011(2) 00100(3) 00101(4)
        let bytes = [0b1010_0110, 0b0100_0010, 0b1000_0000];
        let mut r = BitReader::new(&bytes);
        assert_eq!(r.read_se(), Some(0));
        assert_eq!(r.read_se(), Some(1));
        assert_eq!(r.read_se(), Some(-1));
        assert_eq!(r.read_se(), Some(2));
        assert_eq!(r.read_se(), Some(-2));
    }
}