pixelflux/encoders/
bits.rs

1/*
2 * This Source Code Form is subject to the terms of the Mozilla Public
3 * License, v. 2.0. If a copy of the MPL was not distributed with this
4 * file, You can obtain one at https://mozilla.org/MPL/2.0/.
5 */
6
7//! A most-significant-bit-first writer for the headers a VA-API session packs itself: the
8//! fixed-width and Exp-Golomb fields of H.264, HEVC, and AV1 syntax, the RBSP trailing bits,
9//! and the emulation prevention an H.26x NAL unit needs on its way out.
10
11/// Bits appended most significant first.
12#[derive(Default)]
13pub struct BitWriter {
14    bytes: Vec<u8>,
15    bits: usize,
16}
17
18impl BitWriter {
19    pub fn new() -> Self {
20        Self::default()
21    }
22
23    /// The number of bits written so far.
24    pub fn len(&self) -> usize {
25        self.bits
26    }
27
28    pub fn is_empty(&self) -> bool {
29        self.bits == 0
30    }
31
32    /// Whether the writer stands on a byte boundary.
33    pub fn aligned(&self) -> bool {
34        self.bits.is_multiple_of(8)
35    }
36
37    /// The low `count` bits of `value`, most significant first.
38    pub fn u(&mut self, count: u32, value: u64) {
39        for shift in (0..count).rev() {
40            if self.bits.is_multiple_of(8) {
41                self.bytes.push(0);
42            }
43            if (value >> shift) & 1 != 0 {
44                let index = self.bits >> 3;
45                self.bytes[index] |= 1 << (7 - (self.bits & 7));
46            }
47            self.bits += 1;
48        }
49    }
50
51    /// One flag bit.
52    pub fn flag(&mut self, value: bool) {
53        self.u(1, value as u64);
54    }
55
56    /// An unsigned Exp-Golomb code.
57    pub fn ue(&mut self, value: u32) {
58        let v = value as u64 + 1;
59        let zeros = 63 - v.leading_zeros();
60        self.u(zeros, 0);
61        self.u(zeros + 1, v);
62    }
63
64    /// A signed Exp-Golomb code.
65    pub fn se(&mut self, value: i32) {
66        let mapped = if value > 0 {
67            2 * value as u32 - 1
68        } else {
69            (-2 * value) as u32
70        };
71        self.ue(mapped);
72    }
73
74    /// The RBSP stop bit and the zero bits that align what follows.
75    pub fn trailing_bits(&mut self) {
76        self.flag(true);
77        self.align();
78    }
79
80    /// Zero bits up to the next byte boundary.
81    pub fn align(&mut self) {
82        while !self.bits.is_multiple_of(8) {
83            self.flag(false);
84        }
85    }
86
87    /// Whole bytes appended at a byte boundary.
88    pub fn bytes(&mut self, bytes: &[u8]) {
89        debug_assert!(self.aligned());
90        self.bytes.extend_from_slice(bytes);
91        self.bits += 8 * bytes.len();
92    }
93
94    /// The bytes written; the last one is padded with zero bits.
95    pub fn into_bytes(self) -> Vec<u8> {
96        self.bytes
97    }
98
99    pub fn as_bytes(&self) -> &[u8] {
100        &self.bytes
101    }
102}
103
104/// An RBSP as a NAL unit payload: an emulation prevention byte ahead of every byte of three
105/// or less that follows two zeros.
106pub fn escape(rbsp: &[u8]) -> Vec<u8> {
107    let mut out = Vec::with_capacity(rbsp.len() + rbsp.len() / 64 + 4);
108    let mut zeros = 0;
109    for &byte in rbsp {
110        if zeros >= 2 && byte <= 3 {
111            out.push(3);
112            zeros = 0;
113        }
114        out.push(byte);
115        zeros = if byte == 0 { zeros + 1 } else { 0 };
116    }
117    out
118}
119
120/// One Annex-B NAL unit: a four-byte start code, the header bytes, and the escaped payload.
121pub fn nal_unit(header: &[u8], rbsp: &[u8]) -> Vec<u8> {
122    let mut out = vec![0, 0, 0, 1];
123    out.extend_from_slice(header);
124    out.extend(escape(rbsp));
125    out
126}
127
128#[cfg(test)]
129mod tests {
130    use super::*;
131
132    #[test]
133    fn exp_golomb_and_fixed_fields_pack_as_the_standards_read_them() {
134        let mut w = BitWriter::new();
135        w.ue(0);
136        w.ue(1);
137        w.ue(2);
138        w.ue(3);
139        w.se(1);
140        w.se(-1);
141        w.u(3, 0b101);
142        w.trailing_bits();
143        assert_eq!(w.as_bytes(), &[0b1010_0110, 0b0100_0100, 0b1110_1100]);
144        assert_eq!(w.len(), 24);
145    }
146
147    #[test]
148    fn emulation_prevention_guards_every_start_code() {
149        assert_eq!(
150            escape(&[0, 0, 1, 0, 0, 0, 5, 0, 0, 3]),
151            vec![0, 0, 3, 1, 0, 0, 3, 0, 5, 0, 0, 3, 3]
152        );
153        assert_eq!(
154            nal_unit(&[0x67], &[0, 0, 2]),
155            vec![0, 0, 0, 1, 0x67, 0, 0, 3, 2]
156        );
157    }
158}
pixelflux/encoders/bits.rs

pixelflux/encoders/
bits.rs
