gamut_bitstream/
bitwriter.rs

1//! Most-significant-bit-first bit writer for AV1 uncompressed headers (AV1 §4, §8.1).
2
3/// A most-significant-bit-first bit writer.
4///
5/// AV1 uncompressed headers and OBU framing are described with `f(n)` fixed-width fields that a
6/// decoder reads most-significant-bit first (AV1 §8.1, "Parsing process for `f(n)`"); this writer
7/// is the encoder-side mirror. Bits accumulate into an internal byte buffer; the buffer is byte
8/// aligned exactly when [`BitWriter::is_byte_aligned`] returns `true`.
9#[derive(Debug, Default, Clone)]
10pub struct BitWriter {
11    bytes: Vec<u8>,
12    /// Bits already filled in the current (last) partial byte, `0..=7`. `0` means byte aligned.
13    bit_pos: u8,
14}
15
16impl BitWriter {
17    /// Creates an empty writer.
18    #[must_use]
19    pub fn new() -> Self {
20        Self::default()
21    }
22
23    /// Writes a single bit (`f(1)`); only the low bit of `bit` is used.
24    pub fn put_bit(&mut self, bit: u8) {
25        if self.bit_pos == 0 {
26            self.bytes.push(0);
27        }
28        if bit & 1 != 0 {
29            let last = self.bytes.len() - 1;
30            self.bytes[last] |= 1 << (7 - self.bit_pos);
31        }
32        self.bit_pos = (self.bit_pos + 1) & 7;
33    }
34
35    /// Writes the low `n` bits of `value`, most-significant bit first (`f(n)`). `n` must be `0..=32`.
36    pub fn put_bits(&mut self, value: u32, n: u32) {
37        debug_assert!(n <= 32);
38        for i in (0..n).rev() {
39            self.put_bit(((value >> i) & 1) as u8);
40        }
41    }
42
43    /// Number of bits written so far.
44    #[must_use]
45    pub fn bit_len(&self) -> usize {
46        if self.bit_pos == 0 {
47            self.bytes.len() * 8
48        } else {
49            (self.bytes.len() - 1) * 8 + self.bit_pos as usize
50        }
51    }
52
53    /// Returns `true` when the next bit would start a fresh byte.
54    #[must_use]
55    pub fn is_byte_aligned(&self) -> bool {
56        self.bit_pos == 0
57    }
58
59    /// Pads with zero bits up to the next byte boundary (a no-op when already aligned).
60    pub fn byte_align(&mut self) {
61        while self.bit_pos != 0 {
62            self.put_bit(0);
63        }
64    }
65
66    /// Appends whole bytes verbatim. Requires the writer to be byte aligned.
67    ///
68    /// # Panics
69    ///
70    /// Debug builds assert byte alignment; misuse in release silently misaligns the stream.
71    pub fn put_bytes(&mut self, data: &[u8]) {
72        debug_assert!(self.is_byte_aligned(), "put_bytes requires byte alignment");
73        self.bytes.extend_from_slice(data);
74    }
75
76    /// Borrows the bytes written so far.
77    #[must_use]
78    pub fn as_bytes(&self) -> &[u8] {
79        &self.bytes
80    }
81
82    /// Consumes the writer, returning the byte buffer (trailing bits in the final partial byte are
83    /// zero-filled).
84    #[must_use]
85    pub fn into_bytes(self) -> Vec<u8> {
86        self.bytes
87    }
88}
89
90#[cfg(test)]
91mod tests {
92    use super::*;
93
94    /// Reads `n` bits MSB-first starting at `bit_pos`; mirrors the AV1 `f(n)` decoder so tests can
95    /// assert the exact bit layout a real decoder would observe.
96    fn read_bits(bytes: &[u8], bit_pos: &mut usize, n: u32) -> u32 {
97        let mut x = 0u32;
98        for _ in 0..n {
99            let byte = bytes[*bit_pos >> 3];
100            let bit = (byte >> (7 - (*bit_pos & 7))) & 1;
101            x = (x << 1) | u32::from(bit);
102            *bit_pos += 1;
103        }
104        x
105    }
106
107    #[test]
108    fn put_bits_is_msb_first() {
109        let mut w = BitWriter::new();
110        w.put_bits(0b101, 3);
111        w.put_bits(0b01, 2);
112        assert_eq!(w.bit_len(), 5);
113        // 0b101_01 packed MSB-first into one byte: 1010_1000 = 0xA8.
114        assert_eq!(w.as_bytes(), &[0xA8]);
115    }
116
117    #[test]
118    fn byte_align_pads_with_zeros() {
119        let mut w = BitWriter::new();
120        w.put_bit(1);
121        w.byte_align();
122        assert!(w.is_byte_aligned());
123        assert_eq!(w.into_bytes(), &[0x80]);
124    }
125
126    #[test]
127    fn put_bytes_after_alignment() {
128        let mut w = BitWriter::new();
129        w.put_bits(0xF, 4);
130        w.byte_align();
131        w.put_bytes(&[0x12, 0x34]);
132        assert_eq!(w.into_bytes(), &[0xF0, 0x12, 0x34]);
133    }
134
135    #[test]
136    fn roundtrip_various_widths() {
137        let mut w = BitWriter::new();
138        let fields = [(0x1, 1u32), (0x2A, 6), (0xFFFF, 16), (0x0, 3), (0x7, 3)];
139        for &(v, n) in &fields {
140            w.put_bits(v, n);
141        }
142        let bytes = w.into_bytes();
143        let mut pos = 0usize;
144        for &(v, n) in &fields {
145            assert_eq!(read_bits(&bytes, &mut pos, n), v, "field {v:#x}/{n}");
146        }
147    }
148}
gamut_bitstream/bitwriter.rs

gamut_bitstream/
bitwriter.rs
