1mod cabac;
26#[cfg(feature = "profile")]
28pub use cabac::bin_census;
29mod mb16;
30mod params;
31
32pub use params::{Pps, Sps};
33pub use mb16::{MvField, MV_DUMP};
34
35#[doc(hidden)]
38fn abl_deblock() -> bool {
41 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
42 *ON.get_or_init(|| std::env::var("RFF_ABL_DEBLOCK").map_or(false, |v| v != "0"))
43}
44
45pub mod cabac_test {
46 pub use crate::cabac::Cabac;
47 pub use crate::mb16::b_inter_shape;
48 pub use crate::mb16::parse_mb_type_b;
49}
50
51use mb16::{FrameDecoder, GridPool, WeightTable};
52use rusty_h264_common::bit_reader::OutOfData;
53use rusty_h264_common::nal::{emulation_unprevent, split_annex_b};
54use rusty_h264_common::{BitReader, NalUnitType, YuvFrame};
55
56#[derive(Debug, Clone, PartialEq, Eq)]
58pub enum DecodeError {
59 Truncated,
61 MissingParameterSet,
63 Unsupported(&'static str),
65}
66
67impl From<OutOfData> for DecodeError {
68 fn from(_: OutOfData) -> Self {
69 DecodeError::Truncated
70 }
71}
72
73impl core::fmt::Display for DecodeError {
74 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
75 match self {
76 DecodeError::Truncated => f.write_str("bitstream truncated"),
77 DecodeError::MissingParameterSet => f.write_str("slice before SPS/PPS"),
78 DecodeError::Unsupported(s) => write!(f, "unsupported coding tool: {s}"),
79 }
80 }
81}
82
83impl std::error::Error for DecodeError {}
84
85pub(crate) type Ref = std::sync::Arc<RefFrame>;
94
95#[derive(Debug, Clone, Default)]
96#[allow(dead_code)]
97pub(crate) struct RefFrame {
98 pub py: Vec<u8>,
104 pub pu: Vec<u8>,
105 pub pv: Vec<u8>,
106 pub cw: usize,
107 pub ch: usize,
108 pub frame_num: u32,
110 pub poc: i32,
112 pub mv: Vec<(i32, i32)>,
116 pub ref_idx: Vec<i32>,
117 pub mv1: Vec<(i32, i32)>,
124 pub ref_idx1: Vec<i32>,
125 pub ref_poc: Vec<i32>,
129 pub w4: usize,
130 pub long_term: bool,
134 pub long_term_idx: u32,
135}
136
137pub(crate) const LPAD: usize = 16;
141pub(crate) const CPAD: usize = 8;
142
143impl RefFrame {
144 #[inline]
145 pub fn lstride(&self) -> usize {
146 self.cw + 2 * LPAD
147 }
148 #[inline]
149 pub fn cstride(&self) -> usize {
150 self.cw / 2 + 2 * CPAD
151 }
152}
153
154#[derive(Clone, Copy)]
156enum Mmco {
157 Unref(u32),
159 UnrefLong(u32),
161 AssignLong(u32, u32),
163 MaxLong(u32),
165 Reset,
167 CurrentLong(u32),
169}
170
171struct PendingPic {
175 fd: mb16::FrameDecoder,
176 frame_num: u32,
177 poc: i32,
178 next_mb: usize,
179 total_mb: usize,
180 slice_count: u16,
181 deblock: bool,
182 filter_offset_a: i32,
183 filter_offset_b: i32,
184 crop_r: usize,
185 crop_b: usize,
186 max_refs: usize,
187 log2_max_frame_num: u32,
188 is_reference: bool,
191 idr_long_term: bool,
192 mmco_ops: Vec<Mmco>,
193}
194
195fn no_pool() -> bool {
197 use std::sync::atomic::{AtomicU8, Ordering};
198 static ON: AtomicU8 = AtomicU8::new(0);
199 match ON.load(Ordering::Relaxed) {
200 0 => {
201 let v = std::env::var_os("RS_H264_NO_POOL").is_some_and(|v| v == "1");
202 ON.store(if v { 1 } else { 2 }, Ordering::Relaxed);
203 v
204 }
205 n => n == 1,
206 }
207}
208
209#[derive(Default)]
212pub struct Decoder {
213 sps: std::collections::HashMap<u32, Sps>,
216 pps: std::collections::HashMap<u32, Pps>,
217 refs: Vec<Ref>,
219 cur: Option<PendingPic>,
221 poc: PocState,
225 last_poc: i32,
227 prev_ref_frame_num: u32,
230 grid_pool: GridPool,
233 plane_pool: Vec<Vec<u8>>,
236 retired: Vec<Ref>,
241}
242
243#[derive(Default)]
245struct PocState {
246 prev_msb: i32,
247 prev_lsb: i32,
248 prev_frame_num: u32,
249 prev_frame_num_offset: i64,
250}
251
252impl Decoder {
253 pub fn new() -> Self {
255 Self::default()
256 }
257
258 pub fn decode(&mut self, annex_b: &[u8]) -> Result<Option<YuvFrame>, DecodeError> {
261 let _g = rusty_h264_common::prof::scope(rusty_h264_common::prof::Stage::Total);
262 let mut frame = None;
263 for nal in split_annex_b(annex_b) {
264 if nal.is_empty() {
265 continue;
266 }
267 let nal_type = NalUnitType::from_id(nal[0]);
268 let rbsp = emulation_unprevent(&nal[1..]);
269 match nal_type {
270 NalUnitType::Sps => {
271 let s = Sps::parse(&rbsp)?;
272 self.sps.insert(s.seq_parameter_set_id, s);
273 }
274 NalUnitType::Pps => {
275 let p = Pps::parse(&rbsp)?;
276 self.pps.insert(p.pic_parameter_set_id, p);
277 }
278 NalUnitType::IdrSlice | NalUnitType::NonIdrSlice => {
279 let nal_ref_idc = (nal[0] >> 5) & 3;
280 let is_idr = nal_type == NalUnitType::IdrSlice;
281 if let Some(f) = self.decode_slice(&rbsp, is_idr, nal_ref_idc)? {
282 frame = Some(f);
283 }
284 }
285 _ => {} }
287 }
288 Ok(frame)
289 }
290
291 pub fn decode_stream(&mut self, annex_b: &[u8]) -> Result<Vec<YuvFrame>, DecodeError> {
299 let mut out = Vec::new();
300 let mut gop: Vec<(i32, YuvFrame)> = Vec::new();
301 for au in split_access_units(annex_b) {
302 if au_is_idr(au) {
303 flush_gop(&mut gop, &mut out); }
305 if let Some(frame) = self.decode(au)? {
306 gop.push((self.last_poc, frame));
307 }
308 }
309 flush_gop(&mut gop, &mut out);
310 Ok(out)
311 }
312
313 fn decode_slice(
314 &mut self,
315 rbsp: &[u8],
316 is_idr: bool,
317 nal_ref_idc: u8,
318 ) -> Result<Option<YuvFrame>, DecodeError> {
319 let mut r = BitReader::new(rbsp);
320 let first_mb_in_slice = r.read_ue()? as usize;
322 let slice_type = r.read_ue()?;
323 let is_p = matches!(slice_type, 0 | 5);
324 let is_b = matches!(slice_type, 1 | 6);
325 let is_i = matches!(slice_type, 2 | 7);
326 if !is_p && !is_b && !is_i {
327 return Err(DecodeError::Unsupported("SP/SI slices"));
328 }
329 let pic_parameter_set_id = r.read_ue()?;
331 let pps = self.pps.get(&pic_parameter_set_id).cloned().ok_or(DecodeError::MissingParameterSet)?;
332 let sps = self.sps.get(&pps.seq_parameter_set_id).cloned().ok_or(DecodeError::MissingParameterSet)?;
333 let sps = &sps;
334 let pps = &pps;
335 let cabac = pps.entropy_coding_mode_flag;
339 let frame_num = r.read_bits(sps.log2_max_frame_num)?;
340 if is_idr {
341 let _idr_pic_id = r.read_ue()?;
342 }
343 let mut poc_lsb = 0u32;
346 let mut delta_poc_bottom = 0i32;
347 if sps.pic_order_cnt_type == 0 {
348 poc_lsb = r.read_bits(sps.log2_max_pic_order_cnt_lsb)?;
349 if pps.bottom_field_pic_order_present {
350 delta_poc_bottom = r.read_se()?;
351 }
352 } else if sps.pic_order_cnt_type == 1 && !sps.delta_pic_order_always_zero {
353 let _delta_pic_order_cnt_0 = r.read_se()?;
354 if pps.bottom_field_pic_order_present {
355 let _delta_pic_order_cnt_1 = r.read_se()?;
356 }
357 }
358 let pic_poc = if first_mb_in_slice == 0 {
361 self.compute_poc(sps, is_idr, nal_ref_idc, frame_num, poc_lsb, delta_poc_bottom)
362 } else {
363 self.cur.as_ref().map_or(0, |p| p.poc)
364 };
365 if pps.redundant_pic_cnt_present_flag {
370 let redundant_pic_cnt = r.read_ue()?;
371 if redundant_pic_cnt != 0 {
372 return Ok(None);
373 }
374 }
375 if std::env::var_os("RH264_DUMP_MB").is_some() {
376 eprintln!(
377 "SLICE fn={frame_num} poc={pic_poc} nal_ref_idc={nal_ref_idc} is_p={is_p} is_b={is_b} first_mb={first_mb_in_slice}"
378 );
379 }
380 let direct_spatial = if is_b { r.read_bit()? } else { true };
382 let mut num_ref_idx_l0 = pps.num_ref_idx_l0_default as usize;
383 let mut num_ref_idx_l1 = pps.num_ref_idx_l1_default as usize;
384 let mut reorder_l0: Vec<(u32, u32)> = Vec::new();
385 let mut reorder_l1: Vec<(u32, u32)> = Vec::new();
386 if is_p || is_b {
387 if r.read_bit()? {
389 num_ref_idx_l0 = (r.read_ue()? + 1) as usize;
390 if is_b {
391 num_ref_idx_l1 = (r.read_ue()? + 1) as usize;
392 }
393 }
394 if r.read_bit()? {
396 parse_ref_pic_list_modification(&mut r, &mut reorder_l0)?;
397 }
398 if is_b && r.read_bit()? {
399 parse_ref_pic_list_modification(&mut r, &mut reorder_l1)?;
401 }
402 }
403 let weights = if is_p && pps.weighted_pred {
408 Some(parse_pred_weight_table(&mut r, num_ref_idx_l0, 0, false)?)
409 } else if is_b && pps.weighted_bipred_idc == 1 {
410 return Err(DecodeError::Unsupported("explicit B weighted prediction"));
411 } else {
412 None
413 };
414 let mut idr_long_term = false;
418 let mut mmco_ops: Vec<Mmco> = Vec::new();
419 if nal_ref_idc == 0 {
420 } else if is_idr {
422 let _no_output_of_prior_pics = r.read_bit()?;
423 idr_long_term = r.read_bit()?; } else if r.read_bit()? {
425 loop {
427 let op = r.read_ue()?;
428 match op {
429 0 => break,
430 1 => mmco_ops.push(Mmco::Unref(r.read_ue()?)),
431 2 => mmco_ops.push(Mmco::UnrefLong(r.read_ue()?)),
432 3 => {
433 let diff = r.read_ue()?;
434 let idx = r.read_ue()?;
435 mmco_ops.push(Mmco::AssignLong(diff, idx));
436 }
437 4 => mmco_ops.push(Mmco::MaxLong(r.read_ue()?)),
438 5 => mmco_ops.push(Mmco::Reset),
439 6 => mmco_ops.push(Mmco::CurrentLong(r.read_ue()?)),
440 _ => return Err(DecodeError::Unsupported("invalid MMCO")),
441 }
442 if mmco_ops.len() > 128 {
443 return Err(DecodeError::Truncated);
444 }
445 }
446 }
447 let cabac_init_idc = if cabac && !is_i {
451 let v = r.read_ue()?;
452 if v > 2 {
453 return Err(DecodeError::Unsupported("invalid cabac_init_idc"));
454 }
455 v
456 } else {
457 0
458 };
459 let slice_qp_delta = r.read_se()?;
460 let mut deblock = true;
465 let (mut filter_offset_a, mut filter_offset_b) = (0i32, 0i32);
466 if pps.deblocking_filter_control_present_flag {
467 let disable_deblocking_filter_idc = r.read_ue()?;
468 deblock = disable_deblocking_filter_idc != 1;
471 if disable_deblocking_filter_idc != 1 {
472 filter_offset_a = r.read_se()? * 2;
474 filter_offset_b = r.read_se()? * 2;
475 }
476 }
477 let slice_qp = (pps.pic_init_qp + slice_qp_delta).clamp(0, 51) as u8;
478
479 if first_mb_in_slice == 0 && !is_idr && sps.gaps_in_frame_num_allowed {
482 self.insert_frame_num_gaps(
483 frame_num,
484 1u32 << sps.log2_max_frame_num,
485 sps.max_num_ref_frames.max(1) as usize,
486 sps.pic_width_in_mbs * 16,
487 sps.pic_height_in_mbs * 16,
488 );
489 }
490
491 let max_fn = 1u32 << sps.log2_max_frame_num;
494 let (ref_list0, ref_list1) = if is_b {
495 build_ref_list_b(
496 &self.refs, pic_poc, frame_num, max_fn,
497 num_ref_idx_l0, num_ref_idx_l1, &reorder_l0, &reorder_l1,
498 )?
499 } else if is_p {
500 (build_ref_list_p(&self.refs, frame_num, max_fn, num_ref_idx_l0, &reorder_l0)?, Vec::new())
501 } else {
502 (Vec::new(), Vec::new())
503 };
504 if first_mb_in_slice == 0 {
508 if is_idr {
509 self.refs.clear();
510 }
511 let _g_setup = rusty_h264_common::prof::scope(rusty_h264_common::prof::Stage::DecSetup);
522 let mut fd = FrameDecoder::with_pool(
523 sps.pic_width_in_mbs,
524 sps.pic_height_in_mbs,
525 slice_qp,
526 pps.chroma_qp_index_offset,
527 ref_list0,
528 num_ref_idx_l0,
529 pps.constrained_intra_pred_flag,
530 pps.transform_8x8_mode_flag,
531 sps.profile_idc != 66, if no_pool() { GridPool::default() } else { std::mem::take(&mut self.grid_pool) },
536 );
537 if is_b {
538 fd.set_b_context(
539 ref_list1,
540 num_ref_idx_l1,
541 direct_spatial,
542 pic_poc,
543 pps.weighted_bipred_idc,
544 sps.direct_8x8_inference,
545 );
546 }
547 if sps.has_scaling || pps.pic_scaling_matrix_present {
548 let (s4, s8) = resolve_scaling(sps, pps);
549 fd.set_scaling(s4, s8);
550 }
551 if let Some(w) = weights {
552 fd.set_weights(w);
553 }
554 if let Some(prev) = self.cur.take() {
562 if prev.next_mb < prev.total_mb {
563 return Err(DecodeError::Truncated);
564 }
565 }
566 self.cur = Some(PendingPic {
567 fd,
568 frame_num,
569 poc: pic_poc,
570 next_mb: 0,
571 total_mb: sps.pic_width_in_mbs * sps.pic_height_in_mbs,
572 slice_count: 0,
573 deblock,
574 filter_offset_a,
575 filter_offset_b,
576 crop_r: sps.frame_crop_right as usize,
577 crop_b: sps.frame_crop_bottom as usize,
578 max_refs: sps.max_num_ref_frames.max(1) as usize,
579 log2_max_frame_num: sps.log2_max_frame_num,
580 is_reference: nal_ref_idc != 0,
581 idr_long_term,
582 mmco_ops,
583 });
584 } else {
585 let Some(pic) = self.cur.as_mut() else {
587 return Err(DecodeError::Unsupported("slice continues a missing picture"));
588 };
589 pic.fd.begin_slice(slice_qp, ref_list0, num_ref_idx_l0);
590 if is_b {
591 pic.fd.set_b_context(
592 ref_list1,
593 num_ref_idx_l1,
594 direct_spatial,
595 pic.poc,
596 pps.weighted_bipred_idc,
597 sps.direct_8x8_inference,
598 );
599 }
600 if sps.has_scaling || pps.pic_scaling_matrix_present {
601 let (s4, s8) = resolve_scaling(sps, pps);
602 pic.fd.set_scaling(s4, s8);
603 }
604 if let Some(w) = weights {
605 pic.fd.set_weights(w);
606 }
607 pic.deblock = deblock;
609 pic.filter_offset_a = filter_offset_a;
610 pic.filter_offset_b = filter_offset_b;
611 pic.idr_long_term |= idr_long_term;
612 pic.mmco_ops.extend(mmco_ops);
613 }
614
615 let pic = self.cur.as_mut().expect("pending picture set above");
616 pic.fd.set_deblock_params(deblock && !abl_deblock(), filter_offset_a, filter_offset_b);
620 let first = first_mb_in_slice.min(pic.total_mb);
621 let next = if cabac {
622 r.align_to_byte().map_err(|_| DecodeError::Truncated)?;
624 let (data, start) = (r.data(), r.bit_pos() / 8);
625 pic.fd
626 .decode_slice_data_cabac(data, start, slice_qp, cabac_init_idc, is_i, is_p, first)
627 } else {
628 pic.fd.decode_slice_data(&mut r, is_p, first)
629 }
630 .map_err(|e| match e {
631 mb16::MbError::Truncated => DecodeError::Truncated,
632 mb16::MbError::Unsupported(s) => DecodeError::Unsupported(s),
633 })?;
634 pic.next_mb = next;
635 pic.slice_count += 1;
636 if std::env::var_os("RH264_DUMP_MB").is_some() {
637 eprintln!(
638 " slice decoded {}/{} MBs{}",
639 next,
640 pic.total_mb,
641 if next < pic.total_mb { " <-- INCOMPLETE" } else { "" }
642 );
643 }
644
645 if pic.next_mb < pic.total_mb {
646 return Ok(None); }
648
649 let pic = self.cur.take().expect("pending picture");
651 let PendingPic {
652 mut fd,
653 frame_num,
654 poc,
655 deblock,
656 filter_offset_a,
657 filter_offset_b,
658 crop_r,
659 crop_b,
660 max_refs,
661 log2_max_frame_num,
662 is_reference,
663 idr_long_term,
664 mmco_ops,
665 ..
666 } = pic;
667 self.last_poc = poc;
668 if deblock && !abl_deblock() {
676 fd.deblock(filter_offset_a, filter_offset_b);
677 }
678 let reference = if is_reference {
681 let _dg = rusty_h264_common::prof::scope(rusty_h264_common::prof::Stage::DpbClone);
682 Some(fd.as_reference_pooled(&mut self.plane_pool))
683 } else {
684 None
686 };
687 let _fg = rusty_h264_common::prof::scope(rusty_h264_common::prof::Stage::Finalize);
688 if let Some(mut reference) = reference {
689 reference.frame_num = frame_num;
690 reference.poc = poc;
691 if std::env::var_os("RH264_DUMP_MB").is_some() {
692 eprintln!("DPB-ADD fn={frame_num} poc={poc}");
693 }
694 if idr_long_term {
695 reference.long_term = true;
696 reference.long_term_idx = 0;
697 }
698 self.prev_ref_frame_num =
700 self.apply_ref_marking(reference, &mmco_ops, frame_num, log2_max_frame_num, max_refs);
701 }
702 let (frame, pool) = fd.into_frame_recycle(crop_r, crop_b);
703 self.grid_pool = pool;
704 self.reclaim_retired();
705 Ok(Some(frame))
706 }
707
708 fn reclaim_retired(&mut self) {
713 for arc in self.retired.drain(..) {
714 if let Ok(rf) = std::sync::Arc::try_unwrap(arc) {
715 self.plane_pool.push(rf.py);
716 self.plane_pool.push(rf.pu);
717 self.plane_pool.push(rf.pv);
718 }
719 }
720 self.plane_pool.truncate(18);
723 }
724
725 fn insert_frame_num_gaps(&mut self, frame_num: u32, max_fn: u32, max_refs: usize, w: usize, h: usize) {
731 if max_fn == 0 {
732 return;
733 }
734 let start = (self.prev_ref_frame_num + 1) % max_fn;
735 let gap = (frame_num + max_fn - start) % max_fn;
736 if gap == 0 {
737 return;
738 }
739 let cap = max_refs.max(1);
745 let n = (gap as usize).min(cap);
746 let (cw, ch) = (w, h);
747 let mut expected = (frame_num + max_fn - n as u32) % max_fn;
748 for _ in 0..n {
749 self.refs.insert(
750 0,
751 std::sync::Arc::new(RefFrame {
752 py: vec![128; (cw + 2 * LPAD) * (ch + 2 * LPAD)],
754 pu: vec![128; (cw / 2 + 2 * CPAD) * (ch / 2 + 2 * CPAD)],
755 pv: vec![128; (cw / 2 + 2 * CPAD) * (ch / 2 + 2 * CPAD)],
756 cw,
757 ch,
758 frame_num: expected,
759 poc: 0,
760 mv: Vec::new(),
761 ref_idx: Vec::new(),
762 mv1: Vec::new(),
763 ref_idx1: Vec::new(),
764 ref_poc: Vec::new(),
765 w4: 0,
766 long_term: false,
767 long_term_idx: 0,
768 }),
769 );
770 self.refs.truncate(cap);
771 expected = (expected + 1) % max_fn;
772 }
773 self.prev_ref_frame_num = (frame_num + max_fn - 1) % max_fn;
774 }
775
776 pub fn last_poc(&self) -> i32 {
780 self.last_poc
781 }
782
783 fn compute_poc(
787 &mut self,
788 sps: &Sps,
789 is_idr: bool,
790 nal_ref_idc: u8,
791 frame_num: u32,
792 poc_lsb: u32,
793 delta_bottom: i32,
794 ) -> i32 {
795 match sps.pic_order_cnt_type {
796 0 => {
797 let max_lsb = 1i32 << sps.log2_max_pic_order_cnt_lsb;
798 let (prev_msb, prev_lsb) =
799 if is_idr { (0, 0) } else { (self.poc.prev_msb, self.poc.prev_lsb) };
800 let lsb = poc_lsb as i32;
801 let msb = if lsb < prev_lsb && prev_lsb - lsb >= max_lsb / 2 {
802 prev_msb + max_lsb
803 } else if lsb > prev_lsb && lsb - prev_lsb > max_lsb / 2 {
804 prev_msb - max_lsb
805 } else {
806 prev_msb
807 };
808 let top = msb + lsb;
809 let poc = top.min(top + delta_bottom);
810 if nal_ref_idc != 0 {
811 self.poc.prev_msb = msb;
812 self.poc.prev_lsb = lsb;
813 }
814 poc
815 }
816 2 => {
817 let max_fn = 1i64 << sps.log2_max_frame_num;
818 let offset = if is_idr {
819 0
820 } else if self.poc.prev_frame_num > frame_num {
821 self.poc.prev_frame_num_offset + max_fn
822 } else {
823 self.poc.prev_frame_num_offset
824 };
825 let poc = if is_idr {
826 0
827 } else {
828 2 * (offset + frame_num as i64) - i64::from(nal_ref_idc == 0)
829 };
830 self.poc.prev_frame_num_offset = offset;
831 self.poc.prev_frame_num = frame_num;
832 poc as i32
833 }
834 _ => {
835 self.poc.prev_frame_num = frame_num;
836 frame_num as i32 * 2
837 }
838 }
839 }
840
841 fn apply_ref_marking(
852 &mut self,
853 mut reference: RefFrame,
854 ops: &[Mmco],
855 frame_num: u32,
856 log2_max_frame_num: u32,
857 max_refs: usize,
858 ) -> u32 {
859 let max = 1i64 << log2_max_frame_num;
860 let curr = frame_num as i64;
861 let pic_num = |rf: &RefFrame| -> i64 {
862 if (rf.frame_num as i64) > curr {
863 rf.frame_num as i64 - max
864 } else {
865 rf.frame_num as i64
866 }
867 };
868
869 if ops.is_empty() {
870 let out_fn = reference.frame_num;
874 self.refs.insert(0, std::sync::Arc::new(reference));
875 while self.refs.len() > max_refs {
876 match self.refs.iter().rposition(|r| !r.long_term) {
877 Some(pos) => {
878 let evicted = self.refs.remove(pos);
881 self.retired.push(evicted);
882 }
883 None => break,
884 }
885 }
886 return out_fn;
887 }
888
889 for &op in ops {
891 match op {
892 Mmco::Unref(diff) => {
893 let target = curr - (diff as i64 + 1);
894 self.refs.retain(|r| r.long_term || pic_num(r) != target);
895 }
896 Mmco::UnrefLong(ltpn) => {
897 self.refs.retain(|r| !(r.long_term && r.long_term_idx == ltpn));
898 }
899 Mmco::AssignLong(diff, idx) => {
900 let target = curr - (diff as i64 + 1);
901 self.refs.retain(|r| !(r.long_term && r.long_term_idx == idx));
902 for r in self.refs.iter_mut() {
903 if !r.long_term && pic_num(r) == target {
904 let r = std::sync::Arc::make_mut(r);
906 r.long_term = true;
907 r.long_term_idx = idx;
908 }
909 }
910 }
911 Mmco::MaxLong(max_plus1) => {
912 self.refs.retain(|r| !(r.long_term && r.long_term_idx + 1 > max_plus1));
913 }
914 Mmco::Reset => {
915 self.refs.clear();
916 reference.frame_num = 0;
917 }
918 Mmco::CurrentLong(idx) => {
919 self.refs.retain(|r| !(r.long_term && r.long_term_idx == idx));
920 reference.long_term = true;
921 reference.long_term_idx = idx;
922 }
923 }
924 }
925 let out_fn = reference.frame_num;
926 self.refs.insert(0, std::sync::Arc::new(reference));
927 let cap = max_refs.max(16);
929 if self.refs.len() > cap {
930 self.refs.truncate(cap);
931 }
932 out_fn
933 }
934}
935
936fn flush_gop(gop: &mut Vec<(i32, YuvFrame)>, out: &mut Vec<YuvFrame>) {
938 gop.sort_by_key(|(poc, _)| *poc);
939 out.extend(gop.drain(..).map(|(_, f)| f));
940}
941
942fn au_is_idr(au: &[u8]) -> bool {
944 split_annex_b(au)
945 .iter()
946 .any(|n| !n.is_empty() && NalUnitType::from_id(n[0]) == NalUnitType::IdrSlice)
947}
948
949pub fn split_access_units(stream: &[u8]) -> Vec<&[u8]> {
957 let mut codes: Vec<(usize, bool)> = Vec::new();
959 let mut i = 0;
960 while i + 3 <= stream.len() {
961 if stream[i] == 0 && stream[i + 1] == 0 && stream[i + 2] == 1 {
962 let nal_type = NalUnitType::from_id(stream.get(i + 3).copied().unwrap_or(0));
963 let is_vcl = matches!(nal_type, NalUnitType::IdrSlice | NalUnitType::NonIdrSlice);
964 let sc = if i > 0 && stream[i - 1] == 0 { i - 1 } else { i };
966 codes.push((sc, is_vcl));
967 i += 3;
968 } else {
969 i += 1;
970 }
971 }
972 if codes.is_empty() {
973 return vec![stream];
974 }
975 let mut aus = Vec::new();
976 let mut start = codes[0].0;
977 for k in 0..codes.len() {
978 if codes[k].1 {
979 let end = codes.get(k + 1).map_or(stream.len(), |c| c.0);
980 aus.push(&stream[start..end]);
981 start = end;
982 }
983 }
984 aus
985}
986
987fn parse_pred_weight_table(
990 r: &mut BitReader,
991 num_l0: usize,
992 num_l1: usize,
993 is_b: bool,
994) -> Result<WeightTable, DecodeError> {
995 let luma_log2_denom = r.read_ue()? as i32;
996 let chroma_log2_denom = r.read_ue()? as i32;
997 if !(0..=7).contains(&luma_log2_denom) || !(0..=7).contains(&chroma_log2_denom) {
1001 return Err(DecodeError::Unsupported("invalid weight denom"));
1002 }
1003 let mut wt = WeightTable {
1004 luma_log2_denom,
1005 chroma_log2_denom,
1006 ..Default::default()
1007 };
1008 let lists: &[(usize, usize)] = if is_b {
1009 &[(0, num_l0), (1, num_l1)]
1010 } else {
1011 &[(0, num_l0)]
1012 };
1013 for &(list, n) in lists {
1014 let mut luma = Vec::with_capacity(n);
1015 let mut chroma = Vec::with_capacity(n);
1016 for _ in 0..n {
1017 let (mut lw, mut lo) = (1 << luma_log2_denom, 0);
1018 if r.read_bit()? {
1019 lw = r.read_se()?;
1020 lo = r.read_se()?;
1021 }
1022 luma.push((lw, lo));
1023 let mut ch = [(1 << chroma_log2_denom, 0); 2];
1024 if r.read_bit()? {
1025 for slot in ch.iter_mut() {
1026 *slot = (r.read_se()?, r.read_se()?);
1027 }
1028 }
1029 chroma.push(ch);
1030 }
1031 wt.luma[list] = luma;
1032 wt.chroma[list] = chroma;
1033 }
1034 Ok(wt)
1035}
1036
1037fn resolve_scaling(sps: &Sps, pps: &Pps) -> ([[i32; 16]; 6], [[i32; 64]; 2]) {
1042 use crate::params::{
1043 DEFAULT_4X4_INTER, DEFAULT_4X4_INTRA, DEFAULT_8X8_INTER, DEFAULT_8X8_INTRA,
1044 };
1045 const ZZ4: [usize; 16] = [0, 1, 4, 8, 5, 2, 3, 6, 9, 12, 13, 10, 7, 11, 14, 15];
1046 const ZZ8: [usize; 64] = [
1048 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,
1049 20, 13, 6, 7, 14, 21, 28, 35, 42, 49, 56, 57, 50, 43, 36, 29, 22, 15, 23, 30, 37, 44, 51,
1050 58, 59, 52, 45, 38, 31, 39, 46, 53, 60, 61, 54, 47, 55, 62, 63,
1051 ];
1052 let mut z4 = [[16u8; 16]; 6];
1055 for i in 0..6 {
1056 z4[i] = if pps.pic_scaling_matrix_present {
1057 if pps.scaling_present_4x4[i] {
1058 pps.scaling_4x4[i]
1059 } else {
1060 match i {
1061 0 if sps.has_scaling => sps.scaling_4x4[0],
1062 0 => DEFAULT_4X4_INTRA,
1063 3 if sps.has_scaling => sps.scaling_4x4[3],
1064 3 => DEFAULT_4X4_INTER,
1065 _ => z4[i - 1],
1066 }
1067 }
1068 } else {
1069 sps.scaling_4x4[i]
1070 };
1071 }
1072 let mut z8 = [[16u8; 64]; 2];
1073 for (i, list) in z8.iter_mut().enumerate() {
1074 *list = if pps.pic_scaling_matrix_present {
1075 if pps.scaling_present_8x8[i] {
1076 pps.scaling_8x8[i]
1077 } else if sps.has_scaling {
1078 sps.scaling_8x8[i]
1079 } else if i == 0 {
1080 DEFAULT_8X8_INTRA
1081 } else {
1082 DEFAULT_8X8_INTER
1083 }
1084 } else {
1085 sps.scaling_8x8[i]
1086 };
1087 }
1088 let mut out4 = [[16i32; 16]; 6];
1089 for (li, list) in out4.iter_mut().enumerate() {
1090 for k in 0..16 {
1091 list[ZZ4[k]] = z4[li][k] as i32;
1092 }
1093 }
1094 let mut out8 = [[16i32; 64]; 2];
1095 for (li, list) in out8.iter_mut().enumerate() {
1096 for k in 0..64 {
1097 list[ZZ8[k]] = z8[li][k] as i32;
1098 }
1099 }
1100 (out4, out8)
1101}
1102
1103fn parse_ref_pic_list_modification(
1106 r: &mut BitReader,
1107 out: &mut Vec<(u32, u32)>,
1108) -> Result<(), DecodeError> {
1109 loop {
1110 let idc = r.read_ue()?;
1111 if idc == 3 {
1112 break;
1113 }
1114 if idc > 3 {
1115 return Err(DecodeError::Unsupported("invalid ref_pic_list_modification"));
1116 }
1117 let val = r.read_ue()?; out.push((idc, val));
1119 if out.len() > 64 {
1120 return Err(DecodeError::Truncated); }
1122 }
1123 Ok(())
1124}
1125
1126fn build_ref_list_p(
1130 dpb: &[Ref],
1131 curr_frame_num: u32,
1132 max_frame_num: u32,
1133 num_active: usize,
1134 mods: &[(u32, u32)],
1135) -> Result<Vec<Ref>, DecodeError> {
1136 let curr = curr_frame_num as i64;
1137 let max = max_frame_num as i64;
1138 let pic_num = |fnum: u32| -> i64 {
1139 let f = fnum as i64;
1140 if f > curr { f - max } else { f }
1141 };
1142 let mut init: Vec<Ref> = dpb.iter().filter(|r| !r.long_term).cloned().collect();
1143 init.sort_by_key(|rf| core::cmp::Reverse(pic_num(rf.frame_num)));
1144 let mut long: Vec<Ref> = dpb.iter().filter(|r| r.long_term).cloned().collect();
1145 long.sort_by_key(|rf| rf.long_term_idx);
1146 init.extend(long);
1147 apply_list_modification(init, curr_frame_num, max_frame_num, num_active, mods)
1148}
1149
1150#[allow(clippy::too_many_arguments)]
1155fn build_ref_list_b(
1156 dpb: &[Ref],
1157 curr_poc: i32,
1158 curr_frame_num: u32,
1159 max_frame_num: u32,
1160 num0: usize,
1161 num1: usize,
1162 mods0: &[(u32, u32)],
1163 mods1: &[(u32, u32)],
1164) -> Result<(Vec<Ref>, Vec<Ref>), DecodeError> {
1165 let mut less: Vec<Ref> =
1166 dpb.iter().filter(|r| !r.long_term && r.poc < curr_poc).cloned().collect();
1167 let mut greater: Vec<Ref> =
1168 dpb.iter().filter(|r| !r.long_term && r.poc > curr_poc).cloned().collect();
1169 let mut long: Vec<Ref> = dpb.iter().filter(|r| r.long_term).cloned().collect();
1170 less.sort_by_key(|r| core::cmp::Reverse(r.poc)); greater.sort_by_key(|r| r.poc); long.sort_by_key(|r| r.long_term_idx);
1173
1174 let mut init0 = less.clone();
1175 init0.extend(greater.clone());
1176 init0.extend(long.clone());
1177 let mut init1 = greater;
1178 init1.extend(less);
1179 init1.extend(long);
1180
1181 let eq_len = num0.min(num1).min(init0.len()).min(init1.len());
1184 if num1 > 1
1185 && init1.len() > 1
1186 && (0..eq_len).all(|i| same_picture(&init0[i], &init1[i]))
1187 && eq_len == num1.min(init1.len())
1188 && eq_len == num0.min(init0.len())
1189 {
1190 init1.swap(0, 1);
1191 }
1192
1193 let list0 = apply_list_modification(init0, curr_frame_num, max_frame_num, num0, mods0)?;
1194 let list1 = apply_list_modification(init1, curr_frame_num, max_frame_num, num1, mods1)?;
1195 Ok((list0, list1))
1196}
1197
1198fn same_picture(a: &RefFrame, b: &RefFrame) -> bool {
1200 a.long_term == b.long_term
1201 && if a.long_term { a.long_term_idx == b.long_term_idx } else { a.poc == b.poc }
1202}
1203
1204fn apply_list_modification(
1209 init: Vec<Ref>,
1210 curr_frame_num: u32,
1211 max_frame_num: u32,
1212 num_active: usize,
1213 mods: &[(u32, u32)],
1214) -> Result<Vec<Ref>, DecodeError> {
1215 if mods.is_empty() {
1216 let mut init = init;
1217 init.truncate(num_active.max(1));
1218 return Ok(init);
1219 }
1220 let curr = curr_frame_num as i64;
1221 let max = max_frame_num as i64;
1222 let mut list = init.clone();
1223 let mut pic_num_pred = curr;
1224 let mut refidx = 0usize;
1225 for &(idc, val) in mods {
1226 let matches: Box<dyn Fn(&RefFrame) -> bool> = if idc == 2 {
1227 Box::new(move |r: &RefFrame| r.long_term && r.long_term_idx == val)
1228 } else {
1229 let abs_diff = (val as i64) + 1;
1230 let no_wrap = if idc == 0 {
1231 let x = pic_num_pred - abs_diff;
1232 if x < 0 { x + max } else { x }
1233 } else {
1234 let x = pic_num_pred + abs_diff;
1235 if x >= max { x - max } else { x }
1236 };
1237 pic_num_pred = no_wrap;
1238 let target = if no_wrap > curr { no_wrap - max } else { no_wrap };
1239 Box::new(move |r: &RefFrame| {
1240 let pn = if r.frame_num as i64 > curr {
1241 r.frame_num as i64 - max
1242 } else {
1243 r.frame_num as i64
1244 };
1245 !r.long_term && pn == target
1246 })
1247 };
1248 let found = init.iter().find(|r| matches(r)).cloned();
1249 let Some(found) = found else {
1250 if std::env::var_os("RH264_DUMP_MB").is_some() {
1251 let cand: Vec<String> = init
1252 .iter()
1253 .map(|r| {
1254 let pn = if r.frame_num as i64 > curr {
1255 r.frame_num as i64 - max
1256 } else {
1257 r.frame_num as i64
1258 };
1259 format!("(fn={} poc={} lt={} picnum={})", r.frame_num, r.poc, r.long_term, pn)
1260 })
1261 .collect();
1262 eprintln!(
1263 "MODFAIL idc={idc} val={val} curr_frame_num={curr} max={max} init={}",
1264 cand.join(" ")
1265 );
1266 }
1267 return Err(DecodeError::Truncated); };
1269 if refidx > list.len() {
1270 break;
1271 }
1272 list.insert(refidx, found);
1273 if let Some(dup) = list.iter().enumerate().skip(refidx + 1).find(|(_, r)| matches(r)).map(|(i, _)| i) {
1274 list.remove(dup);
1275 }
1276 refidx += 1;
1277 if refidx >= num_active {
1278 break;
1279 }
1280 }
1281 list.truncate(num_active.max(1));
1282 Ok(list)
1283}
1284
1285#[cfg(test)]
1286mod tests {
1287 use super::*;
1288
1289 fn ref_at(poc: i32, fnum: u32) -> Ref {
1290 std::sync::Arc::new(RefFrame {
1291 py: vec![],
1292 pu: vec![],
1293 pv: vec![],
1294 cw: 0,
1295 ch: 0,
1296 frame_num: fnum,
1297 poc,
1298 mv: Vec::new(),
1299 ref_idx: Vec::new(),
1300 mv1: Vec::new(),
1301 ref_idx1: Vec::new(),
1302 ref_poc: Vec::new(),
1303 w4: 0,
1304 long_term: false,
1305 long_term_idx: 0,
1306 })
1307 }
1308
1309 #[test]
1310 fn b_ref_lists_ordered_by_poc() {
1311 let dpb = vec![ref_at(8, 4), ref_at(6, 3), ref_at(2, 1), ref_at(0, 0)];
1313 let (l0, l1) = build_ref_list_b(&dpb, 4, 5, 16, 4, 4, &[], &[]).unwrap();
1314 assert_eq!(l0.iter().map(|r| r.poc).collect::<Vec<_>>(), vec![2, 0, 6, 8]);
1316 assert_eq!(l1.iter().map(|r| r.poc).collect::<Vec<_>>(), vec![6, 8, 2, 0]);
1318 }
1319
1320 #[test]
1321 fn b_ref_list1_swap_when_equal() {
1322 let dpb = vec![ref_at(4, 2), ref_at(2, 1), ref_at(0, 0)];
1325 let (l0, l1) = build_ref_list_b(&dpb, 6, 3, 16, 3, 3, &[], &[]).unwrap();
1326 assert_eq!(l0.iter().map(|r| r.poc).collect::<Vec<_>>(), vec![4, 2, 0]);
1327 assert_eq!(l1.iter().map(|r| r.poc).collect::<Vec<_>>(), vec![2, 4, 0]);
1328 }
1329
1330 #[test]
1331 fn frame_num_gaps_insert_placeholders() {
1332 let mut d = Decoder::new();
1333 d.prev_ref_frame_num = 2;
1334 d.insert_frame_num_gaps(5, 16, 8, 16, 16);
1336 let fns: Vec<u32> = d.refs.iter().map(|r| r.frame_num).collect();
1337 assert_eq!(fns, vec![4, 3], "most-recent placeholder at the front");
1338 assert_eq!(d.prev_ref_frame_num, 4);
1339 assert!(d.refs.iter().all(|r| r.py.iter().all(|&p| p == 128)), "grey fill");
1340 }
1341
1342 #[test]
1343 fn frame_num_gaps_wrap_and_noop() {
1344 let mut d = Decoder::new();
1346 d.prev_ref_frame_num = 14;
1347 d.insert_frame_num_gaps(1, 16, 8, 16, 16);
1348 assert_eq!(d.refs.iter().map(|r| r.frame_num).collect::<Vec<_>>(), vec![0, 15]);
1349 let mut d = Decoder::new();
1351 d.prev_ref_frame_num = 3;
1352 d.insert_frame_num_gaps(4, 16, 8, 16, 16);
1353 assert!(d.refs.is_empty());
1354 }
1355
1356 #[test]
1357 fn missing_param_sets_errors() {
1358 let mut d = Decoder::new();
1359 let nal = rusty_h264_common::NalUnit::new(3, NalUnitType::IdrSlice, vec![0x88, 0x80]);
1362 let err = d.decode(&nal.to_annex_b()).unwrap_err();
1363 assert_eq!(err, DecodeError::MissingParameterSet);
1364 }
1365}