use crate::av1::common::{leb128_size, parse_leb128};
use bitstream_io::{BitRead, BitReader, Endianness};
use std::io::{self, Read, Seek};
#[derive(Default, Debug, Clone, Copy, PartialEq, Eq)]
pub struct UnsizedObu {
pub obu_type: ObuType,
pub has_extension: bool,
pub temporal_id: u8,
pub spatial_id: u8,
pub header_len: u32,
pub is_fragment: bool,
}
#[derive(Default, Debug, Clone, Copy, PartialEq, Eq)]
pub struct SizedObu {
pub obu_type: ObuType,
pub has_extension: bool,
pub has_size_field: bool,
pub temporal_id: u8,
pub spatial_id: u8,
pub size: u32,
pub leb_size: u32,
pub header_len: u32,
pub is_fragment: bool,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ObuType {
Reserved,
SequenceHeader,
TemporalDelimiter,
FrameHeader,
TileGroup,
Metadata,
Frame,
RedundantFrameHeader,
TileList,
Padding,
}
impl Default for ObuType {
fn default() -> Self {
Self::Reserved
}
}
impl UnsizedObu {
pub fn parse<R, E>(reader: &mut BitReader<R, E>) -> io::Result<Self>
where
R: Read + Seek,
E: Endianness,
{
if reader.read_bit()? {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
"forbidden bit in OBU header is set",
));
}
let obu_type = reader.read::<u8>(4)?.into();
let has_extension = reader.read_bit()?;
if reader.read_bit()? {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
"did not expect size field",
));
}
let _ = reader.read_bit()?;
let (temporal_id, spatial_id) = if has_extension {
(reader.read::<u8>(3)?, reader.read::<u8>(2)?)
} else {
(0, 0)
};
reader.byte_align();
Ok(Self {
obu_type,
has_extension,
temporal_id,
spatial_id,
header_len: has_extension as u32 + 1,
is_fragment: false,
})
}
pub fn as_sized(&self, size: u32, leb_size: u32) -> SizedObu {
SizedObu {
obu_type: self.obu_type,
has_extension: self.has_extension,
has_size_field: false,
temporal_id: self.temporal_id,
spatial_id: self.spatial_id,
size,
leb_size,
header_len: self.header_len,
is_fragment: self.is_fragment,
}
}
}
impl SizedObu {
pub fn parse<R, E>(reader: &mut BitReader<R, E>) -> io::Result<Self>
where
R: Read + Seek,
E: Endianness,
{
if reader.read_bit()? {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
"forbidden bit in OBU header is set",
));
}
let obu_type = reader.read::<u8>(4)?.into();
let has_extension = reader.read_bit()?;
if !reader.read_bit()? {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
"expected a size field",
));
}
let _ = reader.read_bit()?;
let (temporal_id, spatial_id) = if has_extension {
(reader.read::<u8>(3)?, reader.read::<u8>(2)?)
} else {
(0, 0)
};
reader.byte_align();
let size = parse_leb128(reader)?;
let leb_size = leb128_size(size) as u32;
Ok(Self {
obu_type,
has_extension,
has_size_field: true,
temporal_id,
spatial_id,
size,
leb_size,
header_len: has_extension as u32 + 1,
is_fragment: false,
})
}
pub fn full_size(&self) -> u32 {
self.size + self.leb_size + self.header_len
}
pub fn partial_size(&self) -> u32 {
self.size + self.header_len
}
}
impl From<u8> for ObuType {
fn from(n: u8) -> Self {
assert!(n < 16);
match n {
1 => Self::SequenceHeader,
2 => Self::TemporalDelimiter,
3 => Self::FrameHeader,
4 => Self::TileGroup,
5 => Self::Metadata,
6 => Self::Frame,
7 => Self::RedundantFrameHeader,
8 => Self::TileList,
15 => Self::Padding,
_ => Self::Reserved,
}
}
}
impl From<ObuType> for u8 {
fn from(ty: ObuType) -> Self {
match ty {
ObuType::Reserved => 0,
ObuType::SequenceHeader => 1,
ObuType::TemporalDelimiter => 2,
ObuType::FrameHeader => 3,
ObuType::TileGroup => 4,
ObuType::Metadata => 5,
ObuType::Frame => 6,
ObuType::RedundantFrameHeader => 7,
ObuType::TileList => 8,
ObuType::Padding => 15,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use bitstream_io::{BigEndian, BitRead, BitReader};
use once_cell::sync::Lazy;
use std::io::Cursor;
#[allow(clippy::type_complexity)]
static OBUS: Lazy<Vec<(SizedObu, Vec<u8>, u64, UnsizedObu, Vec<u8>)>> = Lazy::new(|| {
vec![
(
SizedObu {
obu_type: ObuType::TemporalDelimiter,
has_extension: false,
has_size_field: true,
temporal_id: 0,
spatial_id: 0,
size: 0,
leb_size: 1,
header_len: 1,
is_fragment: false,
},
vec![0b0001_0010, 0b0000_0000],
2,
UnsizedObu {
obu_type: ObuType::TemporalDelimiter,
has_extension: false,
temporal_id: 0,
spatial_id: 0,
header_len: 1,
is_fragment: false,
},
vec![0b0001_0000],
),
(
SizedObu {
obu_type: ObuType::Padding,
has_extension: false,
has_size_field: true,
temporal_id: 0,
spatial_id: 0,
size: 10,
leb_size: 1,
header_len: 1,
is_fragment: false,
},
vec![0b0111_1010, 0b0000_1010, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
2,
UnsizedObu {
obu_type: ObuType::Padding,
has_extension: false,
temporal_id: 0,
spatial_id: 0,
header_len: 1,
is_fragment: false,
},
vec![0b0111_1000, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
),
(
SizedObu {
obu_type: ObuType::Frame,
has_extension: true,
has_size_field: true,
temporal_id: 4,
spatial_id: 3,
size: 5,
leb_size: 1,
header_len: 2,
is_fragment: false,
},
vec![0b0011_0110, 0b1001_1000, 0b0000_0101, 1, 2, 3, 4, 5],
3,
UnsizedObu {
obu_type: ObuType::Frame,
has_extension: true,
temporal_id: 4,
spatial_id: 3,
header_len: 2,
is_fragment: false,
},
vec![0b0011_0100, 0b1001_1000, 1, 2, 3, 4, 5],
),
]
});
#[test]
fn test_parse() {
for (idx, (sized_obu, sized_bytes, expected_position, unsized_obu, unsized_bytes)) in
(*OBUS).iter().enumerate()
{
println!("running test {}...", idx);
{
println!(" parsing sized...");
let mut reader = BitReader::endian(Cursor::new(&sized_bytes), BigEndian);
assert_eq!(SizedObu::parse(&mut reader).unwrap(), *sized_obu);
assert!(reader.byte_aligned());
assert_eq!(reader.into_reader().position(), *expected_position);
};
{
println!(" parsing unsized...");
let mut reader = BitReader::endian(Cursor::new(&unsized_bytes), BigEndian);
assert_eq!(UnsizedObu::parse(&mut reader).unwrap(), *unsized_obu);
assert!(reader.byte_aligned());
assert_eq!(
reader.into_reader().position(),
unsized_obu.header_len as u64
);
}
}
}
#[test]
fn test_conversion() {
for (idx, (sized_obu, _, _, unsized_obu, _)) in (*OBUS).iter().enumerate() {
println!("running test {}...", idx);
assert_eq!(
unsized_obu.as_sized(sized_obu.size, sized_obu.leb_size),
SizedObu {
has_size_field: false,
..*sized_obu
},
);
}
}
#[test]
fn test_parse_rtp_obu() {
let obus = [
(
SizedObu {
obu_type: ObuType::TemporalDelimiter,
has_extension: false,
has_size_field: false,
temporal_id: 0,
spatial_id: 0,
size: 0,
leb_size: 1,
header_len: 1,
is_fragment: false,
},
vec![0b0001_0000],
),
(
SizedObu {
obu_type: ObuType::Padding,
has_extension: false,
has_size_field: false,
temporal_id: 0,
spatial_id: 0,
size: 10,
leb_size: 1,
header_len: 1,
is_fragment: false,
},
vec![0b0111_1000, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
),
(
SizedObu {
obu_type: ObuType::Frame,
has_extension: true,
has_size_field: false,
temporal_id: 4,
spatial_id: 3,
size: 5,
leb_size: 1,
header_len: 2,
is_fragment: false,
},
vec![0b0011_0100, 0b1001_1000, 1, 2, 3, 4, 5],
),
];
for (idx, (sized_obu, rtp_bytes)) in obus.into_iter().enumerate() {
println!("running test {}...", idx);
let mut reader = BitReader::endian(Cursor::new(&rtp_bytes), BigEndian);
let unsized_obu = UnsizedObu::parse(&mut reader).unwrap();
assert_eq!(
unsized_obu.as_sized(sized_obu.size, sized_obu.leb_size),
sized_obu
);
}
}
}