arcly-stream 0.1.5

//! Shared bit/byte writer for synthesizing codec headers in unit tests.
// Each codec test uses a different subset of writers.
#![allow(dead_code)]

/// Big-endian bit writer mirroring [`BitReader`](super::bitreader::BitReader),
/// with Exp-Golomb (`ue`) and the AV1/VP9 helpers (`le`, `uvlc`) the parser
/// tests need to build known headers.
#[derive(Default)]
pub(crate) struct BitWriter {
    bits: Vec<u8>,
}

impl BitWriter {
    /// Append a single bit.
    pub fn bit(&mut self, b: u8) {
        self.bits.push(b & 1);
    }

    /// Append `n` bits of `v`, MSB first (`f(n)`).
    pub fn bits(&mut self, v: u32, n: u32) {
        for i in (0..n).rev() {
            self.bit(((v >> i) & 1) as u8);
        }
    }

    /// Append an unsigned Exp-Golomb code (`ue(v)`).
    pub fn ue(&mut self, v: u32) {
        let n = v + 1;
        let len = 32 - n.leading_zeros();
        for _ in 0..(len - 1) {
            self.bit(0);
        }
        self.bits(n, len);
    }

    /// Append an AV1 unsigned variable-length code (`uvlc`).
    #[allow(dead_code)]
    pub fn uvlc(&mut self, v: u32) {
        let n = v + 1;
        let len = 32 - n.leading_zeros();
        for _ in 0..(len - 1) {
            self.bit(0);
        }
        self.bit(1);
        if len > 1 {
            self.bits(v + 1 - (1 << (len - 1)), len - 1);
        }
    }

    /// Pad with zero bits to the next byte boundary.
    pub fn align(&mut self) {
        while self.bits.len() % 8 != 0 {
            self.bit(0);
        }
    }

    /// Materialize the written bits into bytes (zero-padded MSB-first).
    pub fn bytes(&self) -> Vec<u8> {
        let mut out = vec![0u8; self.bits.len().div_ceil(8)];
        for (i, &b) in self.bits.iter().enumerate() {
            if b == 1 {
                out[i / 8] |= 1 << (7 - (i % 8));
            }
        }
        out
    }
}