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; and the split of an
10//! Annex-B stream into its NAL units on the way in. Nothing here reaches outside the module,
11//! so a fuzz build (`fuzz/`) compiles it on its own.
12
13/// Bits appended most significant first.
14#[derive(Default)]
15pub struct BitWriter {
16    bytes: Vec<u8>,
17    bits: usize,
18}
19
20impl BitWriter {
21    pub fn new() -> Self {
22        Self::default()
23    }
24
25    /// The number of bits written so far.
26    pub fn len(&self) -> usize {
27        self.bits
28    }
29
30    pub fn is_empty(&self) -> bool {
31        self.bits == 0
32    }
33
34    /// Whether the writer stands on a byte boundary.
35    pub fn aligned(&self) -> bool {
36        self.bits.is_multiple_of(8)
37    }
38
39    /// The low `count` bits of `value`, most significant first.
40    pub fn u(&mut self, count: u32, value: u64) {
41        for shift in (0..count).rev() {
42            if self.bits.is_multiple_of(8) {
43                self.bytes.push(0);
44            }
45            if (value >> shift) & 1 != 0 {
46                let index = self.bits >> 3;
47                self.bytes[index] |= 1 << (7 - (self.bits & 7));
48            }
49            self.bits += 1;
50        }
51    }
52
53    /// One flag bit.
54    pub fn flag(&mut self, value: bool) {
55        self.u(1, u64::from(value));
56    }
57
58    /// An unsigned Exp-Golomb code.
59    pub fn ue(&mut self, value: u32) {
60        let v = u64::from(value) + 1;
61        let zeros = 63 - v.leading_zeros();
62        self.u(zeros, 0);
63        self.u(zeros + 1, v);
64    }
65
66    /// A signed Exp-Golomb code.
67    pub fn se(&mut self, value: i32) {
68        let mapped = if value > 0 {
69            2 * value as u32 - 1
70        } else {
71            (-2 * value) as u32
72        };
73        self.ue(mapped);
74    }
75
76    /// The RBSP stop bit and the zero bits that align what follows.
77    pub fn trailing_bits(&mut self) {
78        self.flag(true);
79        self.align();
80    }
81
82    /// Zero bits up to the next byte boundary.
83    pub fn align(&mut self) {
84        while !self.bits.is_multiple_of(8) {
85            self.flag(false);
86        }
87    }
88
89    /// Whole bytes appended at a byte boundary.
90    pub fn bytes(&mut self, bytes: &[u8]) {
91        debug_assert!(self.aligned());
92        self.bytes.extend_from_slice(bytes);
93        self.bits += 8 * bytes.len();
94    }
95
96    /// The bytes written; the last one is padded with zero bits.
97    pub fn into_bytes(self) -> Vec<u8> {
98        self.bytes
99    }
100
101    pub fn as_bytes(&self) -> &[u8] {
102        &self.bytes
103    }
104}
105
106/// An RBSP as a NAL unit payload: an emulation prevention byte ahead of every byte of three
107/// or less that follows two zeros.
108pub fn escape(rbsp: &[u8]) -> Vec<u8> {
109    let mut out = Vec::with_capacity(rbsp.len() + rbsp.len() / 64 + 4);
110    let mut zeros = 0;
111    for &byte in rbsp {
112        if zeros >= 2 && byte <= 3 {
113            out.push(3);
114            zeros = 0;
115        }
116        out.push(byte);
117        zeros = if byte == 0 { zeros + 1 } else { 0 };
118    }
119    out
120}
121
122/// One Annex-B NAL unit: a four-byte start code, the header bytes, and the escaped payload.
123pub fn nal_unit(header: &[u8], rbsp: &[u8]) -> Vec<u8> {
124    let mut out = vec![0, 0, 0, 1];
125    out.extend_from_slice(header);
126    out.extend(escape(rbsp));
127    out
128}
129
130/// The NAL units of an Annex-B stream, each without its start code and never empty.
131pub(crate) fn annexb_nals(au: &[u8]) -> impl Iterator<Item = &[u8]> {
132    let mut starts = Vec::new();
133    let mut i = 0usize;
134    while i + 2 < au.len() {
135        if au[i] == 0 && au[i + 1] == 0 && au[i + 2] == 1 {
136            starts.push(i + 3);
137            i += 3;
138        } else {
139            i += 1;
140        }
141    }
142    let mut nals = Vec::with_capacity(starts.len());
143    for (n, &start) in starts.iter().enumerate() {
144        let mut end = starts.get(n + 1).map_or(au.len(), |&next| next - 3);
145        while end > start && au[end - 1] == 0 {
146            end -= 1;
147        }
148        if end > start {
149            nals.push(&au[start..end]);
150        }
151    }
152    nals.into_iter()
153}
154
155#[cfg(test)]
156mod tests {
157    use super::*;
158
159    #[test]
160    fn exp_golomb_and_fixed_fields_pack_as_the_standards_read_them() {
161        let mut w = BitWriter::new();
162        w.ue(0);
163        w.ue(1);
164        w.ue(2);
165        w.ue(3);
166        w.se(1);
167        w.se(-1);
168        w.u(3, 0b101);
169        w.trailing_bits();
170        assert_eq!(w.as_bytes(), &[0b1010_0110, 0b0100_0100, 0b1110_1100]);
171        assert_eq!(w.len(), 24);
172    }
173
174    #[test]
175    fn emulation_prevention_guards_every_start_code() {
176        assert_eq!(
177            escape(&[0, 0, 1, 0, 0, 0, 5, 0, 0, 3]),
178            vec![0, 0, 3, 1, 0, 0, 3, 0, 5, 0, 0, 3, 3]
179        );
180        assert_eq!(
181            nal_unit(&[0x67], &[0, 0, 2]),
182            vec![0, 0, 0, 1, 0x67, 0, 0, 3, 2]
183        );
184    }
185}
pixelflux/encoders/bits.rs

pixelflux/encoders/
bits.rs
