use crate::annexb::iter_annexb_nals;
use crate::bitreader::BitReader;
use crate::ir::{DisplayChangeSet, DisplayGeometry, DisplayInfo};
use crate::nal::NalCodec;
use crate::nalu_types::HEVC_SPS_NUT;
use crate::sps::decode_avc_slice_sps;
use alloc::collections::BTreeMap;
use alloc::sync::Arc;
use alloc::vec::Vec;
const AVC_SLICE: u8 = 1;
const AVC_IDR: u8 = 5;
const AVC_SPS: u8 = 7;
const AVC_PPS: u8 = 8;
const AVC_NAL_TYPE_MASK: u8 = 0x1F;
const AVC_PREFIX_NAL: u8 = 14;
const SLICE_PREFIX_BYTES: usize = 64;
const HEVC_PPS: u8 = 34;
const HEVC_LAST_VCL: u8 = 31;
const HEVC_FIRST_IRAP: u8 = 16;
const HEVC_BLA_N_LP: u8 = 18;
const HEVC_IDR_W_RADL: u8 = 19;
const HEVC_IDR_N_LP: u8 = 20;
const BASE_LAYER: u8 = 0;
const MAX_UNRESOLVED_PPS: usize = 512;
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub(crate) enum PpsCoverage {
#[default]
Known,
PpsUnseen,
SpsUnseen,
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub(crate) struct Activation {
pub(crate) sps_ids: Vec<u8>,
pub(crate) coverage: PpsCoverage,
pub(crate) random_access: bool,
}
#[non_exhaustive]
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ActivationOutcome {
Unknown,
Unchanged,
Changed {
display: Arc<DisplayInfo>,
initial: bool,
random_access: bool,
changed: DisplayChangeSet,
},
Unseen,
}
fn avc_sps_id(nal: &[u8]) -> Option<u8> {
decode_avc_slice_sps(nal).ok().map(|s| s.sps_id)
}
fn avc_pps_ids(nal: &[u8]) -> Option<(u8, u8)> {
crate::slice_header::decode_avc_slice_pps(nal)
.ok()
.map(|p| (p.pps_id, p.sps_id))
}
fn hevc_sps_id(nal: &[u8]) -> Option<u8> {
let mut r = BitReader::with_unescape(nal.get(2..)?, "H.265 SPS").ok()?;
r.read_bits(4, "sps_video_parameter_set_id").ok()?;
r.read_bits(3, "sps_max_sub_layers_minus1").ok()?;
r.read_flag("sps_temporal_id_nesting_flag").ok()?;
u8::try_from(r.read_ue("sps_seq_parameter_set_id").ok()?).ok()
}
fn hevc_pps_ids(nal: &[u8]) -> Option<(u8, u8)> {
let mut r = BitReader::with_unescape(nal.get(2..)?, "H.265 PPS").ok()?;
let pps_id = u8::try_from(r.read_ue("pps_pic_parameter_set_id").ok()?).ok()?;
let sps_id = u8::try_from(r.read_ue("pps_seq_parameter_set_id").ok()?).ok()?;
Some((pps_id, sps_id))
}
fn hevc_slice_pps_id(nal: &[u8]) -> Option<u8> {
let t = crate::nal::nal_unit_type(NalCodec::Hevc, nal)?;
if t > HEVC_LAST_VCL || (nal[0] & 0x01) != BASE_LAYER {
return None;
}
let end = nal.len().min(2 + SLICE_PREFIX_BYTES);
let mut r = BitReader::with_unescape(nal.get(2..end)?, "H.265 slice header").ok()?;
if !r.read_flag("first_slice_segment_in_pic_flag").ok()? {
return None;
}
if (HEVC_FIRST_IRAP..=HEVC_LAST_VCL).contains(&t) {
r.read_flag("no_output_of_prior_pics_flag").ok()?;
}
u8::try_from(r.read_ue("slice_pic_parameter_set_id").ok()?).ok()
}
fn avc_slice_pps_id(nal: &[u8]) -> Option<u8> {
let end = nal.len().min(1 + SLICE_PREFIX_BYTES);
let mut r = BitReader::with_unescape(nal.get(1..end)?, "H.264 slice header").ok()?;
r.read_ue("first_mb_in_slice").ok()?;
r.read_ue("slice_type").ok()?;
u8::try_from(r.read_ue("pic_parameter_set_id").ok()?).ok()
}
fn avc_geometry(nal: &[u8]) -> Option<DisplayGeometry> {
DisplayGeometry::from_avc_sps(nal).ok()
}
fn hevc_geometry(nal: &[u8]) -> Option<DisplayGeometry> {
DisplayGeometry::from_hevc_sps(nal).ok()
}
#[non_exhaustive]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TrackCodec {
Avc,
Hevc,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
struct PpsBinding {
sps_id: u8,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum UnseenPart {
Pps,
Sps,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
struct Unresolved {
pps_id: u8,
part: UnseenPart,
}
#[derive(Debug, Clone)]
pub struct SpsTracker {
codec: TrackCodec,
sps: BTreeMap<u8, DisplayGeometry>,
pps: BTreeMap<u8, PpsBinding>,
state: Option<Arc<DisplayInfo>>,
unresolved: Vec<Unresolved>,
}
impl SpsTracker {
pub fn new(codec: TrackCodec) -> Self {
Self {
codec,
sps: BTreeMap::new(),
pps: BTreeMap::new(),
state: None,
unresolved: Vec::new(),
}
}
pub fn state(&self) -> Option<&Arc<DisplayInfo>> {
self.state.as_ref()
}
pub fn seed_state(&mut self, geometry: Option<DisplayGeometry>) {
if self.state.is_none() {
self.state = DisplayInfo::from_geometry(geometry).map(Arc::new);
}
}
pub fn observe(&mut self, au: &[u8]) -> ActivationOutcome {
let mut activation = Activation::default();
let mut slice_pps: Vec<Option<u8>> = Vec::new();
for nal in iter_annexb_nals(au) {
match self.codec {
TrackCodec::Avc => self.observe_avc_nal(nal, &mut slice_pps, &mut activation),
TrackCodec::Hevc => self.observe_hevc_nal(nal, &mut slice_pps, &mut activation),
}
}
if slice_pps.is_empty() {
return ActivationOutcome::Unknown;
}
self.resolve_slices(&slice_pps, &mut activation);
self.finish(activation)
}
fn observe_avc_nal(
&mut self,
nal: &[u8],
slice_pps: &mut Vec<Option<u8>>,
activation: &mut Activation,
) {
let Some(&header) = nal.first() else {
return;
};
match header & AVC_NAL_TYPE_MASK {
AVC_SPS => {
if let (Some(id), Some(g)) = (avc_sps_id(nal), avc_geometry(nal)) {
self.sps.insert(id, g);
}
}
AVC_PPS => {
if let Some((pps_id, sps_id)) = avc_pps_ids(nal) {
self.pps.insert(pps_id, PpsBinding { sps_id });
}
}
AVC_SLICE | AVC_IDR => {
activation.random_access |= header & AVC_NAL_TYPE_MASK == AVC_IDR;
slice_pps.push(avc_slice_pps_id(nal));
}
AVC_PREFIX_NAL => {
if let Some(inner) = nal.get(1..) {
self.observe_avc_nal(inner, slice_pps, activation);
}
}
_ => {}
}
}
fn observe_hevc_nal(
&mut self,
nal: &[u8],
slice_pps: &mut Vec<Option<u8>>,
activation: &mut Activation,
) {
let Some(t) = crate::nal::nal_unit_type(NalCodec::Hevc, nal) else {
return;
};
if (nal[0] & 0x01) != BASE_LAYER {
return;
}
match t {
HEVC_SPS_NUT => {
if let (Some(id), Some(g)) = (hevc_sps_id(nal), hevc_geometry(nal)) {
self.sps.insert(id, g);
}
}
HEVC_PPS => {
if let Some((pps_id, sps_id)) = hevc_pps_ids(nal) {
self.pps.insert(pps_id, PpsBinding { sps_id });
}
}
0..=HEVC_LAST_VCL => {
activation.random_access |=
matches!(t, HEVC_IDR_W_RADL | HEVC_IDR_N_LP | HEVC_BLA_N_LP);
slice_pps.push(hevc_slice_pps_id(nal));
}
_ => {}
}
}
fn resolve_slices(&mut self, slice_pps: &[Option<u8>], activation: &mut Activation) {
for pps_id in slice_pps.iter().flatten() {
match self.lookup(*pps_id) {
Ok(sps_id) => {
if !activation.sps_ids.contains(&sps_id) {
activation.sps_ids.push(sps_id);
}
}
Err(part) => {
activation.coverage = match (activation.coverage, part) {
(PpsCoverage::PpsUnseen, _) => PpsCoverage::PpsUnseen,
(_, UnseenPart::Sps) => PpsCoverage::SpsUnseen,
(_, UnseenPart::Pps) => PpsCoverage::PpsUnseen,
};
self.remember(*pps_id, part);
}
}
}
}
fn lookup(&self, pps_id: u8) -> Result<u8, UnseenPart> {
let pps = self.pps.get(&pps_id).ok_or(UnseenPart::Pps)?;
if !self.sps.contains_key(&pps.sps_id) {
return Err(UnseenPart::Sps);
}
Ok(pps.sps_id)
}
fn remember(&mut self, pps_id: u8, part: UnseenPart) {
let entry = Unresolved { pps_id, part };
if self.unresolved.contains(&entry) {
return;
}
if self.unresolved.len() >= MAX_UNRESOLVED_PPS {
self.unresolved.remove(0);
}
self.unresolved.push(entry);
}
fn forget_resolved(&mut self) {
let unresolved = core::mem::take(&mut self.unresolved);
self.unresolved = unresolved
.into_iter()
.filter(|u| self.lookup(u.pps_id).is_err())
.collect();
}
fn finish(&mut self, activation: Activation) -> ActivationOutcome {
if activation.sps_ids.is_empty() {
return if activation.coverage == PpsCoverage::Known {
ActivationOutcome::Unknown
} else {
ActivationOutcome::Unseen
};
}
self.forget_resolved();
let geometry = self.activated_geometry(&activation.sps_ids);
let display = Arc::new(DisplayInfo::new(geometry));
let unchanged = self
.state
.as_ref()
.is_some_and(|state| state.as_ref() == display.as_ref());
if unchanged {
return ActivationOutcome::Unchanged;
}
let initial = self.state.is_none();
let previous = self.state.replace(display.clone());
let changed = if initial {
DisplayChangeSet::initial()
} else {
DisplayInfo::changes_from(&display, previous.as_deref())
};
ActivationOutcome::Changed {
display,
initial,
random_access: activation.random_access,
changed,
}
}
fn activated_geometry(&self, ids: &[u8]) -> Option<DisplayGeometry> {
ids.first().and_then(|id| self.sps.get(id)).cloned()
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::ir::AspectRatioIdc;
use alloc::vec;
#[derive(Default)]
struct W {
out: Vec<u8>,
cur: u8,
n: u8,
}
impl W {
fn u(&mut self, bits: u32, v: u32) {
for i in (0..bits).rev() {
self.cur = (self.cur << 1) | (((v >> i) & 1) as u8);
self.n += 1;
if self.n == 8 {
self.out.push(self.cur);
self.cur = 0;
self.n = 0;
}
}
}
fn ue(&mut self, v: u32) {
let x = v + 1;
let len = 32 - x.leading_zeros();
self.u(len - 1, 0);
self.u(len, x);
}
fn se(&mut self, v: i32) {
self.ue(if v > 0 {
(2 * v - 1) as u32
} else {
(-2 * v) as u32
});
}
fn flag(&mut self, b: bool) {
self.u(1, u32::from(b));
}
fn rbsp_trailing_bits(mut self) -> Vec<u8> {
self.u(1, 1);
while self.n != 0 {
self.u(1, 0);
}
self.out
}
}
fn sps_nal(sps_id: u32, mb: u32, idc: u32, extended: Option<(u32, u32)>, level: u8) -> Vec<u8> {
let mut w = W::default();
w.ue(sps_id);
w.ue(0); w.ue(2); w.ue(1); w.flag(false); w.ue(mb - 1); w.ue(mb - 1); w.flag(true); w.flag(true); w.flag(false); w.flag(true); w.flag(true); w.u(8, idc);
if let Some((sw, sh)) = extended {
w.u(16, sw);
w.u(16, sh);
}
w.flag(false); w.flag(true); w.u(3, 5); w.flag(false); w.flag(false); let mut v = vec![0x67u8, 66, 0x00, level];
v.extend(w.rbsp_trailing_bits());
v
}
fn pps_nal(pps_id: u32, sps_id: u32) -> Vec<u8> {
let mut w = W::default();
w.ue(pps_id);
w.ue(sps_id);
w.flag(false); w.flag(false); w.ue(0); w.ue(0); w.ue(0); w.flag(false); w.u(2, 0); w.se(0); w.se(0); w.se(0); w.flag(false); w.flag(false); w.flag(false); let mut v = vec![0x68u8];
v.extend(w.rbsp_trailing_bits());
v
}
fn idr_slice_nal(pps_id: u32) -> Vec<u8> {
let mut w = W::default();
w.ue(0); w.ue(2); w.ue(pps_id);
w.ue(0); w.ue(0); let mut v = vec![0x65u8];
v.extend(w.rbsp_trailing_bits());
v
}
fn p_slice_nal(pps_id: u32) -> Vec<u8> {
let mut w = W::default();
w.ue(0); w.ue(0); w.ue(pps_id);
w.ue(1); let mut v = vec![0x41u8]; v.extend(w.rbsp_trailing_bits());
v
}
fn annexb(nals: &[Vec<u8>]) -> Vec<u8> {
let mut out = Vec::new();
for n in nals {
out.extend_from_slice(&[0, 0, 0, 1]);
out.extend_from_slice(n);
}
out
}
fn change(outcome: &ActivationOutcome) -> &DisplayChangeSet {
match outcome {
ActivationOutcome::Changed { changed, .. } => changed,
other => panic!("expected a change, got {other:?}"),
}
}
fn sar(outcome: &ActivationOutcome) -> Option<AspectRatioIdc> {
match outcome {
ActivationOutcome::Changed { display, .. } => display.geometry.as_ref()?.aspect_ratio,
_ => None,
}
}
fn tracker() -> SpsTracker {
SpsTracker::new(TrackCodec::Avc)
}
fn base_au() -> Vec<u8> {
annexb(&[sps_nal(0, 8, 4, None, 30), pps_nal(0, 0), idr_slice_nal(0)])
}
#[test]
fn first_activation_is_the_initial_change() {
let mut t = tracker();
let outcome = t.observe(&base_au());
let ActivationOutcome::Changed {
initial, display, ..
} = &outcome
else {
panic!("expected the initial change, got {outcome:?}");
};
assert!(*initial);
assert_eq!(
display.sample_aspect_ratio().map(|s| s.as_pair()),
Some((16, 11))
);
assert_eq!(*change(&outcome), DisplayChangeSet::initial());
assert_eq!(sar(&outcome), Some(AspectRatioIdc::Sar16x11));
}
#[test]
fn an_identical_repeat_changes_nothing() {
let mut t = tracker();
t.observe(&base_au());
assert!(matches!(
t.observe(&base_au()),
ActivationOutcome::Unchanged
));
assert!(matches!(
t.observe(&annexb(&[p_slice_nal(0)])),
ActivationOutcome::Unchanged
));
}
#[test]
fn a_geometry_neutral_sps_change_is_not_a_change() {
let mut t = tracker();
t.observe(&base_au());
let b = annexb(&[sps_nal(0, 8, 4, None, 40), pps_nal(0, 0), idr_slice_nal(0)]);
assert!(matches!(t.observe(&b), ActivationOutcome::Unchanged));
}
#[test]
fn a_sar_only_change_with_the_same_dar_sets_aspect_ratio() {
let mut t = tracker();
t.observe(&base_au());
let b = annexb(&[
sps_nal(0, 8, 255, Some((64, 45)), 30),
pps_nal(0, 0),
idr_slice_nal(0),
]);
let outcome = t.observe(&b);
let c = change(&outcome);
assert!(c.aspect_ratio && !c.initial, "{c:?}");
assert!(!c.output_rect && !c.video_signal && !c.overscan, "{c:?}");
assert_eq!(
sar(&outcome),
Some(AspectRatioIdc::ExtendedSar {
sar_width: 64,
sar_height: 45
})
);
}
#[test]
fn a_new_id_sent_ahead_of_use_fires_only_at_the_referencing_au() {
let mut t = tracker();
t.observe(&base_au());
assert!(matches!(
t.observe(&annexb(&[sps_nal(1, 16, 1, None, 30)])),
ActivationOutcome::Unknown
));
assert!(matches!(
t.observe(&annexb(&[p_slice_nal(0)])),
ActivationOutcome::Unchanged
));
let outcome = t.observe(&annexb(&[pps_nal(1, 1), idr_slice_nal(1)]));
let ActivationOutcome::Changed {
initial,
random_access,
display,
changed,
} = &outcome
else {
panic!("expected a change, got {outcome:?}");
};
assert!(!*initial && *random_access);
assert!(changed.aspect_ratio && changed.output_rect, "{changed:?}");
assert_eq!(display.geometry.as_ref().map(|g| g.coded_width), Some(256));
}
#[test]
fn a_mid_cvs_change_is_not_random_access() {
let mut t = tracker();
t.observe(&base_au());
let outcome = t.observe(&annexb(&[sps_nal(0, 8, 1, None, 30), p_slice_nal(0)]));
assert!(matches!(
outcome,
ActivationOutcome::Changed {
random_access: false,
..
}
));
}
#[test]
fn an_unseen_pps_holds_the_state_and_binds_on_a_later_au() {
let mut t = tracker();
t.observe(&base_au());
assert!(matches!(
t.observe(&annexb(&[idr_slice_nal(1)])),
ActivationOutcome::Unseen
));
assert_eq!(
t.state()
.and_then(|s| s.sample_aspect_ratio())
.map(|s| s.as_pair()),
Some((16, 11))
);
let outcome = t.observe(&annexb(&[
sps_nal(1, 8, 1, None, 30),
pps_nal(1, 1),
idr_slice_nal(1),
]));
assert!(
matches!(outcome, ActivationOutcome::Changed { .. }),
"{outcome:?}"
);
assert_eq!(sar(&outcome), Some(AspectRatioIdc::Square));
assert!(matches!(
t.observe(&annexb(&[p_slice_nal(1)])),
ActivationOutcome::Unchanged
));
}
#[test]
fn a_pps_naming_an_unseen_sps_is_held_too() {
let mut t = tracker();
t.observe(&base_au());
assert!(matches!(
t.observe(&annexb(&[pps_nal(2, 3), idr_slice_nal(2)])),
ActivationOutcome::Unseen
));
let outcome = t.observe(&annexb(&[sps_nal(3, 8, 1, None, 30), idr_slice_nal(2)]));
assert!(
matches!(outcome, ActivationOutcome::Changed { .. }),
"{outcome:?}"
);
}
#[test]
fn a_non_video_access_unit_resolves_nothing() {
let mut t = tracker();
assert!(matches!(
t.observe(&[0x00, 0x00, 0x01, 0xb3]),
ActivationOutcome::Unknown
));
}
}