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
35pub fn edc_stats_report() {
37 mb16::edcstat::report();
38}
39
40#[doc(hidden)]
43fn abl_deblock() -> bool {
46 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
47 *ON.get_or_init(|| std::env::var("RFF_ABL_DEBLOCK").map_or(false, |v| v != "0"))
48}
49
50pub mod cabac_test {
51 pub use crate::cabac::Cabac;
52 pub use crate::mb16::b_inter_shape;
53 pub use crate::mb16::parse_mb_type_b;
54}
55
56use mb16::{FrameDecoder, GridPool, WeightTable};
57use rusty_h264_common::bit_reader::OutOfData;
58use rusty_h264_common::nal::{emulation_unprevent, split_annex_b};
59use rusty_h264_common::{BitReader, NalUnitType, YuvFrame};
60
61#[derive(Debug, Clone, PartialEq, Eq)]
63pub enum DecodeError {
64 Truncated,
66 MissingParameterSet,
68 Unsupported(&'static str),
70}
71
72impl From<OutOfData> for DecodeError {
73 fn from(_: OutOfData) -> Self {
74 DecodeError::Truncated
75 }
76}
77
78impl core::fmt::Display for DecodeError {
79 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
80 match self {
81 DecodeError::Truncated => f.write_str("bitstream truncated"),
82 DecodeError::MissingParameterSet => f.write_str("slice before SPS/PPS"),
83 DecodeError::Unsupported(s) => write!(f, "unsupported coding tool: {s}"),
84 }
85 }
86}
87
88impl std::error::Error for DecodeError {}
89
90pub(crate) type Ref = std::sync::Arc<RefFrame>;
99
100#[derive(Debug, Clone, Default)]
101#[allow(dead_code)]
102pub(crate) struct RefFrame {
103 pub py: Vec<u8>,
109 pub pu: Vec<u8>,
110 pub pv: Vec<u8>,
111 pub cw: usize,
112 pub ch: usize,
113 pub frame_num: u32,
115 pub poc: i32,
117 pub mv: Vec<(i32, i32)>,
121 pub ref_idx: Vec<i32>,
122 pub mv1: Vec<(i32, i32)>,
129 pub ref_idx1: Vec<i32>,
130 pub ref_poc: Vec<i32>,
134 pub w4: usize,
135 pub long_term: bool,
139 pub long_term_idx: u32,
140}
141
142pub(crate) const LPAD: usize = 16;
146pub(crate) const CPAD: usize = 8;
147
148impl RefFrame {
149 #[inline]
150 pub fn lstride(&self) -> usize {
151 self.cw + 2 * LPAD
152 }
153 #[inline]
154 pub fn cstride(&self) -> usize {
155 self.cw / 2 + 2 * CPAD
156 }
157}
158
159#[derive(Clone, Copy)]
161enum Mmco {
162 Unref(u32),
164 UnrefLong(u32),
166 AssignLong(u32, u32),
168 MaxLong(u32),
170 Reset,
172 CurrentLong(u32),
174}
175
176struct PendingPic {
180 fd: mb16::FrameDecoder,
181 frame_num: u32,
182 poc: i32,
183 next_mb: usize,
184 total_mb: usize,
185 slice_count: u16,
186 deblock: bool,
187 filter_offset_a: i32,
188 filter_offset_b: i32,
189 crop_r: usize,
190 crop_b: usize,
191 max_refs: usize,
192 log2_max_frame_num: u32,
193 is_reference: bool,
196 idr_long_term: bool,
197 mmco_ops: Vec<Mmco>,
198}
199
200fn no_pool() -> bool {
202 use std::sync::atomic::{AtomicU8, Ordering};
203 static ON: AtomicU8 = AtomicU8::new(0);
204 match ON.load(Ordering::Relaxed) {
205 0 => {
206 let v = std::env::var_os("RS_H264_NO_POOL").is_some_and(|v| v == "1");
207 ON.store(if v { 1 } else { 2 }, Ordering::Relaxed);
208 v
209 }
210 n => n == 1,
211 }
212}
213
214#[derive(Default)]
217pub struct Decoder {
218 sps: std::collections::HashMap<u32, Sps>,
221 pps: std::collections::HashMap<u32, Pps>,
222 refs: Vec<Ref>,
224 cur: Option<PendingPic>,
226 poc: PocState,
230 last_poc: i32,
232 prev_ref_frame_num: u32,
235 grid_pool: GridPool,
238 plane_pool: Vec<Vec<u8>>,
241 retired: Vec<Ref>,
246}
247
248#[derive(Default)]
250struct PocState {
251 prev_msb: i32,
252 prev_lsb: i32,
253 prev_frame_num: u32,
254 prev_frame_num_offset: i64,
255}
256
257impl Decoder {
258 pub fn new() -> Self {
260 Self::default()
261 }
262
263 pub fn decode(&mut self, annex_b: &[u8]) -> Result<Option<YuvFrame>, DecodeError> {
266 let _g = rusty_h264_common::prof::scope(rusty_h264_common::prof::Stage::Total);
267 let mut frame = None;
268 let nals = {
272 let _s = rusty_h264_common::prof::scope(rusty_h264_common::prof::Stage::DecNalSplit);
273 split_annex_b(annex_b)
274 };
275 for nal in nals {
276 if nal.is_empty() {
277 continue;
278 }
279 let nal_type = NalUnitType::from_id(nal[0]);
280 let rbsp = {
281 let _s = rusty_h264_common::prof::scope(rusty_h264_common::prof::Stage::DecRbsp);
282 emulation_unprevent(&nal[1..])
283 };
284 match nal_type {
285 NalUnitType::Sps => {
286 let s = Sps::parse(&rbsp)?;
287 self.sps.insert(s.seq_parameter_set_id, s);
288 }
289 NalUnitType::Pps => {
290 let p = Pps::parse(&rbsp)?;
291 self.pps.insert(p.pic_parameter_set_id, p);
292 }
293 NalUnitType::IdrSlice | NalUnitType::NonIdrSlice => {
294 let nal_ref_idc = (nal[0] >> 5) & 3;
295 let is_idr = nal_type == NalUnitType::IdrSlice;
296 if let Some(f) = self.decode_slice(&rbsp, is_idr, nal_ref_idc)? {
297 frame = Some(f);
298 }
299 }
300 _ => {} }
302 }
303 Ok(frame)
304 }
305
306 pub fn decode_stream(&mut self, annex_b: &[u8]) -> Result<Vec<YuvFrame>, DecodeError> {
314 let mut out = Vec::new();
315 let mut gop: Vec<(i32, YuvFrame)> = Vec::new();
316 for au in split_access_units(annex_b) {
317 if au_is_idr(au) {
318 flush_gop(&mut gop, &mut out); }
320 if let Some(frame) = self.decode(au)? {
321 gop.push((self.last_poc, frame));
322 }
323 }
324 flush_gop(&mut gop, &mut out);
325 Ok(out)
326 }
327
328 fn decode_slice(
329 &mut self,
330 rbsp: &[u8],
331 is_idr: bool,
332 nal_ref_idc: u8,
333 ) -> Result<Option<YuvFrame>, DecodeError> {
334 let mut r = BitReader::new(rbsp);
335 let first_mb_in_slice = r.read_ue()? as usize;
337 let slice_type = r.read_ue()?;
338 let is_p = matches!(slice_type, 0 | 5);
339 let is_b = matches!(slice_type, 1 | 6);
340 let is_i = matches!(slice_type, 2 | 7);
341 if !is_p && !is_b && !is_i {
342 return Err(DecodeError::Unsupported("SP/SI slices"));
343 }
344 let pic_parameter_set_id = r.read_ue()?;
346 let pps = self.pps.get(&pic_parameter_set_id).cloned().ok_or(DecodeError::MissingParameterSet)?;
347 let sps = self.sps.get(&pps.seq_parameter_set_id).cloned().ok_or(DecodeError::MissingParameterSet)?;
348 let sps = &sps;
349 let pps = &pps;
350 let cabac = pps.entropy_coding_mode_flag;
354 let frame_num = r.read_bits(sps.log2_max_frame_num)?;
355 if is_idr {
356 let _idr_pic_id = r.read_ue()?;
357 }
358 let mut poc_lsb = 0u32;
361 let mut delta_poc_bottom = 0i32;
362 if sps.pic_order_cnt_type == 0 {
363 poc_lsb = r.read_bits(sps.log2_max_pic_order_cnt_lsb)?;
364 if pps.bottom_field_pic_order_present {
365 delta_poc_bottom = r.read_se()?;
366 }
367 } else if sps.pic_order_cnt_type == 1 && !sps.delta_pic_order_always_zero {
368 let _delta_pic_order_cnt_0 = r.read_se()?;
369 if pps.bottom_field_pic_order_present {
370 let _delta_pic_order_cnt_1 = r.read_se()?;
371 }
372 }
373 let pic_poc = if first_mb_in_slice == 0 {
376 self.compute_poc(sps, is_idr, nal_ref_idc, frame_num, poc_lsb, delta_poc_bottom)
377 } else {
378 self.cur.as_ref().map_or(0, |p| p.poc)
379 };
380 if pps.redundant_pic_cnt_present_flag {
385 let redundant_pic_cnt = r.read_ue()?;
386 if redundant_pic_cnt != 0 {
387 return Ok(None);
388 }
389 }
390 if std::env::var_os("RH264_DUMP_MB").is_some() {
391 eprintln!(
392 "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}"
393 );
394 }
395 let direct_spatial = if is_b { r.read_bit()? } else { true };
397 let mut num_ref_idx_l0 = pps.num_ref_idx_l0_default as usize;
398 let mut num_ref_idx_l1 = pps.num_ref_idx_l1_default as usize;
399 let mut reorder_l0: Vec<(u32, u32)> = Vec::new();
400 let mut reorder_l1: Vec<(u32, u32)> = Vec::new();
401 if is_p || is_b {
402 if r.read_bit()? {
404 num_ref_idx_l0 = (r.read_ue()? + 1) as usize;
405 if is_b {
406 num_ref_idx_l1 = (r.read_ue()? + 1) as usize;
407 }
408 }
409 if r.read_bit()? {
411 parse_ref_pic_list_modification(&mut r, &mut reorder_l0)?;
412 }
413 if is_b && r.read_bit()? {
414 parse_ref_pic_list_modification(&mut r, &mut reorder_l1)?;
416 }
417 }
418 let weights = if is_p && pps.weighted_pred {
423 Some(parse_pred_weight_table(&mut r, num_ref_idx_l0, 0, false)?)
424 } else if is_b && pps.weighted_bipred_idc == 1 {
425 return Err(DecodeError::Unsupported("explicit B weighted prediction"));
426 } else {
427 None
428 };
429 let mut idr_long_term = false;
433 let mut mmco_ops: Vec<Mmco> = Vec::new();
434 if nal_ref_idc == 0 {
435 } else if is_idr {
437 let _no_output_of_prior_pics = r.read_bit()?;
438 idr_long_term = r.read_bit()?; } else if r.read_bit()? {
440 loop {
442 let op = r.read_ue()?;
443 match op {
444 0 => break,
445 1 => mmco_ops.push(Mmco::Unref(r.read_ue()?)),
446 2 => mmco_ops.push(Mmco::UnrefLong(r.read_ue()?)),
447 3 => {
448 let diff = r.read_ue()?;
449 let idx = r.read_ue()?;
450 mmco_ops.push(Mmco::AssignLong(diff, idx));
451 }
452 4 => mmco_ops.push(Mmco::MaxLong(r.read_ue()?)),
453 5 => mmco_ops.push(Mmco::Reset),
454 6 => mmco_ops.push(Mmco::CurrentLong(r.read_ue()?)),
455 _ => return Err(DecodeError::Unsupported("invalid MMCO")),
456 }
457 if mmco_ops.len() > 128 {
458 return Err(DecodeError::Truncated);
459 }
460 }
461 }
462 let cabac_init_idc = if cabac && !is_i {
466 let v = r.read_ue()?;
467 if v > 2 {
468 return Err(DecodeError::Unsupported("invalid cabac_init_idc"));
469 }
470 v
471 } else {
472 0
473 };
474 let slice_qp_delta = r.read_se()?;
475 let mut deblock = true;
480 let (mut filter_offset_a, mut filter_offset_b) = (0i32, 0i32);
481 if pps.deblocking_filter_control_present_flag {
482 let disable_deblocking_filter_idc = r.read_ue()?;
483 deblock = disable_deblocking_filter_idc != 1;
486 if disable_deblocking_filter_idc != 1 {
487 filter_offset_a = r.read_se()? * 2;
489 filter_offset_b = r.read_se()? * 2;
490 }
491 }
492 let slice_qp = (pps.pic_init_qp + slice_qp_delta).clamp(0, 51) as u8;
493
494 if first_mb_in_slice == 0 && !is_idr && sps.gaps_in_frame_num_allowed {
497 self.insert_frame_num_gaps(
498 frame_num,
499 1u32 << sps.log2_max_frame_num,
500 sps.max_num_ref_frames.max(1) as usize,
501 sps.pic_width_in_mbs * 16,
502 sps.pic_height_in_mbs * 16,
503 );
504 }
505
506 let max_fn = 1u32 << sps.log2_max_frame_num;
509 let (ref_list0, ref_list1) = if is_b {
510 build_ref_list_b(
511 &self.refs, pic_poc, frame_num, max_fn,
512 num_ref_idx_l0, num_ref_idx_l1, &reorder_l0, &reorder_l1,
513 )?
514 } else if is_p {
515 (build_ref_list_p(&self.refs, frame_num, max_fn, num_ref_idx_l0, &reorder_l0)?, Vec::new())
516 } else {
517 (Vec::new(), Vec::new())
518 };
519 if first_mb_in_slice == 0 {
523 if is_idr {
524 self.refs.clear();
525 }
526 let _g_setup = rusty_h264_common::prof::scope(rusty_h264_common::prof::Stage::DecSetup);
537 let mut fd = FrameDecoder::with_pool(
538 sps.pic_width_in_mbs,
539 sps.pic_height_in_mbs,
540 slice_qp,
541 pps.chroma_qp_index_offset,
542 ref_list0,
543 num_ref_idx_l0,
544 pps.constrained_intra_pred_flag,
545 pps.transform_8x8_mode_flag,
546 sps.profile_idc != 66, if no_pool() { GridPool::default() } else { std::mem::take(&mut self.grid_pool) },
551 );
552 if is_b {
553 fd.set_b_context(
554 ref_list1,
555 num_ref_idx_l1,
556 direct_spatial,
557 pic_poc,
558 pps.weighted_bipred_idc,
559 sps.direct_8x8_inference,
560 );
561 }
562 if sps.has_scaling || pps.pic_scaling_matrix_present {
563 let (s4, s8) = resolve_scaling(sps, pps);
564 fd.set_scaling(s4, s8);
565 }
566 if let Some(w) = weights {
567 fd.set_weights(w);
568 }
569 if let Some(prev) = self.cur.take() {
577 if prev.next_mb < prev.total_mb {
578 return Err(DecodeError::Truncated);
579 }
580 }
581 self.cur = Some(PendingPic {
582 fd,
583 frame_num,
584 poc: pic_poc,
585 next_mb: 0,
586 total_mb: sps.pic_width_in_mbs * sps.pic_height_in_mbs,
587 slice_count: 0,
588 deblock,
589 filter_offset_a,
590 filter_offset_b,
591 crop_r: sps.frame_crop_right as usize,
592 crop_b: sps.frame_crop_bottom as usize,
593 max_refs: sps.max_num_ref_frames.max(1) as usize,
594 log2_max_frame_num: sps.log2_max_frame_num,
595 is_reference: nal_ref_idc != 0,
596 idr_long_term,
597 mmco_ops,
598 });
599 } else {
600 let Some(pic) = self.cur.as_mut() else {
602 return Err(DecodeError::Unsupported("slice continues a missing picture"));
603 };
604 pic.fd.begin_slice(slice_qp, ref_list0, num_ref_idx_l0);
605 if is_b {
606 pic.fd.set_b_context(
607 ref_list1,
608 num_ref_idx_l1,
609 direct_spatial,
610 pic.poc,
611 pps.weighted_bipred_idc,
612 sps.direct_8x8_inference,
613 );
614 }
615 if sps.has_scaling || pps.pic_scaling_matrix_present {
616 let (s4, s8) = resolve_scaling(sps, pps);
617 pic.fd.set_scaling(s4, s8);
618 }
619 if let Some(w) = weights {
620 pic.fd.set_weights(w);
621 }
622 pic.deblock = deblock;
624 pic.filter_offset_a = filter_offset_a;
625 pic.filter_offset_b = filter_offset_b;
626 pic.idr_long_term |= idr_long_term;
627 pic.mmco_ops.extend(mmco_ops);
628 }
629
630 let pic = self.cur.as_mut().expect("pending picture set above");
631 pic.fd.set_deblock_params(deblock && !abl_deblock(), filter_offset_a, filter_offset_b);
635 let first = first_mb_in_slice.min(pic.total_mb);
636 let next = if cabac {
637 r.align_to_byte().map_err(|_| DecodeError::Truncated)?;
639 let (data, start) = (r.data(), r.bit_pos() / 8);
640 pic.fd
641 .decode_slice_data_cabac(data, start, slice_qp, cabac_init_idc, is_i, is_p, first)
642 } else {
643 pic.fd.decode_slice_data(&mut r, is_p, first)
644 }
645 .map_err(|e| match e {
646 mb16::MbError::Truncated => DecodeError::Truncated,
647 mb16::MbError::Unsupported(s) => DecodeError::Unsupported(s),
648 })?;
649 pic.next_mb = next;
650 pic.slice_count += 1;
651 if std::env::var_os("RH264_DUMP_MB").is_some() {
652 eprintln!(
653 " slice decoded {}/{} MBs{}",
654 next,
655 pic.total_mb,
656 if next < pic.total_mb { " <-- INCOMPLETE" } else { "" }
657 );
658 }
659
660 if pic.next_mb < pic.total_mb {
661 return Ok(None); }
663
664 let pic = self.cur.take().expect("pending picture");
666 let PendingPic {
667 mut fd,
668 frame_num,
669 poc,
670 deblock,
671 filter_offset_a,
672 filter_offset_b,
673 crop_r,
674 crop_b,
675 max_refs,
676 log2_max_frame_num,
677 is_reference,
678 idr_long_term,
679 mmco_ops,
680 ..
681 } = pic;
682 self.last_poc = poc;
683 if deblock && !abl_deblock() {
691 fd.deblock(filter_offset_a, filter_offset_b);
692 }
693 let reference = if is_reference {
696 let _dg = rusty_h264_common::prof::scope(rusty_h264_common::prof::Stage::DpbClone);
697 Some(fd.as_reference_pooled(&mut self.plane_pool))
698 } else {
699 None
701 };
702 let _fg = rusty_h264_common::prof::scope(rusty_h264_common::prof::Stage::Finalize);
703 if let Some(mut reference) = reference {
704 reference.frame_num = frame_num;
705 reference.poc = poc;
706 if std::env::var_os("RH264_DUMP_MB").is_some() {
707 eprintln!("DPB-ADD fn={frame_num} poc={poc}");
708 }
709 if idr_long_term {
710 reference.long_term = true;
711 reference.long_term_idx = 0;
712 }
713 self.prev_ref_frame_num =
715 self.apply_ref_marking(reference, &mmco_ops, frame_num, log2_max_frame_num, max_refs);
716 }
717 let (frame, pool) = fd.into_frame_recycle(crop_r, crop_b);
718 self.grid_pool = pool;
719 self.reclaim_retired();
720 Ok(Some(frame))
721 }
722
723 fn reclaim_retired(&mut self) {
728 for arc in self.retired.drain(..) {
729 if let Ok(rf) = std::sync::Arc::try_unwrap(arc) {
730 self.plane_pool.push(rf.py);
731 self.plane_pool.push(rf.pu);
732 self.plane_pool.push(rf.pv);
733 }
734 }
735 self.plane_pool.truncate(18);
738 }
739
740 fn insert_frame_num_gaps(&mut self, frame_num: u32, max_fn: u32, max_refs: usize, w: usize, h: usize) {
746 if max_fn == 0 {
747 return;
748 }
749 let start = (self.prev_ref_frame_num + 1) % max_fn;
750 let gap = (frame_num + max_fn - start) % max_fn;
751 if gap == 0 {
752 return;
753 }
754 let cap = max_refs.max(1);
760 let n = (gap as usize).min(cap);
761 let (cw, ch) = (w, h);
762 let mut expected = (frame_num + max_fn - n as u32) % max_fn;
763 for _ in 0..n {
764 self.refs.insert(
765 0,
766 std::sync::Arc::new(RefFrame {
767 py: vec![128; (cw + 2 * LPAD) * (ch + 2 * LPAD)],
769 pu: vec![128; (cw / 2 + 2 * CPAD) * (ch / 2 + 2 * CPAD)],
770 pv: vec![128; (cw / 2 + 2 * CPAD) * (ch / 2 + 2 * CPAD)],
771 cw,
772 ch,
773 frame_num: expected,
774 poc: 0,
775 mv: Vec::new(),
776 ref_idx: Vec::new(),
777 mv1: Vec::new(),
778 ref_idx1: Vec::new(),
779 ref_poc: Vec::new(),
780 w4: 0,
781 long_term: false,
782 long_term_idx: 0,
783 }),
784 );
785 self.refs.truncate(cap);
786 expected = (expected + 1) % max_fn;
787 }
788 self.prev_ref_frame_num = (frame_num + max_fn - 1) % max_fn;
789 }
790
791 pub fn last_poc(&self) -> i32 {
795 self.last_poc
796 }
797
798 fn compute_poc(
802 &mut self,
803 sps: &Sps,
804 is_idr: bool,
805 nal_ref_idc: u8,
806 frame_num: u32,
807 poc_lsb: u32,
808 delta_bottom: i32,
809 ) -> i32 {
810 match sps.pic_order_cnt_type {
811 0 => {
812 let max_lsb = 1i32 << sps.log2_max_pic_order_cnt_lsb;
813 let (prev_msb, prev_lsb) =
814 if is_idr { (0, 0) } else { (self.poc.prev_msb, self.poc.prev_lsb) };
815 let lsb = poc_lsb as i32;
816 let msb = if lsb < prev_lsb && prev_lsb - lsb >= max_lsb / 2 {
817 prev_msb + max_lsb
818 } else if lsb > prev_lsb && lsb - prev_lsb > max_lsb / 2 {
819 prev_msb - max_lsb
820 } else {
821 prev_msb
822 };
823 let top = msb + lsb;
824 let poc = top.min(top + delta_bottom);
825 if nal_ref_idc != 0 {
826 self.poc.prev_msb = msb;
827 self.poc.prev_lsb = lsb;
828 }
829 poc
830 }
831 2 => {
832 let max_fn = 1i64 << sps.log2_max_frame_num;
833 let offset = if is_idr {
834 0
835 } else if self.poc.prev_frame_num > frame_num {
836 self.poc.prev_frame_num_offset + max_fn
837 } else {
838 self.poc.prev_frame_num_offset
839 };
840 let poc = if is_idr {
841 0
842 } else {
843 2 * (offset + frame_num as i64) - i64::from(nal_ref_idc == 0)
844 };
845 self.poc.prev_frame_num_offset = offset;
846 self.poc.prev_frame_num = frame_num;
847 poc as i32
848 }
849 _ => {
850 self.poc.prev_frame_num = frame_num;
851 frame_num as i32 * 2
852 }
853 }
854 }
855
856 fn apply_ref_marking(
867 &mut self,
868 mut reference: RefFrame,
869 ops: &[Mmco],
870 frame_num: u32,
871 log2_max_frame_num: u32,
872 max_refs: usize,
873 ) -> u32 {
874 let max = 1i64 << log2_max_frame_num;
875 let curr = frame_num as i64;
876 let pic_num = |rf: &RefFrame| -> i64 {
877 if (rf.frame_num as i64) > curr {
878 rf.frame_num as i64 - max
879 } else {
880 rf.frame_num as i64
881 }
882 };
883
884 if ops.is_empty() {
885 let out_fn = reference.frame_num;
889 self.refs.insert(0, std::sync::Arc::new(reference));
890 while self.refs.len() > max_refs {
891 match self.refs.iter().rposition(|r| !r.long_term) {
892 Some(pos) => {
893 let evicted = self.refs.remove(pos);
896 self.retired.push(evicted);
897 }
898 None => break,
899 }
900 }
901 return out_fn;
902 }
903
904 for &op in ops {
906 match op {
907 Mmco::Unref(diff) => {
908 let target = curr - (diff as i64 + 1);
909 self.refs.retain(|r| r.long_term || pic_num(r) != target);
910 }
911 Mmco::UnrefLong(ltpn) => {
912 self.refs.retain(|r| !(r.long_term && r.long_term_idx == ltpn));
913 }
914 Mmco::AssignLong(diff, idx) => {
915 let target = curr - (diff as i64 + 1);
916 self.refs.retain(|r| !(r.long_term && r.long_term_idx == idx));
917 for r in self.refs.iter_mut() {
918 if !r.long_term && pic_num(r) == target {
919 let r = std::sync::Arc::make_mut(r);
921 r.long_term = true;
922 r.long_term_idx = idx;
923 }
924 }
925 }
926 Mmco::MaxLong(max_plus1) => {
927 self.refs.retain(|r| !(r.long_term && r.long_term_idx + 1 > max_plus1));
928 }
929 Mmco::Reset => {
930 self.refs.clear();
931 reference.frame_num = 0;
932 }
933 Mmco::CurrentLong(idx) => {
934 self.refs.retain(|r| !(r.long_term && r.long_term_idx == idx));
935 reference.long_term = true;
936 reference.long_term_idx = idx;
937 }
938 }
939 }
940 let out_fn = reference.frame_num;
941 self.refs.insert(0, std::sync::Arc::new(reference));
942 let cap = max_refs.max(16);
944 if self.refs.len() > cap {
945 self.refs.truncate(cap);
946 }
947 out_fn
948 }
949}
950
951fn flush_gop(gop: &mut Vec<(i32, YuvFrame)>, out: &mut Vec<YuvFrame>) {
953 gop.sort_by_key(|(poc, _)| *poc);
954 out.extend(gop.drain(..).map(|(_, f)| f));
955}
956
957fn au_is_idr(au: &[u8]) -> bool {
959 split_annex_b(au)
960 .iter()
961 .any(|n| !n.is_empty() && NalUnitType::from_id(n[0]) == NalUnitType::IdrSlice)
962}
963
964pub fn split_access_units(stream: &[u8]) -> Vec<&[u8]> {
972 let mut codes: Vec<(usize, bool)> = Vec::new();
974 let mut i = 0;
975 while i + 3 <= stream.len() {
976 if stream[i] == 0 && stream[i + 1] == 0 && stream[i + 2] == 1 {
977 let nal_type = NalUnitType::from_id(stream.get(i + 3).copied().unwrap_or(0));
978 let is_vcl = matches!(nal_type, NalUnitType::IdrSlice | NalUnitType::NonIdrSlice);
979 let sc = if i > 0 && stream[i - 1] == 0 { i - 1 } else { i };
981 codes.push((sc, is_vcl));
982 i += 3;
983 } else {
984 i += 1;
985 }
986 }
987 if codes.is_empty() {
988 return vec![stream];
989 }
990 let mut aus = Vec::new();
991 let mut start = codes[0].0;
992 for k in 0..codes.len() {
993 if codes[k].1 {
994 let end = codes.get(k + 1).map_or(stream.len(), |c| c.0);
995 aus.push(&stream[start..end]);
996 start = end;
997 }
998 }
999 aus
1000}
1001
1002fn parse_pred_weight_table(
1005 r: &mut BitReader,
1006 num_l0: usize,
1007 num_l1: usize,
1008 is_b: bool,
1009) -> Result<WeightTable, DecodeError> {
1010 let luma_log2_denom = r.read_ue()? as i32;
1011 let chroma_log2_denom = r.read_ue()? as i32;
1012 if !(0..=7).contains(&luma_log2_denom) || !(0..=7).contains(&chroma_log2_denom) {
1016 return Err(DecodeError::Unsupported("invalid weight denom"));
1017 }
1018 let mut wt = WeightTable {
1019 luma_log2_denom,
1020 chroma_log2_denom,
1021 ..Default::default()
1022 };
1023 let lists: &[(usize, usize)] = if is_b {
1024 &[(0, num_l0), (1, num_l1)]
1025 } else {
1026 &[(0, num_l0)]
1027 };
1028 for &(list, n) in lists {
1029 let mut luma = Vec::with_capacity(n);
1030 let mut chroma = Vec::with_capacity(n);
1031 for _ in 0..n {
1032 let (mut lw, mut lo) = (1 << luma_log2_denom, 0);
1033 if r.read_bit()? {
1034 lw = r.read_se()?;
1035 lo = r.read_se()?;
1036 }
1037 luma.push((lw, lo));
1038 let mut ch = [(1 << chroma_log2_denom, 0); 2];
1039 if r.read_bit()? {
1040 for slot in ch.iter_mut() {
1041 *slot = (r.read_se()?, r.read_se()?);
1042 }
1043 }
1044 chroma.push(ch);
1045 }
1046 wt.luma[list] = luma;
1047 wt.chroma[list] = chroma;
1048 }
1049 Ok(wt)
1050}
1051
1052fn resolve_scaling(sps: &Sps, pps: &Pps) -> ([[i32; 16]; 6], [[i32; 64]; 2]) {
1057 use crate::params::{
1058 DEFAULT_4X4_INTER, DEFAULT_4X4_INTRA, DEFAULT_8X8_INTER, DEFAULT_8X8_INTRA,
1059 };
1060 const ZZ4: [usize; 16] = [0, 1, 4, 8, 5, 2, 3, 6, 9, 12, 13, 10, 7, 11, 14, 15];
1061 const ZZ8: [usize; 64] = [
1063 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,
1064 20, 13, 6, 7, 14, 21, 28, 35, 42, 49, 56, 57, 50, 43, 36, 29, 22, 15, 23, 30, 37, 44, 51,
1065 58, 59, 52, 45, 38, 31, 39, 46, 53, 60, 61, 54, 47, 55, 62, 63,
1066 ];
1067 let mut z4 = [[16u8; 16]; 6];
1070 for i in 0..6 {
1071 z4[i] = if pps.pic_scaling_matrix_present {
1072 if pps.scaling_present_4x4[i] {
1073 pps.scaling_4x4[i]
1074 } else {
1075 match i {
1076 0 if sps.has_scaling => sps.scaling_4x4[0],
1077 0 => DEFAULT_4X4_INTRA,
1078 3 if sps.has_scaling => sps.scaling_4x4[3],
1079 3 => DEFAULT_4X4_INTER,
1080 _ => z4[i - 1],
1081 }
1082 }
1083 } else {
1084 sps.scaling_4x4[i]
1085 };
1086 }
1087 let mut z8 = [[16u8; 64]; 2];
1088 for (i, list) in z8.iter_mut().enumerate() {
1089 *list = if pps.pic_scaling_matrix_present {
1090 if pps.scaling_present_8x8[i] {
1091 pps.scaling_8x8[i]
1092 } else if sps.has_scaling {
1093 sps.scaling_8x8[i]
1094 } else if i == 0 {
1095 DEFAULT_8X8_INTRA
1096 } else {
1097 DEFAULT_8X8_INTER
1098 }
1099 } else {
1100 sps.scaling_8x8[i]
1101 };
1102 }
1103 let mut out4 = [[16i32; 16]; 6];
1104 for (li, list) in out4.iter_mut().enumerate() {
1105 for k in 0..16 {
1106 list[ZZ4[k]] = z4[li][k] as i32;
1107 }
1108 }
1109 let mut out8 = [[16i32; 64]; 2];
1110 for (li, list) in out8.iter_mut().enumerate() {
1111 for k in 0..64 {
1112 list[ZZ8[k]] = z8[li][k] as i32;
1113 }
1114 }
1115 (out4, out8)
1116}
1117
1118fn parse_ref_pic_list_modification(
1121 r: &mut BitReader,
1122 out: &mut Vec<(u32, u32)>,
1123) -> Result<(), DecodeError> {
1124 loop {
1125 let idc = r.read_ue()?;
1126 if idc == 3 {
1127 break;
1128 }
1129 if idc > 3 {
1130 return Err(DecodeError::Unsupported("invalid ref_pic_list_modification"));
1131 }
1132 let val = r.read_ue()?; out.push((idc, val));
1134 if out.len() > 64 {
1135 return Err(DecodeError::Truncated); }
1137 }
1138 Ok(())
1139}
1140
1141fn build_ref_list_p(
1145 dpb: &[Ref],
1146 curr_frame_num: u32,
1147 max_frame_num: u32,
1148 num_active: usize,
1149 mods: &[(u32, u32)],
1150) -> Result<Vec<Ref>, DecodeError> {
1151 let curr = curr_frame_num as i64;
1152 let max = max_frame_num as i64;
1153 let pic_num = |fnum: u32| -> i64 {
1154 let f = fnum as i64;
1155 if f > curr { f - max } else { f }
1156 };
1157 let mut init: Vec<Ref> = dpb.iter().filter(|r| !r.long_term).cloned().collect();
1158 init.sort_by_key(|rf| core::cmp::Reverse(pic_num(rf.frame_num)));
1159 let mut long: Vec<Ref> = dpb.iter().filter(|r| r.long_term).cloned().collect();
1160 long.sort_by_key(|rf| rf.long_term_idx);
1161 init.extend(long);
1162 apply_list_modification(init, curr_frame_num, max_frame_num, num_active, mods)
1163}
1164
1165#[allow(clippy::too_many_arguments)]
1170fn build_ref_list_b(
1171 dpb: &[Ref],
1172 curr_poc: i32,
1173 curr_frame_num: u32,
1174 max_frame_num: u32,
1175 num0: usize,
1176 num1: usize,
1177 mods0: &[(u32, u32)],
1178 mods1: &[(u32, u32)],
1179) -> Result<(Vec<Ref>, Vec<Ref>), DecodeError> {
1180 let mut less: Vec<Ref> =
1181 dpb.iter().filter(|r| !r.long_term && r.poc < curr_poc).cloned().collect();
1182 let mut greater: Vec<Ref> =
1183 dpb.iter().filter(|r| !r.long_term && r.poc > curr_poc).cloned().collect();
1184 let mut long: Vec<Ref> = dpb.iter().filter(|r| r.long_term).cloned().collect();
1185 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);
1188
1189 let mut init0 = less.clone();
1190 init0.extend(greater.clone());
1191 init0.extend(long.clone());
1192 let mut init1 = greater;
1193 init1.extend(less);
1194 init1.extend(long);
1195
1196 let eq_len = num0.min(num1).min(init0.len()).min(init1.len());
1199 if num1 > 1
1200 && init1.len() > 1
1201 && (0..eq_len).all(|i| same_picture(&init0[i], &init1[i]))
1202 && eq_len == num1.min(init1.len())
1203 && eq_len == num0.min(init0.len())
1204 {
1205 init1.swap(0, 1);
1206 }
1207
1208 let list0 = apply_list_modification(init0, curr_frame_num, max_frame_num, num0, mods0)?;
1209 let list1 = apply_list_modification(init1, curr_frame_num, max_frame_num, num1, mods1)?;
1210 Ok((list0, list1))
1211}
1212
1213fn same_picture(a: &RefFrame, b: &RefFrame) -> bool {
1215 a.long_term == b.long_term
1216 && if a.long_term { a.long_term_idx == b.long_term_idx } else { a.poc == b.poc }
1217}
1218
1219fn apply_list_modification(
1224 init: Vec<Ref>,
1225 curr_frame_num: u32,
1226 max_frame_num: u32,
1227 num_active: usize,
1228 mods: &[(u32, u32)],
1229) -> Result<Vec<Ref>, DecodeError> {
1230 if mods.is_empty() {
1231 let mut init = init;
1232 init.truncate(num_active.max(1));
1233 return Ok(init);
1234 }
1235 let curr = curr_frame_num as i64;
1236 let max = max_frame_num as i64;
1237 let mut list = init.clone();
1238 let mut pic_num_pred = curr;
1239 let mut refidx = 0usize;
1240 for &(idc, val) in mods {
1241 let matches: Box<dyn Fn(&RefFrame) -> bool> = if idc == 2 {
1242 Box::new(move |r: &RefFrame| r.long_term && r.long_term_idx == val)
1243 } else {
1244 let abs_diff = (val as i64) + 1;
1245 let no_wrap = if idc == 0 {
1246 let x = pic_num_pred - abs_diff;
1247 if x < 0 { x + max } else { x }
1248 } else {
1249 let x = pic_num_pred + abs_diff;
1250 if x >= max { x - max } else { x }
1251 };
1252 pic_num_pred = no_wrap;
1253 let target = if no_wrap > curr { no_wrap - max } else { no_wrap };
1254 Box::new(move |r: &RefFrame| {
1255 let pn = if r.frame_num as i64 > curr {
1256 r.frame_num as i64 - max
1257 } else {
1258 r.frame_num as i64
1259 };
1260 !r.long_term && pn == target
1261 })
1262 };
1263 let found = init.iter().find(|r| matches(r)).cloned();
1264 let Some(found) = found else {
1265 if std::env::var_os("RH264_DUMP_MB").is_some() {
1266 let cand: Vec<String> = init
1267 .iter()
1268 .map(|r| {
1269 let pn = if r.frame_num as i64 > curr {
1270 r.frame_num as i64 - max
1271 } else {
1272 r.frame_num as i64
1273 };
1274 format!("(fn={} poc={} lt={} picnum={})", r.frame_num, r.poc, r.long_term, pn)
1275 })
1276 .collect();
1277 eprintln!(
1278 "MODFAIL idc={idc} val={val} curr_frame_num={curr} max={max} init={}",
1279 cand.join(" ")
1280 );
1281 }
1282 return Err(DecodeError::Truncated); };
1284 if refidx > list.len() {
1285 break;
1286 }
1287 list.insert(refidx, found);
1288 if let Some(dup) = list.iter().enumerate().skip(refidx + 1).find(|(_, r)| matches(r)).map(|(i, _)| i) {
1289 list.remove(dup);
1290 }
1291 refidx += 1;
1292 if refidx >= num_active {
1293 break;
1294 }
1295 }
1296 list.truncate(num_active.max(1));
1297 Ok(list)
1298}
1299
1300#[cfg(test)]
1301mod tests {
1302 use super::*;
1303
1304 fn ref_at(poc: i32, fnum: u32) -> Ref {
1305 std::sync::Arc::new(RefFrame {
1306 py: vec![],
1307 pu: vec![],
1308 pv: vec![],
1309 cw: 0,
1310 ch: 0,
1311 frame_num: fnum,
1312 poc,
1313 mv: Vec::new(),
1314 ref_idx: Vec::new(),
1315 mv1: Vec::new(),
1316 ref_idx1: Vec::new(),
1317 ref_poc: Vec::new(),
1318 w4: 0,
1319 long_term: false,
1320 long_term_idx: 0,
1321 })
1322 }
1323
1324 #[test]
1325 fn b_ref_lists_ordered_by_poc() {
1326 let dpb = vec![ref_at(8, 4), ref_at(6, 3), ref_at(2, 1), ref_at(0, 0)];
1328 let (l0, l1) = build_ref_list_b(&dpb, 4, 5, 16, 4, 4, &[], &[]).unwrap();
1329 assert_eq!(l0.iter().map(|r| r.poc).collect::<Vec<_>>(), vec![2, 0, 6, 8]);
1331 assert_eq!(l1.iter().map(|r| r.poc).collect::<Vec<_>>(), vec![6, 8, 2, 0]);
1333 }
1334
1335 #[test]
1336 fn b_ref_list1_swap_when_equal() {
1337 let dpb = vec![ref_at(4, 2), ref_at(2, 1), ref_at(0, 0)];
1340 let (l0, l1) = build_ref_list_b(&dpb, 6, 3, 16, 3, 3, &[], &[]).unwrap();
1341 assert_eq!(l0.iter().map(|r| r.poc).collect::<Vec<_>>(), vec![4, 2, 0]);
1342 assert_eq!(l1.iter().map(|r| r.poc).collect::<Vec<_>>(), vec![2, 4, 0]);
1343 }
1344
1345 #[test]
1346 fn frame_num_gaps_insert_placeholders() {
1347 let mut d = Decoder::new();
1348 d.prev_ref_frame_num = 2;
1349 d.insert_frame_num_gaps(5, 16, 8, 16, 16);
1351 let fns: Vec<u32> = d.refs.iter().map(|r| r.frame_num).collect();
1352 assert_eq!(fns, vec![4, 3], "most-recent placeholder at the front");
1353 assert_eq!(d.prev_ref_frame_num, 4);
1354 assert!(d.refs.iter().all(|r| r.py.iter().all(|&p| p == 128)), "grey fill");
1355 }
1356
1357 #[test]
1358 fn frame_num_gaps_wrap_and_noop() {
1359 let mut d = Decoder::new();
1361 d.prev_ref_frame_num = 14;
1362 d.insert_frame_num_gaps(1, 16, 8, 16, 16);
1363 assert_eq!(d.refs.iter().map(|r| r.frame_num).collect::<Vec<_>>(), vec![0, 15]);
1364 let mut d = Decoder::new();
1366 d.prev_ref_frame_num = 3;
1367 d.insert_frame_num_gaps(4, 16, 8, 16, 16);
1368 assert!(d.refs.is_empty());
1369 }
1370
1371 #[test]
1372 fn missing_param_sets_errors() {
1373 let mut d = Decoder::new();
1374 let nal = rusty_h264_common::NalUnit::new(3, NalUnitType::IdrSlice, vec![0x88, 0x80]);
1377 let err = d.decode(&nal.to_annex_b()).unwrap_err();
1378 assert_eq!(err, DecodeError::MissingParameterSet);
1379 }
1380}