Skip to main content

pith_digest/
bitreader.rs

1//! A bit-level reader over a byte slice.
2//!
3//! The stream order is the ISO media convention (MP3, H.264, FLAC,
4//! most length-prefixed containers): bit 0 of the stream is the
5//! most-significant bit of byte 0, and a multi-bit field's first
6//! stream bit becomes the result's most-significant bit. DEFLATE's
7//! opposite packing lives behind `pith-inflate`'s own reader, which
8//! is RFC-specific by design.
9
10use crate::error::{Error, Result};
11
12/// Reads bit-level fields from a byte slice, never panicking.
13///
14/// All methods bound-check and return [`crate::Error::Truncated`]
15/// instead of panicking or wrapping, so hostile input surfaces as an
16/// error at the read that needed it.
17#[derive(Debug)]
18pub struct BitReader<'a> {
19    data: &'a [u8],
20    /// Index of the next bit to read, counting from the MSB of byte 0.
21    pos: usize,
22}
23
24impl<'a> BitReader<'a> {
25    /// Creates a reader over `data`, positioned at bit 0.
26    pub const fn new(data: &'a [u8]) -> Self {
27        BitReader { data, pos: 0 }
28    }
29
30    /// The number of bits still readable.
31    pub const fn remaining_bits(&self) -> usize {
32        self.data.len() * 8 - self.pos
33    }
34
35    /// The current bit position from the start of the stream.
36    pub const fn bit_position(&self) -> usize {
37        self.pos
38    }
39
40    /// `true` when the position is on a byte boundary.
41    pub const fn is_aligned(&self) -> bool {
42        self.pos % 8 == 0
43    }
44
45    /// Discards bits up to the next byte boundary. Aligned readers are
46    /// unchanged.
47    pub fn byte_align(&mut self) {
48        self.pos = self.pos.div_ceil(8) * 8;
49    }
50
51    /// Reads `n` bits with the first stream bit becoming the
52    /// most-significant bit of the result, for `n` in `0..=64`.
53    ///
54    /// ```
55    /// use pith_digest::BitReader;
56    ///
57    /// let mut r = BitReader::new(&[0b1010_0110]);
58    /// assert_eq!(r.bits(3).unwrap(), 0b101);
59    /// assert_eq!(r.bits(5).unwrap(), 0b0_0110);
60    /// ```
61    pub fn bits(&mut self, n: usize) -> Result<u64> {
62        if n > 64 {
63            return Err(Error::BadValue("bit width over 64"));
64        }
65        if n > self.remaining_bits() {
66            return Err(Error::truncated("bits", n, self.remaining_bits()));
67        }
68        if n == 0 {
69            return Ok(0);
70        }
71        // A request that fits in the unread tail of the current byte is
72        // the hot case (Huffman side info is read one bit at a time);
73        // the general path assembles the field from whole bytes.
74        if n <= 8 - (self.pos % 8) {
75            let off = self.pos % 8;
76            let byte = self.data[self.pos / 8];
77            // Shift the wanted field to the bottom, then mask it.
78            // `0xFF >> (8 - n)` avoids the u8-shift overflow at n == 8.
79            let v = (byte >> (8 - off - n)) & (0xFFu8 >> (8 - n));
80            self.pos += n;
81            return Ok(u64::from(v));
82        }
83        let byte_pos = self.pos / 8;
84        let head_bits = (8 - self.pos % 8) % 8; // bits to the byte edge
85        let mut value = if head_bits == 0 {
86            0u64
87        } else {
88            u64::from(self.data[byte_pos] & ((1u8 << head_bits) - 1))
89        };
90        let mut left = n - head_bits;
91        let mut i = byte_pos + usize::from(head_bits != 0);
92        while left >= 8 {
93            value = (value << 8) | u64::from(self.data[i]);
94            i += 1;
95            left -= 8;
96        }
97        if left > 0 {
98            let v = self.data[i] >> (8 - left);
99            value = (value << left) | u64::from(v);
100        }
101        self.pos += n;
102        Ok(value)
103    }
104
105    /// Reads one whole byte at the current (must be aligned) position.
106    pub fn byte(&mut self) -> Result<u8> {
107        if !self.is_aligned() {
108            return Err(Error::BadValue("unaligned byte read"));
109        }
110        if self.remaining_bits() < 8 {
111            return Err(Error::truncated("byte", 8, self.remaining_bits()));
112        }
113        let byte = self.data[self.pos / 8];
114        self.pos += 8;
115        Ok(byte)
116    }
117}