oximedia_codec/av1/
obu.rs

1//! AV1 OBU (Open Bitstream Unit) parsing.
2//!
3//! OBUs are the fundamental unit of AV1 bitstreams. Each OBU contains
4//! a header followed by optional payload data.
5
6use crate::error::{CodecError, CodecResult};
7use oximedia_io::BitReader;
8
9/// OBU types as defined in AV1 specification section 5.3.1.
10#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
11pub enum ObuType {
12    /// Reserved (0).
13    Reserved0,
14    /// Sequence header OBU.
15    SequenceHeader,
16    /// Temporal delimiter OBU.
17    TemporalDelimiter,
18    /// Frame header OBU.
19    FrameHeader,
20    /// Tile group OBU.
21    TileGroup,
22    /// Metadata OBU.
23    Metadata,
24    /// Combined frame header and tile group.
25    Frame,
26    /// Redundant frame header OBU.
27    RedundantFrameHeader,
28    /// Tile list OBU.
29    TileList,
30    /// Reserved (9-14).
31    Reserved(u8),
32    /// Padding OBU.
33    Padding,
34}
35
36impl From<u8> for ObuType {
37    fn from(value: u8) -> Self {
38        match value {
39            0 => Self::Reserved0,
40            1 => Self::SequenceHeader,
41            2 => Self::TemporalDelimiter,
42            3 => Self::FrameHeader,
43            4 => Self::TileGroup,
44            5 => Self::Metadata,
45            6 => Self::Frame,
46            7 => Self::RedundantFrameHeader,
47            8 => Self::TileList,
48            15 => Self::Padding,
49            other => Self::Reserved(other),
50        }
51    }
52}
53
54impl From<ObuType> for u8 {
55    fn from(obu_type: ObuType) -> Self {
56        match obu_type {
57            ObuType::Reserved0 => 0,
58            ObuType::SequenceHeader => 1,
59            ObuType::TemporalDelimiter => 2,
60            ObuType::FrameHeader => 3,
61            ObuType::TileGroup => 4,
62            ObuType::Metadata => 5,
63            ObuType::Frame => 6,
64            ObuType::RedundantFrameHeader => 7,
65            ObuType::TileList => 8,
66            ObuType::Padding => 15,
67            ObuType::Reserved(v) => v,
68        }
69    }
70}
71
72/// OBU header as defined in AV1 specification section 5.3.2.
73#[derive(Clone, Debug)]
74pub struct ObuHeader {
75    /// OBU type.
76    pub obu_type: ObuType,
77    /// Extension flag present.
78    pub has_extension: bool,
79    /// Size field present.
80    pub has_size: bool,
81    /// Temporal ID (from extension header).
82    pub temporal_id: u8,
83    /// Spatial ID (from extension header).
84    pub spatial_id: u8,
85}
86
87impl ObuHeader {
88    /// Parse OBU header from bitstream.
89    ///
90    /// # Errors
91    ///
92    /// Returns error if header is malformed.
93    #[allow(clippy::cast_possible_truncation)]
94    pub fn parse(reader: &mut BitReader<'_>) -> CodecResult<Self> {
95        let forbidden = reader.read_bit().map_err(CodecError::Core)?;
96        if forbidden != 0 {
97            return Err(CodecError::InvalidBitstream(
98                "OBU forbidden bit is set".to_string(),
99            ));
100        }
101
102        let obu_type_raw = reader.read_bits(4).map_err(CodecError::Core)? as u8;
103        let obu_type = ObuType::from(obu_type_raw);
104        let has_extension = reader.read_bit().map_err(CodecError::Core)? != 0;
105        let has_size = reader.read_bit().map_err(CodecError::Core)? != 0;
106
107        let reserved = reader.read_bit().map_err(CodecError::Core)?;
108        if reserved != 0 {
109            return Err(CodecError::InvalidBitstream(
110                "OBU reserved bit is set".to_string(),
111            ));
112        }
113
114        let (temporal_id, spatial_id) = if has_extension {
115            let temporal_id = reader.read_bits(3).map_err(CodecError::Core)? as u8;
116            let spatial_id = reader.read_bits(2).map_err(CodecError::Core)? as u8;
117            let _reserved = reader.read_bits(3).map_err(CodecError::Core)?;
118            (temporal_id, spatial_id)
119        } else {
120            (0, 0)
121        };
122
123        Ok(Self {
124            obu_type,
125            has_extension,
126            has_size,
127            temporal_id,
128            spatial_id,
129        })
130    }
131
132    /// Get the header size in bytes.
133    #[must_use]
134    pub const fn header_size(&self) -> usize {
135        if self.has_extension {
136            2
137        } else {
138            1
139        }
140    }
141
142    /// Serialize OBU header to bytes.
143    #[must_use]
144    pub fn to_bytes(&self) -> Vec<u8> {
145        let mut bytes = Vec::with_capacity(2);
146        let obu_type_val: u8 = self.obu_type.into();
147        let first_byte = (obu_type_val << 3)
148            | (u8::from(self.has_extension) << 2)
149            | (u8::from(self.has_size) << 1);
150        bytes.push(first_byte);
151
152        if self.has_extension {
153            let ext_byte = (self.temporal_id << 5) | (self.spatial_id << 3);
154            bytes.push(ext_byte);
155        }
156
157        bytes
158    }
159}
160
161/// Parse LEB128 encoded unsigned integer.
162#[allow(clippy::cast_possible_truncation)]
163pub fn parse_leb128(reader: &mut BitReader<'_>) -> CodecResult<u64> {
164    let mut value: u64 = 0;
165    let mut shift = 0;
166
167    loop {
168        let byte = reader.read_bits(8).map_err(CodecError::Core)? as u8;
169        value |= u64::from(byte & 0x7F) << shift;
170
171        if byte & 0x80 == 0 {
172            break;
173        }
174
175        shift += 7;
176        if shift >= 64 {
177            return Err(CodecError::InvalidBitstream(
178                "LEB128 value overflow".to_string(),
179            ));
180        }
181    }
182
183    Ok(value)
184}
185
186/// Encode a value as LEB128.
187#[must_use]
188pub fn encode_leb128(mut value: u64) -> Vec<u8> {
189    let mut bytes = Vec::new();
190
191    loop {
192        let mut byte = (value & 0x7F) as u8;
193        value >>= 7;
194
195        if value != 0 {
196            byte |= 0x80;
197        }
198        bytes.push(byte);
199
200        if value == 0 {
201            break;
202        }
203    }
204
205    bytes
206}
207
208/// Parse a complete OBU from data.
209#[allow(clippy::cast_possible_truncation)]
210pub fn parse_obu(data: &[u8]) -> CodecResult<(ObuHeader, &[u8], usize)> {
211    let mut reader = BitReader::new(data);
212    let header = ObuHeader::parse(&mut reader)?;
213
214    let size = if header.has_size {
215        parse_leb128(&mut reader)? as usize
216    } else {
217        let header_bytes = reader.bits_read().div_ceil(8);
218        data.len().saturating_sub(header_bytes)
219    };
220
221    let header_bytes = reader.bits_read().div_ceil(8);
222    let total_size = header_bytes + size;
223
224    if total_size > data.len() {
225        return Err(CodecError::InvalidBitstream(format!(
226            "OBU size {} exceeds available data {}",
227            total_size,
228            data.len()
229        )));
230    }
231
232    let payload = &data[header_bytes..header_bytes + size];
233    Ok((header, payload, total_size))
234}
235
236/// Iterator over OBUs in a temporal unit.
237pub struct ObuIterator<'a> {
238    data: &'a [u8],
239    offset: usize,
240}
241
242impl<'a> ObuIterator<'a> {
243    /// Create a new OBU iterator.
244    #[must_use]
245    pub const fn new(data: &'a [u8]) -> Self {
246        Self { data, offset: 0 }
247    }
248}
249
250impl<'a> Iterator for ObuIterator<'a> {
251    type Item = CodecResult<(ObuHeader, &'a [u8])>;
252
253    fn next(&mut self) -> Option<Self::Item> {
254        if self.offset >= self.data.len() {
255            return None;
256        }
257
258        match parse_obu(&self.data[self.offset..]) {
259            Ok((header, payload, total_size)) => {
260                self.offset += total_size;
261                Some(Ok((header, payload)))
262            }
263            Err(e) => Some(Err(e)),
264        }
265    }
266}
267
268#[cfg(test)]
269mod tests {
270    use super::*;
271
272    #[test]
273    fn test_obu_type_from_u8() {
274        assert_eq!(ObuType::from(1), ObuType::SequenceHeader);
275        assert_eq!(ObuType::from(6), ObuType::Frame);
276    }
277
278    #[test]
279    fn test_leb128_single_byte() {
280        let data = [0x7F];
281        let mut reader = BitReader::new(&data);
282        assert_eq!(parse_leb128(&mut reader).expect("should succeed"), 127);
283    }
284
285    #[test]
286    fn test_leb128_multi_byte() {
287        let data = [0x80, 0x01];
288        let mut reader = BitReader::new(&data);
289        assert_eq!(parse_leb128(&mut reader).expect("should succeed"), 128);
290    }
291
292    #[test]
293    fn test_encode_leb128() {
294        assert_eq!(encode_leb128(0), vec![0x00]);
295        assert_eq!(encode_leb128(127), vec![0x7F]);
296        assert_eq!(encode_leb128(128), vec![0x80, 0x01]);
297    }
298}
oximedia_codec/av1/obu.rs

oximedia_codec/av1/
obu.rs
