diff --git a/crates/h264/src/cabac_mb.rs b/crates/h264/src/cabac_mb.rs index 78c257e9b..368750292 100644 --- a/crates/h264/src/cabac_mb.rs +++ b/crates/h264/src/cabac_mb.rs @@ -48,7 +48,7 @@ impl SliceDecoder<'_> { let intra_slice = self.sh.slice_type.is_intra(); loop { ensure!(addr < total, "macroblock address {addr} out of range"); - self.start_mb(addr); + self.start_mb(addr)?; let skip = !intra_slice && self.decode_skip_flag(&mut c); if skip { self.decode_skip()?; diff --git a/crates/h264/src/decoder.rs b/crates/h264/src/decoder.rs index a3244d473..455701f92 100644 --- a/crates/h264/src/decoder.rs +++ b/crates/h264/src/decoder.rs @@ -15,7 +15,7 @@ use crate::slicedec::{PicState, SliceDecoder}; use crate::transform::LevelScale; use crate::{ColorInfo, Picture}; use deckcraft_bitstream::{BitReader, annexb_nals, length_prefixed_nals, unescape_rbsp}; -use std::collections::VecDeque; +use std::collections::{HashMap, VecDeque}; use std::sync::{Arc, Mutex}; /// One slice of a picture job. @@ -90,6 +90,10 @@ pub struct Decoder { ls_cache: Vec<(Arc, Arc)>, /// Pictures leaving the DPB, converted once their decoding job has finished. out_queue: VecDeque, + /// Queued entries per frame id and their summed bytes, kept in step with `out_queue` so the + /// budget check does not rescan it. + queued_ids: HashMap, + queued_bytes: u64, shared: Arc, #[cfg(feature = "threads")] pool: Option, @@ -175,6 +179,8 @@ impl Decoder { active_sps: None, ls_cache: Vec::new(), out_queue: VecDeque::new(), + queued_ids: HashMap::new(), + queued_bytes: 0, shared: Arc::new(Shared::default()), #[cfg(feature = "threads")] pool, @@ -283,7 +289,7 @@ impl Decoder { if !o.frame.is_published(o.frame.mb_h() - 1) && self.out_queue.len() <= max_queue { break; } - let Some(o) = self.out_queue.pop_front() else { break }; + let Some(o) = self.pop_out() else { break }; pics.push(make_picture(&o)); } Ok(pics) @@ -295,7 +301,7 @@ impl Decoder { let mut outs = Vec::new(); self.dpb.flush(&mut outs); self.emit(outs); - let pics = self.out_queue.drain(..).map(|o| make_picture(&o)).collect(); + let pics = std::iter::from_fn(|| self.pop_out()).map(|o| make_picture(&o)).collect(); while let Some(f) = self.in_flight.pop_front() { f.wait_complete(); } @@ -451,13 +457,24 @@ impl Decoder { chroma_planes: if chroma == 0 { 0 } else { 2 }, bytes_per_sample: 1, }; - let pictures = (sps.max_dpb_frames() + 1 + self.max_in_flight) as u64; - ensure!( - shape.fits(pictures, self.memory_budget), - "{}x{} pictures with {pictures} in the decoder exceed the memory budget", - sps.width(), - sps.height() - ); + // Pictures already held count at their own size, not the incoming SPS's: after a + // resolution change, larger pictures of the old size can still be in the DPB or waiting + // for output (one decode call may carry many access units). A buffer shared with a DPB + // entry is counted once. + let pic = shape.bytes().unwrap_or(u64::MAX); + let held = |frames: &mut dyn Iterator>| frames.fold(0u64, |n, f| n.saturating_add(f.bytes())); + // Queued bytes minus the queue entries that share a DPB frame (each distinct DPB frame once). + let mut ids: Vec<(u32, u64)> = self.dpb.entries.iter().map(|e| (e.frame.id, e.frame.bytes())).collect(); + ids.sort_unstable(); + ids.dedup(); + let shared = ids.iter().fold(0u64, |n, (id, b)| n.saturating_add(self.queued_ids.get(id).copied().unwrap_or(0).saturating_mul(*b))); + let queued = self.queued_bytes.saturating_sub(shared); + let in_dpb = held(&mut self.dpb.entries.iter().map(|e| &e.frame)); + let dpb = sps.max_dpb_frames().max(self.dpb.entries.len()) as u64; + // The DPB (the larger of its new slots and what it holds now), the picture being decoded, + // and what is in flight or queued. + let bytes = pic.saturating_mul(dpb).max(in_dpb).saturating_add(pic).saturating_add(pic.saturating_mul(self.max_in_flight as u64).max(queued)); + ensure!(bytes <= self.memory_budget, "{}x{} pictures with {bytes} bytes in the decoder exceed the memory budget", sps.width(), sps.height()); Ok(()) } @@ -562,8 +579,24 @@ impl Decoder { } fn emit(&mut self, outs: Vec) { + for o in &outs { + *self.queued_ids.entry(o.frame.id).or_insert(0) += 1; + self.queued_bytes = self.queued_bytes.saturating_add(o.frame.bytes()); + } self.out_queue.extend(outs); } + + fn pop_out(&mut self) -> Option { + let o = self.out_queue.pop_front()?; + if let Some(n) = self.queued_ids.get_mut(&o.frame.id) { + *n -= 1; + if *n == 0 { + self.queued_ids.remove(&o.frame.id); + } + } + self.queued_bytes = self.queued_bytes.saturating_sub(o.frame.bytes()); + Some(o) + } } /// Decode all slices of one picture and publish its rows. diff --git a/crates/h264/src/dpb.rs b/crates/h264/src/dpb.rs index 46e8f060a..940b1a4c5 100644 --- a/crates/h264/src/dpb.rs +++ b/crates/h264/src/dpb.rs @@ -139,8 +139,9 @@ impl Dpb { let mut unused = (prev_ref_frame_num + 1) % max_frame_num; // Only the last `max_num_ref_frames` "non-existing" frames can survive the sliding window, and // inserting more than that evicts every earlier short-term reference. Skip straight to that tail - // (keeping the POC state walk for every missing frame_num) so a hostile gap of up to 65535 - // frames costs at most `max_num_ref_frames` frame allocations instead of one per missing frame. + // (advancing the POC state across the skipped run in constant time) so a hostile gap of up to + // 65535 frames costs at most `max_num_ref_frames` frame allocations and POC updates instead of + // one per missing frame. let gap = frame_num.wrapping_add(max_frame_num).wrapping_sub(unused) % max_frame_num; let keep = u32::try_from(max_num_ref_frames.max(1)).unwrap_or(u32::MAX); let skip = gap.saturating_sub(keep); @@ -149,10 +150,17 @@ impl Dpb { e.mark = RefMark::Unused; } self.prune(); - for _ in 0..skip { - poc_state(unused); - unused = (unused + 1) % max_frame_num; + // The POC state only advances FrameNumOffset where frame_num decreases, so walking the + // skipped run is the same as visiting its first value, the wrap (if the run crosses + // it) and its last value: constant work however long the gap. + let last = (unused + (skip - 1) % max_frame_num) % max_frame_num; + poc_state(unused); + if last < unused { + poc_state(max_frame_num - 1); + poc_state(0); } + poc_state(last); + unused = (last + 1) % max_frame_num; } let mut guard = 0; while unused != frame_num && guard < max_frame_num { @@ -438,14 +446,15 @@ impl Dpb { } } } else { - while self.entries.len() >= self.capacity { + // Evict until there is room (frame_num gap concealment can leave the DPB over-full), so + // the store never holds more pictures than the memory budget was checked for. + while self.entries.len() >= self.capacity.max(1) { if !self.bump(out) { if let Some(pos) = self.entries.iter().position(|e| !e.needed_for_output) { self.entries.remove(pos); } else { self.entries.remove(0); } - break; } } } @@ -546,6 +555,37 @@ mod tests { (made, walk, left) } + /// (FrameNumOffset, prev frame_num) after `walk`, by the rule of `PocState::update_gap_frame`. + fn poc_after(walk: &[u32], prev: u32, max: u32) -> (u64, u32) { + walk.iter().fold((0, prev), |(off, prev), &f| (if prev > f { off + u64::from(max) } else { off }, f)) + } + + #[test] + fn gap_walk_matches_every_frame_num_for_any_skip() { + for (prev, cur) in [(1, 9), (10, 3), (14, 0), (15, 14), (0, 15), (7, 6)] { + for refs in 1..4 { + let mut dpb = Dpb::new(); + let (_, walk, _) = fill(&mut dpb, prev, cur, 16, refs); + let every: Vec = (1..16).map(|i| (prev + i) % 16).take_while(|&f| f != cur).collect(); + assert_eq!(poc_after(&walk, prev, 16), poc_after(&every, prev, 16), "{prev} -> {cur}, {refs} refs"); + assert_eq!(walk.last(), every.last(), "{prev} -> {cur}, {refs} refs"); + } + } + } + + #[test] + fn reference_store_evicts_down_to_capacity() { + let mut dpb = Dpb::new(); + dpb.capacity = 2; + // Gap concealment can leave more non-output references than the capacity. + for i in 0..5 { + short(&mut dpb, i, i); + } + let mut out = Vec::new(); + assert!(!dpb.bump_for_space(RefMark::Short, 0, &mut out)); + assert!(dpb.entries.len() < 2, "{} pictures left before storing into a 2-picture DPB", dpb.entries.len()); + } + #[test] fn huge_gap_makes_only_the_frames_that_stay_referenced() { let mut dpb = Dpb::new(); @@ -553,8 +593,9 @@ mod tests { short(&mut dpb, 2, 1); let (made, walk, left) = fill(&mut dpb, 1, 65535, 65536, 4); assert_eq!(made, 4); - // The POC state still walks every missing frame_num, in order. - assert_eq!(walk.len(), 65533); + // The POC state ends where walking every missing frame_num would leave it, in a few steps. + assert_eq!(poc_after(&walk, 1, 65536), poc_after(&(2..65535).collect::>(), 1, 65536)); + assert!(walk.len() <= 8, "{} POC updates", walk.len()); assert_eq!(walk.first(), Some(&2)); assert_eq!(walk.last(), Some(&65534)); assert_eq!(left, vec![(65531, true), (65532, true), (65533, true), (65534, true)]); @@ -566,7 +607,9 @@ mod tests { short(&mut dpb, 1, 65000); let (made, walk, left) = fill(&mut dpb, 65000, 3, 65536, 2); assert_eq!(made, 2); - assert_eq!(walk.len(), 538); + let every: Vec = (65001..65536).chain(0..3).collect(); + assert_eq!(poc_after(&walk, 65000, 65536), poc_after(&every, 65000, 65536)); + assert!(walk.len() <= 6, "{} POC updates", walk.len()); assert_eq!(walk.first(), Some(&65001)); assert_eq!(walk.last(), Some(&2)); assert_eq!(left, vec![(1, true), (2, true)]); diff --git a/crates/h264/src/picture.rs b/crates/h264/src/picture.rs index 39add6a08..83331cb95 100644 --- a/crates/h264/src/picture.rs +++ b/crates/h264/src/picture.rs @@ -160,6 +160,14 @@ impl Frame { self.rows.len() } + /// Sample bytes this picture holds (luma and both chroma planes), as the decoder's memory + /// budget counts a picture. + pub fn bytes(&self) -> u64 { + (self.width as u64) + .saturating_mul(self.height as u64) + .saturating_add((self.cwidth as u64).saturating_mul(self.cheight as u64).saturating_mul(2)) + } + /// Publish MB row `r` (ignored if already published). pub fn publish(&self, r: usize, row: FrameRow) { let _ = self.rows[r].set(row); diff --git a/crates/h264/src/slicedec.rs b/crates/h264/src/slicedec.rs index 058ccfe16..649739986 100644 --- a/crates/h264/src/slicedec.rs +++ b/crates/h264/src/slicedec.rs @@ -346,8 +346,11 @@ impl<'a> SliceDecoder<'a> { // Neighbour helpers // ------------------------------------------------------------------------------------------ - /// Set up the current macroblock address and neighbour availability. - pub fn start_mb(&mut self, addr: usize) { + /// Set up the current macroblock address and neighbour availability. A macroblock an earlier + /// slice of this picture already decoded is refused: overlapping slices would otherwise redo + /// the picture's reconstruction once per slice. + pub fn start_mb(&mut self, addr: usize) -> Result<()> { + ensure!(self.pic.mbs[addr].slice_num == u32::MAX, "macroblock {addr} overlaps an earlier slice"); self.mb_addr = addr; self.mb_x = addr % self.mb_w; self.mb_y = addr / self.mb_w; @@ -364,6 +367,7 @@ impl<'a> SliceDecoder<'a> { ]; self.mv_done = 0; self.pic.mbs[addr] = MbState { slice_num: sn, ..MbState::default() }; + Ok(()) } #[inline] @@ -1467,7 +1471,7 @@ impl<'a> SliceDecoder<'a> { let run = r.read_ue()? as usize; ensure!(addr + run <= total, "mb_skip_run out of range"); for _ in 0..run { - self.start_mb(addr); + self.start_mb(addr)?; self.decode_skip()?; self.pic.mb_done(addr); addr += 1; @@ -1477,7 +1481,7 @@ impl<'a> SliceDecoder<'a> { } ensure!(addr < total, "macroblock address {addr} out of range"); } - self.start_mb(addr); + self.start_mb(addr)?; self.decode_mb_cavlc(r)?; self.pic.mb_done(addr); addr += 1; diff --git a/crates/h264/src/synth_tests.rs b/crates/h264/src/synth_tests.rs index 76609dbde..199f5e69d 100644 --- a/crates/h264/src/synth_tests.rs +++ b/crates/h264/src/synth_tests.rs @@ -113,6 +113,7 @@ fn avcc_rejects_a_valid_picture_slice_without_starting_a_picture() { poc_lsb: Some(0), poc_delta: None, wt: None, + first_mb: 0, }, false, ); @@ -153,6 +154,8 @@ struct Pic { /// Explicit weights (luma, cb, cr) as (weight, offset) for refIdx 0 of list 0 and list 1; /// log2 denominators are 5 (luma) and 3 (chroma). wt: Option<[[(i32, i32); 3]; 2]>, + /// first_mb_in_slice (`mbs` then lists the slice's macroblocks from there). + first_mb: u32, } impl Pic { @@ -165,7 +168,7 @@ impl Pic { } fn slice_header(w: &mut BitWriter, p: &Pic, cabac: bool) { - w.write_ue(0); // first_mb + w.write_ue(p.first_mb); w.write_ue(if p.b { 6 } else if p.is_p() { @@ -660,3 +663,63 @@ fn frame_num_gap_is_filled_when_the_sps_allows_gaps() { assert_eq!(out.len(), 2); assert_eq!(dec.stats().frame_num_gaps, 1); } + +/// An IDR picture followed by `pictures` reference P pictures of P_Skip macroblocks. +fn skip_stream(pictures: u32) -> (Vec, Vec>) { + let mut seed = 11u32; + let a: Vec>> = (0..NMB).map(|_| Some(random_mb(&mut seed))).collect(); + let mut head = sps(66, 2, false, PocType::T2); + head.extend(pps(false, false)); + head.extend(slice_nal(&Pic { idr: true, frame_num: 0, ref_idc: 3, mbs: a, ..Default::default() }, false)); + let ps = (1..=pictures).map(|k| slice_nal(&Pic { frame_num: k % 16, ref_idc: 2, mbs: vec![None; NMB], ..Default::default() }, false)).collect(); + (head, ps) +} + +#[test] +fn overlapping_slices_are_not_reconstructed_twice() { + let (head, ps) = skip_stream(1); + let mut once = Decoder::with_threads(1); + let mut out = once.decode(&[head.clone(), ps[0].clone()].concat(), 0).unwrap(); + out.extend(once.flush()); + assert!(once.take_error().is_none()); + assert_eq!(out.len(), 2); + // A second slice of the same picture starting at macroblock 1 covers those macroblocks again. + let again = slice_nal(&Pic { frame_num: 1, ref_idc: 2, mbs: vec![None; NMB - 1], first_mb: 1, ..Default::default() }, false); + let mut twice = Decoder::with_threads(1); + let result = twice.decode(&[head, ps[0].clone(), again].concat(), 0); + twice.flush(); + let err = result.err().or_else(|| twice.take_error()).expect("overlapping slice refused"); + assert!(err.to_string().contains("overlaps"), "{err}"); +} + +#[test] +fn output_queue_counts_against_the_memory_budget() { + let (head, ps) = skip_stream(40); + let sps_nal = sps(66, 2, false, PocType::T2); + let parsed = crate::params::Sps::parse(&deckcraft_bitstream::unescape_rbsp(&sps_nal[5..])).unwrap(); + let one = (NMB * 256 + 2 * NMB * 64) as u64; + let budget = (parsed.max_dpb_frames() + 1 + 3) as u64 * one; + // One access unit per call: output drains as it goes, and the stream fits. + let mut dec = Decoder::with_threads(1); + dec.set_memory_budget(budget); + let mut out = dec.decode(&head, 0).unwrap(); + for (i, p) in ps.iter().enumerate() { + out.extend(dec.decode(p, i as i64 + 1).unwrap()); + } + out.extend(dec.flush()); + assert!(dec.take_error().is_none()); + assert_eq!(out.len(), 41); + // Every access unit in one call: the pictures held for output stop fitting the budget. + let mut dec = Decoder::with_threads(1); + dec.set_memory_budget(budget); + let err = dec.decode(&[head, ps.concat()].concat(), 0).expect_err("output queue is over budget"); + assert!(err.to_string().contains("memory budget"), "{err}"); +} + +/// A held picture is charged at its own size (luma and both 4:2:0 chroma planes), the same +/// count the budget uses for the incoming SPS's pictures. +#[test] +fn a_held_frame_is_charged_at_its_own_size() { + let f = crate::picture::Frame::new(0, 0, 2, 3); + assert_eq!(f.bytes(), (2 * 3 * 256 + 2 * 2 * 3 * 64) as u64); +}