use std::cell::Cell;
use std::io::Cursor;
use std::path::PathBuf;
use std::rc::Rc;
use std::sync::{Arc, Mutex, Weak};
use anyhow::{anyhow, bail, Context as _};
use crate::codec::h264::dpb::{Dpb, DpbEntry, DpbPicRefList, MmcoError, ReferencePicLists};
use crate::codec::h264::parser::{
Level, MaxLongTermFrameIdx, Nalu, NaluType, Parser, Pps, Profile, RefPicListModification, Slice, SliceHeader,
SliceType, Sps,
};
use crate::codec::h264::picture::{Field, IsIdr, PictureData, Reference};
use crate::{
bindings, Buffer, BufferType, Config as VaConfig, Context, Display, DrmPrimeSurfaceDescriptor, H264PicFields,
H264SeqFields, IQMatrix, IQMatrixBufferH264, Image, Picture, PictureH264, PictureParameter,
PictureParameterBufferH264, SliceParameter, SliceParameterBufferH264, Surface, UsageHint, VAConfigAttrib,
VAConfigAttribType, VAEntrypoint, VAProfile, VA_FOURCC_NV12, VA_INVALID_ID, VA_PICTURE_H264_INVALID,
VA_PICTURE_H264_LONG_TERM_REFERENCE, VA_PICTURE_H264_SHORT_TERM_REFERENCE, VA_RT_FORMAT_YUV420,
VA_SLICE_DATA_FLAG_ALL,
};
#[derive(Clone, Debug)]
pub struct Config {
pub device: Option<PathBuf>,
}
impl Config {
pub fn new() -> Self {
Self { device: None }
}
}
impl Default for Config {
fn default() -> Self {
Self::new()
}
}
#[derive(Clone)]
pub struct Frame {
pub timestamp: u64,
pub width: u32,
pub height: u32,
pub data: Vec<u8>,
}
pub struct ExportedFrame {
pub timestamp: u64,
pub width: u32,
pub height: u32,
pub descriptor: DrmPrimeSurfaceDescriptor,
surface: Arc<PooledSurface>,
}
impl ExportedFrame {
pub fn from_surface(surface: Surface<()>, timestamp: u64) -> anyhow::Result<Self> {
surface.sync().map_err(|e| anyhow!("surface sync: {e:?}"))?;
let (width, height) = surface.size();
let descriptor = surface
.export_prime()
.map_err(|e| anyhow!("export the surface: {e:?}"))?;
if descriptor.fourcc != VA_FOURCC_NV12 {
anyhow::bail!("exported fourcc {:#010x}, expected NV12", descriptor.fourcc);
}
Ok(Self {
timestamp,
width,
height,
descriptor,
surface: Arc::new(PooledSurface::unpooled(surface)),
})
}
pub fn download(&self) -> anyhow::Result<Frame> {
read_back(self.surface.get(), self.timestamp, self.width, self.height)
}
}
impl std::fmt::Debug for ExportedFrame {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("ExportedFrame")
.field("timestamp", &self.timestamp)
.field("width", &self.width)
.field("height", &self.height)
.field("fourcc", &self.descriptor.fourcc)
.finish()
}
}
impl std::fmt::Debug for Frame {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Frame")
.field("timestamp", &self.timestamp)
.field("width", &self.width)
.field("height", &self.height)
.field("bytes", &self.data.len())
.finish()
}
}
pub struct Decoder {
parser: Parser,
dpb: Dpb<Handle>,
prev_ref_pic_info: PrevReferencePicInfo,
prev_pic_info: PrevPicInfo,
max_long_term_frame_idx: MaxLongTermFrameIdx,
sequence: Option<Sequence>,
ready: Vec<Handle>,
display: Arc<Display>,
}
impl Decoder {
pub fn new(config: Config) -> anyhow::Result<Self> {
let (device, display) = match config.device {
Some(device) => {
let display =
Display::open_drm_display(&device).map_err(|e| anyhow!("open DRM display {device:?}: {e:?}"))?;
probe(&display)?;
(device, display)
}
None => {
crate::display::open_first(probe).map_err(|e| e.context("find a render node that decodes H.264"))?
}
};
log::info!("opened VA-API H.264 decoder on {device:?}");
Ok(Self {
parser: Parser::default(),
dpb: Dpb::default(),
prev_ref_pic_info: PrevReferencePicInfo::default(),
prev_pic_info: PrevPicInfo::default(),
max_long_term_frame_idx: MaxLongTermFrameIdx::default(),
sequence: None,
ready: Vec::new(),
display,
})
}
pub fn decode(&mut self, access_unit: &[u8], timestamp: u64) -> anyhow::Result<Vec<Frame>> {
self.submit(access_unit, timestamp)?;
self.take_frames()
}
pub fn decode_exported(&mut self, access_unit: &[u8], timestamp: u64) -> anyhow::Result<Vec<ExportedFrame>> {
self.submit(access_unit, timestamp)?;
self.take_exported()
}
fn submit(&mut self, access_unit: &[u8], timestamp: u64) -> anyhow::Result<()> {
let mut cursor = Cursor::new(access_unit);
let mut current: Option<CurrentPicture> = None;
let mut nalus = 0usize;
while let Ok(nalu) = Nalu::next(&mut cursor) {
nalus += 1;
match nalu.header.type_ {
NaluType::Sps => {
self.parser.parse_sps(&nalu).map_err(|e| anyhow!("parse SPS: {e}"))?;
}
NaluType::Pps => {
self.parser.parse_pps(&nalu).map_err(|e| anyhow!("parse PPS: {e}"))?;
}
NaluType::Slice
| NaluType::SliceDpa
| NaluType::SliceDpb
| NaluType::SliceDpc
| NaluType::SliceIdr
| NaluType::SliceExt => {
let slice = self
.parser
.parse_slice_header(nalu)
.map_err(|e| anyhow!("parse slice header: {e}"))?;
let mut picture = match current.take() {
Some(picture) if slice.header.first_mb_in_slice == 0 => {
self.finish_picture(picture)?;
self.begin_picture(&slice, timestamp)?
}
Some(picture) => picture,
None => self.begin_picture(&slice, timestamp)?,
};
self.decode_slice(&mut picture, &slice)?;
current = Some(picture);
}
other => log::trace!("ignoring NAL unit of type {other:?}"),
}
}
if nalus == 0 && !access_unit.is_empty() {
log::warn!(
"access unit of {} bytes holds no Annex-B NAL unit; length-prefixed input has to be converted first",
access_unit.len()
);
}
if let Some(picture) = current.take() {
self.finish_picture(picture)?;
}
Ok(())
}
pub fn flush(&mut self) -> anyhow::Result<Vec<Frame>> {
self.reset();
self.take_frames()
}
pub fn flush_exported(&mut self) -> anyhow::Result<Vec<ExportedFrame>> {
self.reset();
self.take_exported()
}
fn reset(&mut self) {
self.drain();
self.prev_ref_pic_info = PrevReferencePicInfo::default();
self.prev_pic_info = PrevPicInfo::default();
self.max_long_term_frame_idx = MaxLongTermFrameIdx::default();
}
fn take_frames(&mut self) -> anyhow::Result<Vec<Frame>> {
std::mem::take(&mut self.ready).iter().map(download).collect()
}
fn take_exported(&mut self) -> anyhow::Result<Vec<ExportedFrame>> {
std::mem::take(&mut self.ready).iter().map(export).collect()
}
fn drain(&mut self) {
self.ready.extend(self.dpb.drain().into_iter().flatten());
}
fn apply_sps(&mut self, sps: &Sps) -> anyhow::Result<()> {
let info = SequenceInfo::new(sps)?;
let max_dpb_frames = sps.max_dpb_frames();
let max_num_order_frames = sps.max_num_order_frames() as usize;
let max_num_reorder_frames = if max_num_order_frames > max_dpb_frames {
0
} else {
max_num_order_frames
};
self.dpb.set_limits(max_dpb_frames, max_num_reorder_frames);
if self.sequence.as_ref().map(|sequence| &sequence.info) == Some(&info) {
return Ok(());
}
self.drain();
self.sequence = Some(Sequence::new(&self.display, sps, info)?);
Ok(())
}
fn begin_picture(&mut self, slice: &Slice, timestamp: u64) -> anyhow::Result<CurrentPicture> {
let pps = Rc::clone(
self.parser
.get_pps(slice.header.pic_parameter_set_id)
.context("slice refers to an unknown PPS")?,
);
let sps = Rc::clone(&pps.sps);
self.apply_sps(&sps)?;
if slice.nalu.header.idr_pic_flag {
self.prev_ref_pic_info.frame_num = 0;
}
let frame_num = u32::from(slice.header.frame_num);
if frame_num != self.prev_ref_pic_info.frame_num
&& frame_num != (self.prev_ref_pic_info.frame_num + 1) % sps.max_frame_num()
{
self.handle_frame_num_gap(&sps, frame_num, timestamp)?;
}
let pic = self.init_current_pic(slice, &sps, timestamp)?;
let ref_pic_lists = self.dpb.build_ref_pic_lists(&pic);
let sequence = self.sequence.as_ref().expect("sequence built by apply_sps");
let context = Rc::clone(&sequence.context);
let handle = Handle {
surface: Arc::new(sequence.pool.alloc()?),
timestamp,
width: sequence.width,
height: sequence.height,
};
let mut picture = Picture::new(timestamp, Rc::clone(&context), handle.clone());
let pic_param = build_pic_param(&slice.header, &pic, handle.surface.id(), &self.dpb, &sps, &pps);
picture.add_buffer(create_buffer(&context, pic_param)?);
picture.add_buffer(create_buffer(&context, build_iq_matrix(&pps))?);
log::trace!("starting picture POC {}", pic.pic_order_cnt);
Ok(CurrentPicture {
pic,
pps,
handle,
picture,
ref_pic_lists,
})
}
fn decode_slice(&mut self, current: &mut CurrentPicture, slice: &Slice) -> anyhow::Result<()> {
if slice.header.first_mb_in_slice > u32::from(u16::MAX) {
bail!(
"slice starts at macroblock {}, past what a VA slice parameter addresses",
slice.header.first_mb_in_slice
);
}
if slice.header.header_bit_size > usize::from(u16::MAX) {
bail!(
"slice header is {} bits long, past what a VA slice parameter addresses",
slice.header.header_bit_size
);
}
let pps = Rc::clone(
self.parser
.get_pps(slice.header.pic_parameter_set_id)
.context("slice refers to an unknown PPS")?,
);
if SequenceInfo::new(&pps.sps)? != SequenceInfo::new(¤t.pps.sps)? {
bail!("invalid stream: the sequence changed between slices of one picture");
}
current.pps = Rc::clone(&pps);
let (list0, list1) = self.create_ref_pic_lists(¤t.pic, &slice.header, ¤t.ref_pic_lists)?;
let context = &self.sequence.as_ref().expect("sequence built by apply_sps").context;
let slice_param = build_slice_param(&slice.header, slice.nalu.size, &list0, &list1, &pps.sps, &pps);
current.picture.add_buffer(create_buffer(context, slice_param)?);
current.picture.add_buffer(create_buffer(
context,
BufferType::SliceData(slice.nalu.as_ref().to_vec()),
)?);
Ok(())
}
fn finish_picture(&mut self, current: CurrentPicture) -> anyhow::Result<()> {
let CurrentPicture {
mut pic,
pps,
handle,
picture,
..
} = current;
let picture = picture.begin::<()>().map_err(|e| anyhow!("picture begin: {e:?}"))?;
let picture = picture.render().map_err(|e| anyhow!("picture render: {e:?}"))?;
let picture = picture.end().map_err(|e| anyhow!("picture end: {e:?}"))?;
drop(picture.sync::<()>().map_err(|(e, _)| anyhow!("picture sync: {e:?}"))?);
log::trace!("decoded picture POC {}", pic.pic_order_cnt);
if pic.is_ref() {
self.reference_pic_marking(&mut pic, &pps.sps)?;
self.prev_ref_pic_info.fill(&pic);
}
self.prev_pic_info.fill(&pic);
if pic.has_mmco_5 {
self.drain();
}
self.ready.extend(self.dpb.bump_as_needed(&pic).into_iter().flatten());
if pic.is_ref() || self.dpb.has_empty_frame_buffer() {
self.dpb
.store_picture(pic.into_rc(), Some(handle))
.map_err(|e| anyhow!("store picture in DPB: {e}"))?;
} else {
self.ready.push(handle);
}
Ok(())
}
fn init_current_pic(&mut self, slice: &Slice, sps: &Sps, timestamp: u64) -> anyhow::Result<PictureData> {
let mut pic = PictureData::new_from_slice(slice, sps, timestamp, None);
self.compute_pic_order_count(&mut pic, sps)?;
if matches!(pic.is_idr, IsIdr::Yes { .. }) {
if pic.ref_pic_marking.no_output_of_prior_pics_flag {
self.dpb.clear();
} else {
self.drain();
}
}
self.dpb
.update_pic_nums(u32::from(slice.header.frame_num), sps.max_frame_num(), &pic);
Ok(pic)
}
fn handle_frame_num_gap(&mut self, sps: &Sps, frame_num: u32, timestamp: u64) -> anyhow::Result<()> {
if self.dpb.is_empty() {
return Ok(());
}
if !sps.gaps_in_frame_num_value_allowed_flag {
bail!("invalid frame_num {frame_num}: assuming unintentional loss of pictures");
}
log::debug!("frame_num gap up to {frame_num}");
let max_frame_num = sps.max_frame_num();
let mut unused = (self.prev_ref_pic_info.frame_num + 1) % max_frame_num;
while unused != frame_num {
let mut pic = PictureData::new_non_existing(unused, timestamp);
self.compute_pic_order_count(&mut pic, sps)?;
self.dpb.update_pic_nums(unused, max_frame_num, &pic);
self.dpb.sliding_window_marking(&mut pic, sps);
self.ready.extend(self.dpb.bump_as_needed(&pic).into_iter().flatten());
self.dpb
.store_picture(pic.into_rc(), None)
.map_err(|e| anyhow!("store non-existing picture in DPB: {e}"))?;
unused = (unused + 1) % max_frame_num;
}
Ok(())
}
fn compute_pic_order_count(&mut self, pic: &mut PictureData, sps: &Sps) -> anyhow::Result<()> {
match pic.pic_order_cnt_type {
0 => {
let (prev_pic_order_cnt_msb, prev_pic_order_cnt_lsb) = if matches!(pic.is_idr, IsIdr::Yes { .. }) {
(0, 0)
} else if self.prev_ref_pic_info.has_mmco_5 {
if matches!(self.prev_ref_pic_info.field, Field::Bottom) {
(0, 0)
} else {
(0, self.prev_ref_pic_info.top_field_order_cnt)
}
} else {
(
self.prev_ref_pic_info.pic_order_cnt_msb,
self.prev_ref_pic_info.pic_order_cnt_lsb,
)
};
let max_pic_order_cnt_lsb = 1 << (sps.log2_max_pic_order_cnt_lsb_minus4 + 4);
pic.pic_order_cnt_msb = if (pic.pic_order_cnt_lsb < prev_pic_order_cnt_lsb)
&& (prev_pic_order_cnt_lsb - pic.pic_order_cnt_lsb >= max_pic_order_cnt_lsb / 2)
{
prev_pic_order_cnt_msb + max_pic_order_cnt_lsb
} else if (pic.pic_order_cnt_lsb > prev_pic_order_cnt_lsb)
&& (pic.pic_order_cnt_lsb - prev_pic_order_cnt_lsb > max_pic_order_cnt_lsb / 2)
{
prev_pic_order_cnt_msb - max_pic_order_cnt_lsb
} else {
prev_pic_order_cnt_msb
};
if !matches!(pic.field, Field::Bottom) {
pic.top_field_order_cnt = pic.pic_order_cnt_msb + pic.pic_order_cnt_lsb;
}
if !matches!(pic.field, Field::Top) {
pic.bottom_field_order_cnt = if matches!(pic.field, Field::Frame) {
pic.top_field_order_cnt + pic.delta_pic_order_cnt_bottom
} else {
pic.pic_order_cnt_msb + pic.pic_order_cnt_lsb
};
}
}
1 => {
if self.prev_pic_info.has_mmco_5 {
self.prev_pic_info.frame_num_offset = 0;
}
pic.frame_num_offset = if matches!(pic.is_idr, IsIdr::Yes { .. }) {
0
} else if self.prev_pic_info.frame_num > pic.frame_num {
self.prev_pic_info.frame_num_offset + sps.max_frame_num()
} else {
self.prev_pic_info.frame_num_offset
};
let mut abs_frame_num = if sps.num_ref_frames_in_pic_order_cnt_cycle != 0 {
pic.frame_num_offset + pic.frame_num
} else {
0
};
if pic.nal_ref_idc == 0 && abs_frame_num > 0 {
abs_frame_num -= 1;
}
let mut expected_pic_order_cnt = 0;
if abs_frame_num > 0 {
if sps.num_ref_frames_in_pic_order_cnt_cycle == 0 {
bail!("invalid num_ref_frames_in_pic_order_cnt_cycle");
}
let cycle_len = u32::from(sps.num_ref_frames_in_pic_order_cnt_cycle);
let cycles = (abs_frame_num - 1) / cycle_len;
let position_in_cycle = (abs_frame_num - 1) % cycle_len;
expected_pic_order_cnt = cycles as i32 * sps.expected_delta_per_pic_order_cnt_cycle;
for i in 0..=position_in_cycle {
expected_pic_order_cnt += sps.offset_for_ref_frame[i as usize];
}
}
if pic.nal_ref_idc == 0 {
expected_pic_order_cnt += sps.offset_for_non_ref_pic;
}
match pic.field {
Field::Frame => {
pic.top_field_order_cnt = expected_pic_order_cnt + pic.delta_pic_order_cnt0;
pic.bottom_field_order_cnt =
pic.top_field_order_cnt + sps.offset_for_top_to_bottom_field + pic.delta_pic_order_cnt1;
}
Field::Top => {
pic.top_field_order_cnt = expected_pic_order_cnt + pic.delta_pic_order_cnt0;
}
Field::Bottom => {
pic.bottom_field_order_cnt =
expected_pic_order_cnt + sps.offset_for_top_to_bottom_field + pic.delta_pic_order_cnt0;
}
}
}
2 => {
if self.prev_pic_info.has_mmco_5 {
self.prev_pic_info.frame_num_offset = 0;
}
pic.frame_num_offset = if matches!(pic.is_idr, IsIdr::Yes { .. }) {
0
} else if self.prev_pic_info.frame_num > pic.frame_num {
self.prev_pic_info.frame_num_offset + sps.max_frame_num()
} else {
self.prev_pic_info.frame_num_offset
};
let pic_order_cnt = if matches!(pic.is_idr, IsIdr::Yes { .. }) {
0
} else if pic.nal_ref_idc == 0 {
2 * (pic.frame_num_offset + pic.frame_num) as i32 - 1
} else {
2 * (pic.frame_num_offset + pic.frame_num) as i32
};
if matches!(pic.field, Field::Frame | Field::Top) {
pic.top_field_order_cnt = pic_order_cnt;
}
if matches!(pic.field, Field::Frame | Field::Bottom) {
pic.bottom_field_order_cnt = pic_order_cnt;
}
}
other => bail!("invalid pic_order_cnt_type {other}"),
}
pic.pic_order_cnt = match pic.field {
Field::Frame => std::cmp::min(pic.top_field_order_cnt, pic.bottom_field_order_cnt),
Field::Top => pic.top_field_order_cnt,
Field::Bottom => pic.bottom_field_order_cnt,
};
Ok(())
}
fn reference_pic_marking(&mut self, pic: &mut PictureData, sps: &Sps) -> anyhow::Result<()> {
if matches!(pic.is_idr, IsIdr::Yes { .. }) {
self.dpb.mark_all_as_unused_for_ref();
if pic.ref_pic_marking.long_term_reference_flag {
pic.set_reference(Reference::LongTerm, false);
pic.long_term_frame_idx = 0;
self.max_long_term_frame_idx = MaxLongTermFrameIdx::Idx(0);
} else {
pic.set_reference(Reference::ShortTerm, false);
self.max_long_term_frame_idx = MaxLongTermFrameIdx::NoLongTermFrameIndices;
}
return Ok(());
}
if pic.ref_pic_marking.adaptive_ref_pic_marking_mode_flag {
self.handle_memory_management_ops(pic)?;
} else {
self.dpb.sliding_window_marking(pic, sps);
}
Ok(())
}
fn handle_memory_management_ops(&mut self, pic: &mut PictureData) -> Result<(), MmcoError> {
let markings = pic.ref_pic_marking.clone();
for marking in &markings.inner {
match marking.memory_management_control_operation {
0 => break,
1 => self.dpb.mmco_op_1(pic, marking)?,
2 => self.dpb.mmco_op_2(pic, marking)?,
3 => self.dpb.mmco_op_3(pic, marking)?,
4 => self.max_long_term_frame_idx = self.dpb.mmco_op_4(marking),
5 => self.max_long_term_frame_idx = self.dpb.mmco_op_5(pic),
6 => self.dpb.mmco_op_6(pic, marking),
other => return Err(MmcoError::UnknownMmco(other)),
}
}
Ok(())
}
fn create_ref_pic_lists<'a>(
&'a self,
pic: &PictureData,
hdr: &SliceHeader,
lists: &ReferencePicLists,
) -> anyhow::Result<(DpbPicRefList<'a, Handle>, DpbPicRefList<'a, Handle>)> {
let list0 = match hdr.slice_type {
SliceType::P | SliceType::Sp => self.modify_ref_pic_list(pic, hdr, false, &lists.ref_pic_list_p0)?,
SliceType::B => self.modify_ref_pic_list(pic, hdr, false, &lists.ref_pic_list_b0)?,
_ => Vec::new(),
};
let list1 = match hdr.slice_type {
SliceType::B => self.modify_ref_pic_list(pic, hdr, true, &lists.ref_pic_list_b1)?,
_ => Vec::new(),
};
Ok((list0, list1))
}
fn modify_ref_pic_list<'a>(
&'a self,
pic: &PictureData,
hdr: &SliceHeader,
list1: bool,
indices: &[usize],
) -> anyhow::Result<DpbPicRefList<'a, Handle>> {
let (modify, num_ref_idx_active_minus1, modifications) = if list1 {
(
hdr.ref_pic_list_modification_flag_l1,
hdr.num_ref_idx_l1_active_minus1,
&hdr.ref_pic_list_modification_l1,
)
} else {
(
hdr.ref_pic_list_modification_flag_l0,
hdr.num_ref_idx_l0_active_minus1,
&hdr.ref_pic_list_modification_l0,
)
};
let mut list: DpbPicRefList<'a, Handle> = indices
.iter()
.map(|&i| &self.dpb.entries()[i])
.take(usize::from(num_ref_idx_active_minus1) + 1)
.collect();
if !modify {
return Ok(list);
}
let mut pic_num_pred = pic.pic_num;
let mut ref_idx = 0;
for modification in modifications {
match modification.modification_of_pic_nums_idc {
idc @ (0 | 1) => self.short_term_modification(
pic,
&mut list,
num_ref_idx_active_minus1,
hdr.max_pic_num as i32,
idc,
modification,
&mut pic_num_pred,
&mut ref_idx,
)?,
2 => self.long_term_modification(&mut list, num_ref_idx_active_minus1, modification, &mut ref_idx)?,
3 => break,
other => bail!("unexpected modification_of_pic_nums_idc {other}"),
}
}
Ok(list)
}
#[allow(clippy::too_many_arguments)]
fn short_term_modification<'a>(
&'a self,
pic: &PictureData,
list: &mut DpbPicRefList<'a, Handle>,
num_ref_idx_active_minus1: u8,
max_pic_num: i32,
idc: u8,
modification: &RefPicListModification,
pic_num_pred: &mut i32,
ref_idx: &mut usize,
) -> anyhow::Result<()> {
let abs_diff_pic_num = modification.abs_diff_pic_num_minus1 as i32 + 1;
let pic_num_no_wrap = if idc == 0 {
if *pic_num_pred - abs_diff_pic_num < 0 {
*pic_num_pred - abs_diff_pic_num + max_pic_num
} else {
*pic_num_pred - abs_diff_pic_num
}
} else if *pic_num_pred + abs_diff_pic_num >= max_pic_num {
*pic_num_pred + abs_diff_pic_num - max_pic_num
} else {
*pic_num_pred + abs_diff_pic_num
};
*pic_num_pred = pic_num_no_wrap;
let pic_num = if pic_num_no_wrap > pic.pic_num {
pic_num_no_wrap - max_pic_num
} else {
pic_num_no_wrap
};
let entry = self
.dpb
.find_short_term_with_pic_num(pic_num)
.with_context(|| format!("no short-term reference with pic_num {pic_num}"))?;
Self::insert_modification(list, entry, num_ref_idx_active_minus1, ref_idx, |target| {
target.pic_num_f(max_pic_num) != pic_num
});
Ok(())
}
fn long_term_modification<'a>(
&'a self,
list: &mut DpbPicRefList<'a, Handle>,
num_ref_idx_active_minus1: u8,
modification: &RefPicListModification,
ref_idx: &mut usize,
) -> anyhow::Result<()> {
let long_term_pic_num = modification.long_term_pic_num;
let max_long_term_frame_idx = self.max_long_term_frame_idx;
let entry = self
.dpb
.find_long_term_with_long_term_pic_num(long_term_pic_num)
.with_context(|| format!("no long-term reference with long_term_pic_num {long_term_pic_num}"))?;
Self::insert_modification(list, entry, num_ref_idx_active_minus1, ref_idx, |target| {
target.long_term_pic_num_f(max_long_term_frame_idx) != long_term_pic_num
});
Ok(())
}
fn insert_modification<'a>(
list: &mut DpbPicRefList<'a, Handle>,
entry: &'a DpbEntry<Handle>,
num_ref_idx_active_minus1: u8,
ref_idx: &mut usize,
keep: impl Fn(&PictureData) -> bool,
) {
let at = (*ref_idx).min(list.len());
list.insert(at, entry);
*ref_idx = at + 1;
let mut next = *ref_idx;
for current in *ref_idx..=usize::from(num_ref_idx_active_minus1) + 1 {
if current >= list.len() {
break;
}
if keep(&list[current].pic.borrow()) {
list[next] = list[current];
next += 1;
}
}
list.truncate(usize::from(num_ref_idx_active_minus1) + 1);
}
}
struct CurrentPicture {
pic: PictureData,
pps: Rc<Pps>,
handle: Handle,
picture: Picture<crate::PictureNew, Handle>,
ref_pic_lists: ReferencePicLists,
}
#[derive(Clone)]
struct Handle {
surface: Arc<PooledSurface>,
timestamp: u64,
width: u32,
height: u32,
}
impl std::borrow::Borrow<Surface<()>> for Handle {
fn borrow(&self) -> &Surface<()> {
self.surface.get()
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
struct SequenceInfo {
visible: (u32, u32),
coded: (u32, u32),
profile_idc: u8,
bit_depth_luma_minus8: u8,
bit_depth_chroma_minus8: u8,
chroma_format_idc: u8,
frame_mbs_only_flag: bool,
max_dpb_frames: usize,
}
impl SequenceInfo {
fn new(sps: &Sps) -> anyhow::Result<Self> {
if !sps.frame_mbs_only_flag {
bail!("interlaced H.264 is not supported");
}
if sps.chroma_format_idc != 1 {
bail!(
"only 4:2:0 chroma is supported, got chroma_format_idc {}",
sps.chroma_format_idc
);
}
if sps.bit_depth_luma_minus8 != 0 || sps.bit_depth_chroma_minus8 != 0 {
bail!(
"only 8-bit samples are supported, got {}-bit luma and {}-bit chroma",
sps.bit_depth_luma_minus8 + 8,
sps.bit_depth_chroma_minus8 + 8
);
}
if sps.frame_cropping_flag && (sps.frame_crop_left_offset != 0 || sps.frame_crop_top_offset != 0) {
bail!("left and top H.264 cropping is not supported");
}
let (crop_right, crop_bottom) = if sps.frame_cropping_flag {
(sps.frame_crop_right_offset, sps.frame_crop_bottom_offset)
} else {
(0, 0)
};
let width = crop_right
.checked_mul(2)
.and_then(|crop| sps.width().checked_sub(crop))
.context("invalid horizontal H.264 crop")?;
let height = crop_bottom
.checked_mul(2)
.and_then(|crop| sps.height().checked_sub(crop))
.context("invalid vertical H.264 crop")?;
if width == 0 || height == 0 || width % 2 != 0 || height % 2 != 0 {
bail!("visible size {width}x{height} is not a non-zero even 4:2:0 size");
}
Ok(Self {
visible: (width, height),
coded: (sps.width(), sps.height()),
profile_idc: sps.profile_idc,
bit_depth_luma_minus8: sps.bit_depth_luma_minus8,
bit_depth_chroma_minus8: sps.bit_depth_chroma_minus8,
chroma_format_idc: sps.chroma_format_idc,
frame_mbs_only_flag: sps.frame_mbs_only_flag,
max_dpb_frames: sps.max_dpb_frames(),
})
}
}
struct Sequence {
info: SequenceInfo,
width: u32,
height: u32,
pool: SurfacePool,
context: Rc<Context>,
_config: VaConfig,
}
impl Sequence {
fn new(display: &Arc<Display>, sps: &Sps, info: SequenceInfo) -> anyhow::Result<Self> {
let (width, height) = info.visible;
let attrs = vec![VAConfigAttrib {
type_: VAConfigAttribType::VAConfigAttribRTFormat,
value: VA_RT_FORMAT_YUV420,
}];
let config = display
.create_config(attrs, va_profile(sps)?, VAEntrypoint::VAEntrypointVLD)
.map_err(|e| anyhow!("create VA config: {e:?}"))?;
let (coded_width, coded_height) = info.coded;
let context = display
.create_context::<()>(&config, coded_width, coded_height, None, true)
.map_err(|e| anyhow!("create VA context: {e:?}"))?;
log::info!(
"VA-API H.264 sequence: {width}x{height} visible, {coded_width}x{coded_height} coded, \
profile_idc {}, DPB {} frames",
sps.profile_idc,
info.max_dpb_frames
);
let capacity = info.max_dpb_frames + 1;
Ok(Self {
info,
width,
height,
pool: SurfacePool {
display: Arc::clone(display),
width: coded_width,
height: coded_height,
free: Arc::new(FreeList {
surfaces: Mutex::new(Vec::new()),
capacity,
}),
allocated: Cell::new(0),
},
context,
_config: config,
})
}
}
struct SurfacePool {
display: Arc<Display>,
width: u32,
height: u32,
free: Arc<FreeList>,
allocated: Cell<usize>,
}
struct FreeList {
surfaces: Mutex<Vec<Surface<()>>>,
capacity: usize,
}
impl SurfacePool {
fn alloc(&self) -> anyhow::Result<PooledSurface> {
let recycled = self.free.surfaces.lock().expect("poisoned").pop();
let surface = match recycled {
Some(surface) => surface,
None => {
let surface = self
.display
.create_surfaces::<()>(
VA_RT_FORMAT_YUV420,
Some(VA_FOURCC_NV12),
self.width,
self.height,
Some(UsageHint::USAGE_HINT_DECODER),
vec![()],
)
.map_err(|e| anyhow!("create decode surface: {e:?}"))?
.pop()
.context("VA-API returned an empty surface list")?;
self.allocated.set(self.allocated.get() + 1);
log::debug!(
"allocated decode surface {} at {}x{}",
self.allocated.get(),
self.width,
self.height
);
surface
}
};
Ok(PooledSurface {
surface: Some(surface),
free: Arc::downgrade(&self.free),
})
}
}
struct PooledSurface {
surface: Option<Surface<()>>,
free: Weak<FreeList>,
}
impl PooledSurface {
fn unpooled(surface: Surface<()>) -> Self {
Self {
surface: Some(surface),
free: Weak::new(),
}
}
fn get(&self) -> &Surface<()> {
self.surface.as_ref().expect("surface is taken only on drop")
}
fn id(&self) -> bindings::VASurfaceID {
self.get().id()
}
}
impl Drop for PooledSurface {
fn drop(&mut self) {
let (Some(surface), Some(free)) = (self.surface.take(), self.free.upgrade()) else {
return;
};
let mut surfaces = free.surfaces.lock().expect("poisoned");
if surfaces.len() < free.capacity {
surfaces.push(surface);
}
}
}
fn export(handle: &Handle) -> anyhow::Result<ExportedFrame> {
let surface = handle.surface.get();
surface.sync().map_err(|e| anyhow!("surface sync: {e:?}"))?;
let descriptor = surface
.export_prime()
.map_err(|e| anyhow!("export the decode surface: {e:?}"))?;
Ok(ExportedFrame {
timestamp: handle.timestamp,
width: handle.width,
height: handle.height,
descriptor,
surface: Arc::clone(&handle.surface),
})
}
fn download(handle: &Handle) -> anyhow::Result<Frame> {
read_back(handle.surface.get(), handle.timestamp, handle.width, handle.height)
}
fn read_back(surface: &Surface<()>, timestamp: u64, width: u32, height: u32) -> anyhow::Result<Frame> {
surface.sync().map_err(|e| anyhow!("surface sync: {e:?}"))?;
let visible = (width, height);
match Image::derive_from(surface, visible) {
Ok(image) if image.image().format.fourcc == VA_FOURCC_NV12 => {
return pack_nv12(&image, timestamp, width, height);
}
Ok(_) => log::debug!("derived image is not NV12, falling back to vaGetImage"),
Err(e) => log::debug!("vaDeriveImage failed ({e:?}), falling back to vaGetImage"),
}
let format = surface
.display()
.query_image_formats()
.map_err(|e| anyhow!("query image formats: {e:?}"))?
.into_iter()
.find(|format| format.fourcc == VA_FOURCC_NV12)
.context("driver has no NV12 image format")?;
let image = Image::create_from(surface, format, surface.size(), visible)
.map_err(|e| anyhow!("vaGetImage into an NV12 image: {e:?}"))?;
pack_nv12(&image, timestamp, width, height)
}
fn pack_nv12(image: &Image<'_>, timestamp: u64, frame_width: u32, frame_height: u32) -> anyhow::Result<Frame> {
let va_image = image.image();
let src: &[u8] = image.as_ref();
let (width, height) = (frame_width as usize, frame_height as usize);
let chroma_rows = height / 2;
let mut data = vec![0u8; width * (height + chroma_rows)];
let (luma, chroma) = data.split_at_mut(width * height);
copy_plane(
luma,
src,
va_image.offsets[0] as usize,
va_image.pitches[0] as usize,
width,
height,
)?;
copy_plane(
chroma,
src,
va_image.offsets[1] as usize,
va_image.pitches[1] as usize,
width,
chroma_rows,
)?;
Ok(Frame {
timestamp,
width: frame_width,
height: frame_height,
data,
})
}
fn copy_plane(
dst: &mut [u8],
src: &[u8],
offset: usize,
pitch: usize,
width: usize,
rows: usize,
) -> anyhow::Result<()> {
if rows == 0 {
return Ok(());
}
let end = offset + pitch * (rows - 1) + width;
if pitch < width || src.len() < end {
bail!("VA image plane (offset {offset}, pitch {pitch}) does not hold {rows} rows of {width} bytes");
}
for row in 0..rows {
let from = offset + row * pitch;
dst[row * width..][..width].copy_from_slice(&src[from..from + width]);
}
Ok(())
}
pub fn probe(display: &Display) -> anyhow::Result<()> {
let profiles = display
.query_config_profiles()
.map_err(|e| anyhow!("query VA profiles: {e:?}"))?;
let supported = profiles
.into_iter()
.filter(|profile| {
matches!(
*profile,
VAProfile::VAProfileH264ConstrainedBaseline
| VAProfile::VAProfileH264Main
| VAProfile::VAProfileH264High
)
})
.any(|profile| {
display
.query_config_entrypoints(profile)
.is_ok_and(|entrypoints| entrypoints.contains(&VAEntrypoint::VAEntrypointVLD))
});
if !supported {
bail!("driver exposes no H.264 decode entrypoint");
}
Ok(())
}
fn va_profile(sps: &Sps) -> anyhow::Result<VAProfile::Type> {
let profile = Profile::try_from(sps.profile_idc).map_err(|e| anyhow!("{e}"))?;
Ok(match profile {
Profile::Baseline => VAProfile::VAProfileH264ConstrainedBaseline,
Profile::Main | Profile::Extended => VAProfile::VAProfileH264Main,
Profile::High | Profile::High10 | Profile::High422P => VAProfile::VAProfileH264High,
})
}
fn create_buffer(context: &Rc<Context>, buffer: BufferType) -> anyhow::Result<Buffer> {
context
.create_buffer(buffer)
.map_err(|e| anyhow!("create VA buffer: {e:?}"))
}
#[derive(Debug)]
struct PrevReferencePicInfo {
frame_num: u32,
has_mmco_5: bool,
top_field_order_cnt: i32,
pic_order_cnt_msb: i32,
pic_order_cnt_lsb: i32,
field: Field,
}
impl Default for PrevReferencePicInfo {
fn default() -> Self {
Self {
frame_num: 0,
has_mmco_5: false,
top_field_order_cnt: 0,
pic_order_cnt_msb: 0,
pic_order_cnt_lsb: 0,
field: Field::Frame,
}
}
}
impl PrevReferencePicInfo {
fn fill(&mut self, pic: &PictureData) {
self.has_mmco_5 = pic.has_mmco_5;
self.top_field_order_cnt = pic.top_field_order_cnt;
self.pic_order_cnt_msb = pic.pic_order_cnt_msb;
self.pic_order_cnt_lsb = pic.pic_order_cnt_lsb;
self.field = pic.field;
self.frame_num = pic.frame_num;
}
}
#[derive(Debug, Default)]
struct PrevPicInfo {
frame_num: u32,
frame_num_offset: u32,
has_mmco_5: bool,
}
impl PrevPicInfo {
fn fill(&mut self, pic: &PictureData) {
self.frame_num = pic.frame_num;
self.has_mmco_5 = pic.has_mmco_5;
self.frame_num_offset = pic.frame_num_offset;
}
}
fn va_surface_id(handle: &Option<Handle>) -> bindings::VASurfaceID {
match handle {
Some(handle) => handle.surface.id(),
None => VA_INVALID_ID,
}
}
fn fill_va_h264_pic(pic: &PictureData, surface_id: bindings::VASurfaceID) -> PictureH264 {
let mut flags = 0;
let frame_idx = match pic.reference() {
Reference::LongTerm => {
flags |= VA_PICTURE_H264_LONG_TERM_REFERENCE;
pic.long_term_frame_idx
}
Reference::ShortTerm => {
flags |= VA_PICTURE_H264_SHORT_TERM_REFERENCE;
pic.frame_num
}
Reference::None => pic.frame_num,
};
PictureH264::new(
surface_id,
frame_idx,
flags,
pic.top_field_order_cnt,
pic.bottom_field_order_cnt,
)
}
fn build_invalid_va_h264_pic() -> PictureH264 {
PictureH264::new(VA_INVALID_ID, 0, VA_PICTURE_H264_INVALID, 0, 0)
}
fn build_iq_matrix(pps: &Pps) -> BufferType {
const ZIGZAG_4X4: [usize; 16] = [0, 1, 4, 8, 5, 2, 3, 6, 9, 12, 13, 10, 7, 11, 14, 15];
#[rustfmt::skip]
const ZIGZAG_8X8: [usize; 64] = [
0, 1, 8, 16, 9, 2, 3, 10, 17, 24, 32, 25, 18, 11, 4, 5, 12, 19, 26, 33, 40, 48, 41, 34, 27,
20, 13, 6, 7, 14, 21, 28, 35, 42, 49, 56, 57, 50, 43, 36, 29, 22, 15, 23, 30, 37, 44, 51,
58, 59, 52, 45, 38, 31, 39, 46, 53, 60, 61, 54, 47, 55, 62, 63,
];
let mut scaling_list4x4 = [[0u8; 16]; 6];
for (list, src) in scaling_list4x4.iter_mut().zip(pps.scaling_lists_4x4.iter()) {
for (i, &value) in src.iter().enumerate() {
list[ZIGZAG_4X4[i]] = value;
}
}
let mut scaling_list8x8 = [[0u8; 64]; 2];
for (list, src) in scaling_list8x8.iter_mut().zip(pps.scaling_lists_8x8.iter()) {
for (i, &value) in src.iter().enumerate() {
list[ZIGZAG_8X8[i]] = value;
}
}
BufferType::IQMatrix(IQMatrix::H264(IQMatrixBufferH264::new(
scaling_list4x4,
scaling_list8x8,
)))
}
fn build_pic_param(
hdr: &SliceHeader,
pic: &PictureData,
surface_id: bindings::VASurfaceID,
dpb: &Dpb<Handle>,
sps: &Sps,
pps: &Pps,
) -> BufferType {
let curr_pic = fill_va_h264_pic(pic, surface_id);
let short_term = dpb
.short_term_refs_iter()
.filter(|entry| !entry.pic.borrow().nonexisting);
let mut refs: Vec<PictureH264> = short_term
.chain(dpb.long_term_refs_iter())
.take(16)
.map(|entry| fill_va_h264_pic(&entry.pic.borrow(), va_surface_id(&entry.reference)))
.collect();
refs.resize_with(16, build_invalid_va_h264_pic);
let refs: [PictureH264; 16] = refs.try_into().unwrap_or_else(|_| unreachable!("resized to 16"));
let seq_fields = H264SeqFields::new(
sps.chroma_format_idc as u32,
sps.separate_colour_plane_flag as u32,
sps.gaps_in_frame_num_value_allowed_flag as u32,
sps.frame_mbs_only_flag as u32,
sps.mb_adaptive_frame_field_flag as u32,
sps.direct_8x8_inference_flag as u32,
(sps.level_idc >= Level::L3_1) as u32,
sps.log2_max_frame_num_minus4 as u32,
sps.pic_order_cnt_type as u32,
sps.log2_max_pic_order_cnt_lsb_minus4 as u32,
sps.delta_pic_order_always_zero_flag as u32,
);
let pic_fields = H264PicFields::new(
pps.entropy_coding_mode_flag as u32,
pps.weighted_pred_flag as u32,
pps.weighted_bipred_idc as u32,
pps.transform_8x8_mode_flag as u32,
hdr.field_pic_flag as u32,
pps.constrained_intra_pred_flag as u32,
pps.bottom_field_pic_order_in_frame_present_flag as u32,
pps.deblocking_filter_control_present_flag as u32,
pps.redundant_pic_cnt_present_flag as u32,
(pic.nal_ref_idc != 0) as u32,
);
let picture_height_in_mbs_minus1 = sps.pic_height_in_map_units_minus1;
BufferType::PictureParameter(PictureParameter::H264(PictureParameterBufferH264::new(
curr_pic,
refs,
sps.pic_width_in_mbs_minus1,
picture_height_in_mbs_minus1,
sps.bit_depth_luma_minus8,
sps.bit_depth_chroma_minus8,
sps.max_num_ref_frames,
&seq_fields,
0,
0,
0,
pps.pic_init_qp_minus26,
pps.pic_init_qs_minus26,
pps.chroma_qp_index_offset,
pps.second_chroma_qp_index_offset,
&pic_fields,
hdr.frame_num,
)))
}
fn fill_ref_pic_list(list: &[&DpbEntry<Handle>]) -> [PictureH264; 32] {
let mut pics: Vec<PictureH264> = list
.iter()
.take(32)
.map(|entry| fill_va_h264_pic(&entry.pic.borrow(), va_surface_id(&entry.reference)))
.collect();
pics.resize_with(32, build_invalid_va_h264_pic);
pics.try_into().unwrap_or_else(|_| unreachable!("resized to 32"))
}
fn build_slice_param(
hdr: &SliceHeader,
slice_size: usize,
list0: &[&DpbEntry<Handle>],
list1: &[&DpbEntry<Handle>],
sps: &Sps,
pps: &Pps,
) -> BufferType {
let pwt = &hdr.pred_weight_table;
let mut luma_weight_l0 = [0i16; 32];
let mut luma_offset_l0 = [0i16; 32];
let mut chroma_weight_l0 = [[0i16; 2]; 32];
let mut chroma_offset_l0 = [[0i16; 2]; 32];
let mut luma_weight_l1 = [0i16; 32];
let mut luma_offset_l1 = [0i16; 32];
let mut chroma_weight_l1 = [[0i16; 2]; 32];
let mut chroma_offset_l1 = [[0i16; 2]; 32];
let fill_l0 = (pps.weighted_pred_flag && (hdr.slice_type.is_p() || hdr.slice_type.is_sp()))
|| (pps.weighted_bipred_idc == 1 && hdr.slice_type.is_b());
let fill_l1 = pps.weighted_bipred_idc == 1 && hdr.slice_type.is_b();
let chroma = sps.chroma_array_type() != 0;
if fill_l0 {
for i in 0..=usize::from(hdr.num_ref_idx_l0_active_minus1) {
luma_weight_l0[i] = pwt.luma_weight_l0[i];
luma_offset_l0[i] = i16::from(pwt.luma_offset_l0[i]);
if chroma {
for j in 0..2 {
chroma_weight_l0[i][j] = pwt.chroma_weight_l0[i][j];
chroma_offset_l0[i][j] = i16::from(pwt.chroma_offset_l0[i][j]);
}
}
}
}
if fill_l1 {
for i in 0..=usize::from(hdr.num_ref_idx_l1_active_minus1) {
luma_weight_l1[i] = pwt.luma_weight_l1[i];
luma_offset_l1[i] = i16::from(pwt.luma_offset_l1[i]);
if chroma {
for j in 0..2 {
chroma_weight_l1[i][j] = pwt.chroma_weight_l1[i][j];
chroma_offset_l1[i][j] = i16::from(pwt.chroma_offset_l1[i][j]);
}
}
}
}
BufferType::SliceParameter(SliceParameter::H264(SliceParameterBufferH264::new(
slice_size as u32,
0,
VA_SLICE_DATA_FLAG_ALL,
hdr.header_bit_size as u16,
hdr.first_mb_in_slice as u16,
hdr.slice_type as u8,
hdr.direct_spatial_mv_pred_flag as u8,
hdr.num_ref_idx_l0_active_minus1,
hdr.num_ref_idx_l1_active_minus1,
hdr.cabac_init_idc,
hdr.slice_qp_delta,
hdr.disable_deblocking_filter_idc,
hdr.slice_alpha_c0_offset_div2,
hdr.slice_beta_offset_div2,
fill_ref_pic_list(list0),
fill_ref_pic_list(list1),
pwt.luma_log2_weight_denom,
pwt.chroma_log2_weight_denom,
fill_l0 as u8,
luma_weight_l0,
luma_offset_l0,
(fill_l0 && chroma) as u8,
chroma_weight_l0,
chroma_offset_l0,
fill_l1 as u8,
luma_weight_l1,
luma_offset_l1,
(fill_l1 && chroma) as u8,
chroma_weight_l1,
chroma_offset_l1,
)))
}
#[cfg(test)]
mod tests {
use std::os::fd::OwnedFd;
use std::process::Command;
use super::*;
fn progressive_sps() -> Sps {
Sps {
frame_mbs_only_flag: true,
chroma_format_idc: 1,
pic_width_in_mbs_minus1: 19,
pic_height_in_map_units_minus1: 14,
..Default::default()
}
}
#[test]
fn rejects_nonzero_crop_origins() {
for (left, top) in [(1, 0), (0, 1)] {
let sps = Sps {
frame_cropping_flag: true,
frame_crop_left_offset: left,
frame_crop_top_offset: top,
..progressive_sps()
};
assert!(SequenceInfo::new(&sps)
.unwrap_err()
.to_string()
.contains("left and top"));
}
}
#[test]
fn crop_only_updates_change_sequence_identity() {
let original = SequenceInfo::new(&progressive_sps()).unwrap();
for (right, bottom, expected) in [(1, 0, (318, 240)), (0, 1, (320, 238))] {
let cropped = SequenceInfo::new(&Sps {
frame_cropping_flag: true,
frame_crop_right_offset: right,
frame_crop_bottom_offset: bottom,
..progressive_sps()
})
.unwrap();
assert_eq!(original.coded, cropped.coded);
assert_eq!(cropped.visible, expected);
assert_ne!(original, cropped);
}
}
#[test]
fn rejects_invalid_crop_dimensions() {
for crop in [160, 161, u32::MAX] {
assert!(SequenceInfo::new(&Sps {
frame_cropping_flag: true,
frame_crop_right_offset: crop,
..progressive_sps()
})
.is_err());
}
}
#[test]
fn maximum_progressive_height_does_not_overflow_va_parameters() {
let sps = Rc::new(Sps {
pic_height_in_map_units_minus1: u16::MAX,
..progressive_sps()
});
let params = build_pic_param(
&SliceHeader::default(),
&PictureData::new_non_existing(0, 0),
VA_INVALID_ID,
&Dpb::default(),
&sps,
&crate::codec::h264::parser::PpsBuilder::new(Rc::clone(&sps)).build(),
);
let BufferType::PictureParameter(PictureParameter::H264(params)) = params else {
panic!("expected H.264 picture parameters");
};
assert_eq!(params.inner().picture_height_in_mbs_minus1, u16::MAX);
assert!(SequenceInfo::new(&Sps {
frame_mbs_only_flag: false,
..progressive_sps()
})
.is_err());
}
fn encoder_config(width: u32, height: u32) -> crate::encode::Config {
crate::encode::Config::new(width, height, 30, 2_000_000, 30)
}
fn nv12_frame(width: u32, height: u32, step: u32) -> Vec<u8> {
let (w, h) = (width as usize, height as usize);
let mut data = vec![128u8; w * h * 3 / 2];
let x0 = (step * 16) as usize % (w / 2);
for y in h / 4..h / 2 {
for x in x0..x0 + w / 4 {
data[y * w + x] = 235;
}
}
data
}
#[test]
fn exported_pictures_do_not_share_a_surface() {
let (width, height) = (320u32, 240u32);
let Ok(mut encoder) = crate::encode::Encoder::new(encoder_config(width, height)) else {
eprintln!("skipping: no VA-API H.264 encoder");
return;
};
let Ok(mut decoder) = Decoder::new(Config::new()) else {
eprintln!("skipping: no VA-API H.264 decoder");
return;
};
let mut exported = Vec::new();
for step in 0..8 {
let unit = encoder
.encode_nv12(&nv12_frame(width, height, step), step == 0)
.expect("encode a picture");
exported.extend(decoder.decode_exported(&unit, step as u64).expect("decode a picture"));
}
exported.extend(decoder.flush_exported().expect("flush the decoder"));
assert!(!exported.is_empty(), "the decoder produced no pictures");
let first = &exported[0];
eprintln!(
"exported {} pictures: {}x{} fourcc {:#x} modifier {:#x}",
exported.len(),
first.width,
first.height,
first.descriptor.fourcc,
first.descriptor.objects[0].drm_format_modifier
);
let mut seen = Vec::new();
for (index, frame) in exported.iter().enumerate() {
assert_eq!((frame.width, frame.height), (width, height));
assert_eq!(frame.descriptor.fourcc, VA_FOURCC_NV12);
assert_eq!(frame.descriptor.objects.len(), 1, "a composed export is one object");
assert_eq!(frame.descriptor.layers[0].num_planes, 2, "NV12 is two planes");
let inode = inode(&frame.descriptor.objects[0].fd);
assert!(
!seen.contains(&inode),
"picture {index} was decoded into an allocation an earlier one still holds"
);
seen.push(inode);
}
}
#[test]
fn the_render_node_is_found_or_named() {
if let Err(err) = Decoder::new(Config::new()) {
eprintln!("skipping: no VA-API H.264 decoder: {err:#}");
return;
}
let named = Config {
device: Some(PathBuf::from("/dev/null")),
};
assert!(
Decoder::new(named).is_err(),
"a named node that is not a render node was swapped for one that is"
);
}
#[test]
fn flush_returns_the_pictures_the_dpb_still_holds() {
const PICTURES: u64 = 5;
let (width, height) = (320u32, 240u32);
let Ok(mut encoder) = crate::encode::Encoder::new(encoder_config(width, height)) else {
eprintln!("skipping: no VA-API H.264 encoder");
return;
};
let Ok(mut decoder) = Decoder::new(Config::new()) else {
eprintln!("skipping: no VA-API H.264 decoder");
return;
};
let mut streamed = Vec::new();
for step in 0..PICTURES {
let unit = encoder
.encode_nv12(&nv12_frame(width, height, step as u32), step == 0)
.expect("encode a picture");
streamed.extend(decoder.decode(&unit, step).expect("decode a picture"));
}
assert!(
(streamed.len() as u64) < PICTURES,
"the DPB held nothing back, so this test proves nothing"
);
let flushed = decoder.flush().expect("flush the decoder");
let timestamps: Vec<u64> = streamed
.iter()
.chain(&flushed)
.map(|picture| picture.timestamp)
.collect();
assert_eq!(
timestamps,
(0..PICTURES).collect::<Vec<_>>(),
"the stream lost pictures at its end"
);
assert!(decoder.flush().expect("flush an empty decoder").is_empty());
}
#[test]
fn a_dropped_export_is_decoded_into_again() {
const PICTURES: u32 = 60;
let (width, height) = (320u32, 240u32);
let Ok(mut encoder) = crate::encode::Encoder::new(encoder_config(width, height)) else {
eprintln!("skipping: no VA-API H.264 encoder");
return;
};
let Ok(mut exporting) = Decoder::new(Config::new()) else {
eprintln!("skipping: no VA-API H.264 decoder");
return;
};
let mut downloading = Decoder::new(Config::new()).expect("a second decoder");
let mut compared = 0;
for step in 0..PICTURES {
let unit = encoder
.encode_nv12(&nv12_frame(width, height, step), step == 0)
.expect("encode a picture");
let exported = exporting
.decode_exported(&unit, step as u64)
.expect("decode to a descriptor");
let downloaded = downloading.decode(&unit, step as u64).expect("decode to client memory");
assert_eq!(exported.len(), downloaded.len(), "the two decoders disagreed");
for (gpu, cpu) in exported.iter().zip(&downloaded) {
let read_back = gpu.download().expect("read the exported surface back");
assert!(read_back.data == cpu.data, "picture {} came back wrong", gpu.timestamp);
compared += 1;
}
}
assert!(compared > 0, "the decoder produced no pictures");
let pool = &exporting.sequence.as_ref().expect("a sequence").pool;
let allocated = pool.allocated.get();
eprintln!("{allocated} surfaces allocated for {PICTURES} exported pictures");
assert!(
allocated <= pool.free.capacity,
"{allocated} surfaces for {PICTURES} pictures; a full DPB and the picture in flight need {}",
pool.free.capacity
);
}
#[test]
fn the_pool_keeps_a_bounded_number_of_free_surfaces() {
const HELD: u32 = 40;
let (width, height) = (320u32, 240u32);
let Ok(mut encoder) = crate::encode::Encoder::new(encoder_config(width, height)) else {
eprintln!("skipping: no VA-API H.264 encoder");
return;
};
let Ok(mut decoder) = Decoder::new(Config::new()) else {
eprintln!("skipping: no VA-API H.264 decoder");
return;
};
let mut step = 0;
let mut decode = |decoder: &mut Decoder| {
let unit = encoder
.encode_nv12(&nv12_frame(width, height, step), step == 0)
.expect("encode a picture");
step += 1;
decoder.decode_exported(&unit, step as u64).expect("decode a picture")
};
let mut held = Vec::new();
for _ in 0..HELD {
held.extend(decode(&mut decoder));
}
let pool = &decoder.sequence.as_ref().expect("a sequence").pool;
let allocated = pool.allocated.get();
assert!(
allocated >= held.len(),
"{} pictures held on {allocated} surfaces",
held.len()
);
drop(held);
let free = pool.free.surfaces.lock().unwrap().len();
assert!(
free <= pool.free.capacity,
"{free} surfaces left free, more than the capacity of {}",
pool.free.capacity
);
for _ in 0..10 {
drop(decode(&mut decoder));
}
let pool = &decoder.sequence.as_ref().expect("a sequence").pool;
assert_eq!(
pool.allocated.get(),
allocated,
"pictures after the release did not reuse the free surfaces"
);
}
#[test]
fn an_exported_frame_is_send_and_sync() {
const fn assert_send_sync<T: Send + Sync>() {}
assert_send_sync::<ExportedFrame>();
}
#[test]
fn an_exported_picture_downloads_to_the_same_pixels() {
let (width, height) = (320u32, 240u32);
let Ok(mut encoder) = crate::encode::Encoder::new(encoder_config(width, height)) else {
eprintln!("skipping: no VA-API H.264 encoder");
return;
};
let Ok(mut exporting) = Decoder::new(Config::new()) else {
eprintln!("skipping: no VA-API H.264 decoder");
return;
};
let mut downloading = Decoder::new(Config::new()).expect("a second decoder");
let mut exported = Vec::new();
let mut downloaded = Vec::new();
for step in 0..8 {
let unit = encoder
.encode_nv12(&nv12_frame(width, height, step), step == 0)
.expect("encode a picture");
exported.extend(
exporting
.decode_exported(&unit, step as u64)
.expect("decode to a descriptor"),
);
downloaded.extend(downloading.decode(&unit, step as u64).expect("decode to client memory"));
}
exported.extend(exporting.flush_exported().expect("flush the exporting decoder"));
downloaded.extend(downloading.flush().expect("flush the downloading decoder"));
assert!(!exported.is_empty(), "the decoder produced no pictures");
assert_eq!(exported.len(), downloaded.len(), "the two decoders disagreed");
for (index, (gpu, cpu)) in exported.iter().zip(&downloaded).enumerate() {
let read_back = gpu.download().expect("read the exported surface back");
assert_eq!(read_back.timestamp, cpu.timestamp, "picture {index} lost its timestamp");
assert_eq!((read_back.width, read_back.height), (cpu.width, cpu.height));
assert!(
read_back.data == cpu.data,
"picture {index} read back differently than it downloaded"
);
}
}
#[test]
fn decoded_pictures_match_a_software_decoder() {
const WIDTH: u32 = 320;
const HEIGHT: u32 = 240;
const PICTURES: usize = 30;
let Some(stream) = ffmpeg(&[
"-f",
"lavfi",
"-i",
&format!("testsrc2=size={WIDTH}x{HEIGHT}:rate=30"),
"-frames:v",
&PICTURES.to_string(),
"-c:v",
"libx264",
"-profile:v",
"main",
"-bf",
"2",
"-refs",
"3",
"-g",
"15",
"-pix_fmt",
"yuv420p",
"-f",
"h264",
"-",
]) else {
eprintln!("skipping: no ffmpeg with libx264");
return;
};
let path = std::env::temp_dir().join(format!("moq-vaapi-decode-{}.264", std::process::id()));
std::fs::write(&path, &stream).expect("write the elementary stream");
let reference = ffmpeg(&[
"-i",
path.to_str().expect("a UTF-8 temporary path"),
"-pix_fmt",
"nv12",
"-f",
"rawvideo",
"-",
]);
let _ = std::fs::remove_file(&path);
let reference = reference.expect("ffmpeg decodes what it just encoded");
let picture_len = (WIDTH * HEIGHT * 3 / 2) as usize;
assert_eq!(
reference.len(),
PICTURES * picture_len,
"ffmpeg decoded a different count"
);
let Ok(mut decoder) = Decoder::new(Config::new()) else {
eprintln!("skipping: no VA-API H.264 decoder");
return;
};
let units = split_access_units(&stream);
let mut decoded = Vec::new();
for (index, unit) in units.iter().enumerate() {
decoded.extend(decoder.decode(unit, index as u64).expect("decode an access unit"));
}
decoded.extend(decoder.flush().expect("flush the decoder"));
assert_eq!(decoded.len(), PICTURES, "the decoder lost or invented pictures");
let mut timestamps: Vec<u64> = decoded.iter().map(|picture| picture.timestamp).collect();
assert!(
timestamps.windows(2).any(|pair| pair[0] > pair[1]),
"the stream came out in decode order, so nothing here tests reordering"
);
timestamps.sort_unstable();
assert_eq!(
timestamps,
(0..PICTURES as u64).collect::<Vec<_>>(),
"a picture's timestamp did not survive reordering"
);
for (index, picture) in decoded.iter().enumerate() {
assert_eq!((picture.width, picture.height), (WIDTH, HEIGHT));
assert!(
picture.data == reference[index * picture_len..][..picture_len],
"picture {index} of {PICTURES} differs from the software decoder"
);
}
}
#[test]
fn a_new_sequence_decodes_at_its_own_size() {
const PER_SEQUENCE: u32 = 6;
let sizes = [(320u32, 240u32), (192u32, 144u32)];
let Ok(mut decoder) = Decoder::new(Config::new()) else {
eprintln!("skipping: no VA-API H.264 decoder");
return;
};
let mut decoded = Vec::new();
let mut timestamp = 0u64;
for &(width, height) in &sizes {
let Ok(mut encoder) = crate::encode::Encoder::new(encoder_config(width, height)) else {
eprintln!("skipping: no VA-API H.264 encoder");
return;
};
for step in 0..PER_SEQUENCE {
let unit = encoder
.encode_nv12(&nv12_frame(width, height, step), step == 0)
.expect("encode a picture");
decoded.extend(decoder.decode(&unit, timestamp).expect("decode a picture"));
timestamp += 1;
}
}
decoded.extend(decoder.flush().expect("flush the decoder"));
let expected: Vec<(u32, u32)> = sizes
.iter()
.flat_map(|&size| std::iter::repeat_n(size, PER_SEQUENCE as usize))
.collect();
let actual: Vec<(u32, u32)> = decoded.iter().map(|picture| (picture.width, picture.height)).collect();
assert_eq!(actual, expected, "a picture came out at the wrong size or out of order");
let timestamps: Vec<u64> = decoded.iter().map(|picture| picture.timestamp).collect();
assert_eq!(timestamps, (0..timestamp).collect::<Vec<_>>());
for picture in &decoded {
let expected_len = (picture.width * picture.height * 3 / 2) as usize;
assert_eq!(
picture.data.len(),
expected_len,
"a picture was read back at the wrong size"
);
}
}
fn ffmpeg(args: &[&str]) -> Option<Vec<u8>> {
let output = Command::new("ffmpeg").arg("-v").arg("error").args(args).output().ok()?;
if !output.status.success() {
eprintln!("ffmpeg {args:?} failed: {}", String::from_utf8_lossy(&output.stderr));
return None;
}
Some(output.stdout)
}
fn split_access_units(stream: &[u8]) -> Vec<Vec<u8>> {
let starts: Vec<usize> = (0..stream.len().saturating_sub(3))
.filter(|&i| stream[i..i + 3] == [0, 0, 1])
.collect();
let mut units = Vec::new();
let mut current: Vec<u8> = Vec::new();
let mut current_has_picture = false;
for (index, &start) in starts.iter().enumerate() {
let end = starts.get(index + 1).copied().unwrap_or(stream.len());
let nal_type = stream[start + 3] & 0x1f;
let is_slice = (1..=5).contains(&nal_type);
if current_has_picture && (is_slice || matches!(nal_type, 6..=9)) {
units.push(std::mem::take(&mut current));
current_has_picture = false;
}
current.extend_from_slice(&stream[start..end]);
current_has_picture |= is_slice;
}
if !current.is_empty() {
units.push(current);
}
units
}
fn inode(fd: &OwnedFd) -> u64 {
use std::os::unix::fs::MetadataExt as _;
let file = std::fs::File::from(fd.try_clone().expect("duplicate the descriptor"));
file.metadata().expect("stat the descriptor").ino()
}
}