diff --git a/crates/h264/src/decoder.rs b/crates/h264/src/decoder.rs index 19726df2e..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(); } @@ -457,7 +463,12 @@ impl Decoder { // 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())); - let queued = held(&mut self.out_queue.iter().filter(|o| !self.dpb.entries.iter().any(|e| Arc::ptr_eq(&e.frame, &o.frame))).map(|o| &o.frame)); + // 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, @@ -568,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.