chdlady-huffman 0.1.0

//! Bitstream reader and writer for arbitrary-width bit sequences.
/// Reads arbitrary bit sequences (1 to 32 bits) from an in-memory byte slice.
#[derive(Debug, Clone)]
pub struct BitReader<'a> {
    data: &'a [u8],
    buffer: u32,
    bits: usize,
    offset: usize,
    bit_offset: usize,
}

impl<'a> BitReader<'a> {
    /// Creates a new BitReader reading from the provided byte slice.
    pub fn new(data: &'a [u8]) -> Self {
        Self {
            data,
            buffer: 0,
            bits: 0,
            offset: 0,
            bit_offset: 0,
        }
    }

    /// Returns true if more bytes were requested than exist in the source buffer.
    pub fn overflow(&self) -> bool {
        let full_bytes = self.bits / 8;
        if self.offset >= full_bytes {
            (self.offset - full_bytes) > self.data.len()
        } else {
            false
        }
    }

    /// Returns the current read offset in bytes.
    pub fn read_offset(&self) -> usize {
        let mut result = self.offset;
        let mut bits = self.bits;
        while bits >= 8 {
            result = result.saturating_sub(1);
            bits -= 8;
        }
        if self.bit_offset > bits {
            result = result.saturating_add(1);
        }
        result
    }

    /// Peeks up to 32 bits from the bitstream without consuming them.
    #[inline(always)]
    pub fn peek(&mut self, numbits: usize) -> u32 {
        if numbits == 0 {
            return 0;
        }
        debug_assert!(numbits <= 32);

        if numbits > self.bits {
            while self.bits < 32 {
                let mut newbits = 0u32;
                if self.offset < self.data.len() {
                    newbits = (u32::from(self.data[self.offset]) << self.bit_offset) & 0xff;
                }

                if self.bits + 8 > 32 {
                    self.bit_offset = 32 - self.bits;
                    newbits >>= 8 - self.bit_offset;
                    self.buffer |= newbits;
                    self.bits += self.bit_offset;
                } else {
                    self.buffer |= newbits << (24 - self.bits);
                    self.bits += 8 - self.bit_offset;
                    self.bit_offset = 0;
                    self.offset += 1;
                }
            }
        }

        if numbits == 32 {
            self.buffer
        } else {
            self.buffer >> (32 - numbits)
        }
    }

    /// Discards the next `numbits` from the bitstream.
    #[inline(always)]
    pub fn remove(&mut self, numbits: usize) {
        if numbits >= 32 {
            self.buffer = 0;
            self.bits = self.bits.saturating_sub(numbits);
        } else {
            self.buffer <<= numbits;
            self.bits = self.bits.saturating_sub(numbits);
        }
    }

    /// Reads `numbits` (0 to 32) from the bitstream.
    pub fn read(&mut self, numbits: usize) -> u32 {
        let result = self.peek(numbits);
        self.remove(numbits);
        result
    }

    /// Flushes to the nearest byte boundary and returns byte offset.
    pub fn flush(&mut self) -> usize {
        while self.bits >= 8 {
            self.offset = self.offset.saturating_sub(1);
            self.bits -= 8;
        }
        if self.bit_offset > self.bits {
            self.offset = self.offset.saturating_add(1);
        }
        self.bits = 0;
        self.buffer = 0;
        self.bit_offset = 0;
        self.offset
    }
}

/// Writes arbitrary bit sequences (1 to 32 bits) to an output byte stream.
#[derive(Debug, Clone, Default)]
pub struct BitWriter {
    data: Vec<u8>,
    buffer: u32,
    bits: usize,
}

impl BitWriter {
    /// Creates a new, empty BitWriter.
    pub fn new() -> Self {
        Self::default()
    }

    /// Creates a new BitWriter with pre-allocated buffer capacity.
    pub fn with_capacity(capacity: usize) -> Self {
        Self {
            data: Vec::with_capacity(capacity),
            buffer: 0,
            bits: 0,
        }
    }

    /// Writes `numbits` (0 to 32) from `value` to the bitstream.
    pub fn write(&mut self, value: u32, mut numbits: usize) {
        if numbits == 0 {
            return;
        }
        debug_assert!(numbits <= 32);

        let mut newbits = if numbits == 32 {
            value
        } else {
            (value & ((1u32 << numbits) - 1)) << (32 - numbits)
        };

        while self.bits + numbits >= 32 && numbits > 0 {
            while self.bits >= 8 {
                self.data.push((self.buffer >> 24) as u8);
                self.buffer <<= 8;
                self.bits -= 8;
            }

            if self.bits + numbits >= 32 {
                let rem = (32 - self.bits).min(numbits);
                self.buffer |= newbits >> self.bits;
                self.bits += rem;
                if rem == 32 {
                    newbits = 0;
                } else {
                    newbits <<= rem;
                }
                numbits -= rem;
            }
        }

        if numbits > 0 {
            self.buffer |= newbits >> self.bits;
            self.bits += numbits;
        }
    }

    /// Flushes unwritten bits to the nearest byte boundary.
    pub fn flush(&mut self) -> usize {
        while self.bits > 0 {
            self.data.push((self.buffer >> 24) as u8);
            self.buffer <<= 8;
            self.bits = self.bits.saturating_sub(8);
        }
        self.bits = 0;
        self.buffer = 0;
        self.data.len()
    }

    /// Consumes the writer, flushes remaining bits, and returns written bytes.
    pub fn into_bytes(mut self) -> Vec<u8> {
        self.flush();
        self.data
    }

    /// Returns written bytes without consuming the writer.
    pub fn as_bytes(&self) -> &[u8] {
        &self.data
    }
}