1mod cabac;
26mod mb16;
27mod params;
28
29pub use params::{Pps, Sps};
30pub use mb16::{MvField, MV_DUMP};
31
32#[doc(hidden)]
35fn abl_deblock() -> bool {
38 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
39 *ON.get_or_init(|| std::env::var("RFF_ABL_DEBLOCK").map_or(false, |v| v != "0"))
40}
41
42pub mod cabac_test {
43 pub use crate::cabac::Cabac;
44 pub use crate::mb16::b_inter_shape;
45 pub use crate::mb16::parse_mb_type_b;
46}
47
48use mb16::{FrameDecoder, WeightTable};
49use rusty_h264_common::bit_reader::OutOfData;
50use rusty_h264_common::nal::{emulation_unprevent, split_annex_b};
51use rusty_h264_common::{BitReader, NalUnitType, YuvFrame};
52
53#[derive(Debug, Clone, PartialEq, Eq)]
55pub enum DecodeError {
56 Truncated,
58 MissingParameterSet,
60 Unsupported(&'static str),
62}
63
64impl From<OutOfData> for DecodeError {
65 fn from(_: OutOfData) -> Self {
66 DecodeError::Truncated
67 }
68}
69
70impl core::fmt::Display for DecodeError {
71 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
72 match self {
73 DecodeError::Truncated => f.write_str("bitstream truncated"),
74 DecodeError::MissingParameterSet => f.write_str("slice before SPS/PPS"),
75 DecodeError::Unsupported(s) => write!(f, "unsupported coding tool: {s}"),
76 }
77 }
78}
79
80impl std::error::Error for DecodeError {}
81
82pub(crate) type Ref = std::sync::Arc<RefFrame>;
91
92#[derive(Debug, Clone, Default)]
93#[allow(dead_code)]
94pub(crate) struct RefFrame {
95 pub py: Vec<u8>,
101 pub pu: Vec<u8>,
102 pub pv: Vec<u8>,
103 pub cw: usize,
104 pub ch: usize,
105 pub frame_num: u32,
107 pub poc: i32,
109 pub mv: Vec<(i32, i32)>,
113 pub ref_idx: Vec<i32>,
114 pub mv1: Vec<(i32, i32)>,
121 pub ref_idx1: Vec<i32>,
122 pub ref_poc: Vec<i32>,
126 pub w4: usize,
127 pub long_term: bool,
131 pub long_term_idx: u32,
132}
133
134pub(crate) const LPAD: usize = 16;
138pub(crate) const CPAD: usize = 8;
139
140impl RefFrame {
141 #[inline]
142 pub fn lstride(&self) -> usize {
143 self.cw + 2 * LPAD
144 }
145 #[inline]
146 pub fn cstride(&self) -> usize {
147 self.cw / 2 + 2 * CPAD
148 }
149}
150
151#[derive(Clone, Copy)]
153enum Mmco {
154 Unref(u32),
156 UnrefLong(u32),
158 AssignLong(u32, u32),
160 MaxLong(u32),
162 Reset,
164 CurrentLong(u32),
166}
167
168struct PendingPic {
172 fd: mb16::FrameDecoder,
173 frame_num: u32,
174 poc: i32,
175 next_mb: usize,
176 total_mb: usize,
177 slice_count: u16,
178 deblock: bool,
179 filter_offset_a: i32,
180 filter_offset_b: i32,
181 crop_r: usize,
182 crop_b: usize,
183 max_refs: usize,
184 log2_max_frame_num: u32,
185 is_reference: bool,
188 idr_long_term: bool,
189 mmco_ops: Vec<Mmco>,
190}
191
192#[derive(Default)]
195pub struct Decoder {
196 sps: std::collections::HashMap<u32, Sps>,
199 pps: std::collections::HashMap<u32, Pps>,
200 refs: Vec<Ref>,
202 cur: Option<PendingPic>,
204 poc: PocState,
208 last_poc: i32,
210 prev_ref_frame_num: u32,
213}
214
215#[derive(Default)]
217struct PocState {
218 prev_msb: i32,
219 prev_lsb: i32,
220 prev_frame_num: u32,
221 prev_frame_num_offset: i64,
222}
223
224impl Decoder {
225 pub fn new() -> Self {
227 Self::default()
228 }
229
230 pub fn decode(&mut self, annex_b: &[u8]) -> Result<Option<YuvFrame>, DecodeError> {
233 let _g = rusty_h264_common::prof::scope(rusty_h264_common::prof::Stage::Total);
234 let mut frame = None;
235 for nal in split_annex_b(annex_b) {
236 if nal.is_empty() {
237 continue;
238 }
239 let nal_type = NalUnitType::from_id(nal[0]);
240 let rbsp = emulation_unprevent(&nal[1..]);
241 match nal_type {
242 NalUnitType::Sps => {
243 let s = Sps::parse(&rbsp)?;
244 self.sps.insert(s.seq_parameter_set_id, s);
245 }
246 NalUnitType::Pps => {
247 let p = Pps::parse(&rbsp)?;
248 self.pps.insert(p.pic_parameter_set_id, p);
249 }
250 NalUnitType::IdrSlice | NalUnitType::NonIdrSlice => {
251 let nal_ref_idc = (nal[0] >> 5) & 3;
252 let is_idr = nal_type == NalUnitType::IdrSlice;
253 if let Some(f) = self.decode_slice(&rbsp, is_idr, nal_ref_idc)? {
254 frame = Some(f);
255 }
256 }
257 _ => {} }
259 }
260 Ok(frame)
261 }
262
263 pub fn decode_stream(&mut self, annex_b: &[u8]) -> Result<Vec<YuvFrame>, DecodeError> {
271 let mut out = Vec::new();
272 let mut gop: Vec<(i32, YuvFrame)> = Vec::new();
273 for au in split_access_units(annex_b) {
274 if au_is_idr(au) {
275 flush_gop(&mut gop, &mut out); }
277 if let Some(frame) = self.decode(au)? {
278 gop.push((self.last_poc, frame));
279 }
280 }
281 flush_gop(&mut gop, &mut out);
282 Ok(out)
283 }
284
285 fn decode_slice(
286 &mut self,
287 rbsp: &[u8],
288 is_idr: bool,
289 nal_ref_idc: u8,
290 ) -> Result<Option<YuvFrame>, DecodeError> {
291 let mut r = BitReader::new(rbsp);
292 let first_mb_in_slice = r.read_ue()? as usize;
294 let slice_type = r.read_ue()?;
295 let is_p = matches!(slice_type, 0 | 5);
296 let is_b = matches!(slice_type, 1 | 6);
297 let is_i = matches!(slice_type, 2 | 7);
298 if !is_p && !is_b && !is_i {
299 return Err(DecodeError::Unsupported("SP/SI slices"));
300 }
301 let pic_parameter_set_id = r.read_ue()?;
303 let pps = self.pps.get(&pic_parameter_set_id).cloned().ok_or(DecodeError::MissingParameterSet)?;
304 let sps = self.sps.get(&pps.seq_parameter_set_id).cloned().ok_or(DecodeError::MissingParameterSet)?;
305 let sps = &sps;
306 let pps = &pps;
307 let cabac = pps.entropy_coding_mode_flag;
311 let frame_num = r.read_bits(sps.log2_max_frame_num)?;
312 if is_idr {
313 let _idr_pic_id = r.read_ue()?;
314 }
315 let mut poc_lsb = 0u32;
318 let mut delta_poc_bottom = 0i32;
319 if sps.pic_order_cnt_type == 0 {
320 poc_lsb = r.read_bits(sps.log2_max_pic_order_cnt_lsb)?;
321 if pps.bottom_field_pic_order_present {
322 delta_poc_bottom = r.read_se()?;
323 }
324 } else if sps.pic_order_cnt_type == 1 && !sps.delta_pic_order_always_zero {
325 let _delta_pic_order_cnt_0 = r.read_se()?;
326 if pps.bottom_field_pic_order_present {
327 let _delta_pic_order_cnt_1 = r.read_se()?;
328 }
329 }
330 let pic_poc = if first_mb_in_slice == 0 {
333 self.compute_poc(sps, is_idr, nal_ref_idc, frame_num, poc_lsb, delta_poc_bottom)
334 } else {
335 self.cur.as_ref().map_or(0, |p| p.poc)
336 };
337 if pps.redundant_pic_cnt_present_flag {
342 let redundant_pic_cnt = r.read_ue()?;
343 if redundant_pic_cnt != 0 {
344 return Ok(None);
345 }
346 }
347 if std::env::var_os("RH264_DUMP_MB").is_some() {
348 eprintln!(
349 "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}"
350 );
351 }
352 let direct_spatial = if is_b { r.read_bit()? } else { true };
354 let mut num_ref_idx_l0 = pps.num_ref_idx_l0_default as usize;
355 let mut num_ref_idx_l1 = pps.num_ref_idx_l1_default as usize;
356 let mut reorder_l0: Vec<(u32, u32)> = Vec::new();
357 let mut reorder_l1: Vec<(u32, u32)> = Vec::new();
358 if is_p || is_b {
359 if r.read_bit()? {
361 num_ref_idx_l0 = (r.read_ue()? + 1) as usize;
362 if is_b {
363 num_ref_idx_l1 = (r.read_ue()? + 1) as usize;
364 }
365 }
366 if r.read_bit()? {
368 parse_ref_pic_list_modification(&mut r, &mut reorder_l0)?;
369 }
370 if is_b && r.read_bit()? {
371 parse_ref_pic_list_modification(&mut r, &mut reorder_l1)?;
373 }
374 }
375 let weights = if is_p && pps.weighted_pred {
380 Some(parse_pred_weight_table(&mut r, num_ref_idx_l0, 0, false)?)
381 } else if is_b && pps.weighted_bipred_idc == 1 {
382 return Err(DecodeError::Unsupported("explicit B weighted prediction"));
383 } else {
384 None
385 };
386 let mut idr_long_term = false;
390 let mut mmco_ops: Vec<Mmco> = Vec::new();
391 if nal_ref_idc == 0 {
392 } else if is_idr {
394 let _no_output_of_prior_pics = r.read_bit()?;
395 idr_long_term = r.read_bit()?; } else if r.read_bit()? {
397 loop {
399 let op = r.read_ue()?;
400 match op {
401 0 => break,
402 1 => mmco_ops.push(Mmco::Unref(r.read_ue()?)),
403 2 => mmco_ops.push(Mmco::UnrefLong(r.read_ue()?)),
404 3 => {
405 let diff = r.read_ue()?;
406 let idx = r.read_ue()?;
407 mmco_ops.push(Mmco::AssignLong(diff, idx));
408 }
409 4 => mmco_ops.push(Mmco::MaxLong(r.read_ue()?)),
410 5 => mmco_ops.push(Mmco::Reset),
411 6 => mmco_ops.push(Mmco::CurrentLong(r.read_ue()?)),
412 _ => return Err(DecodeError::Unsupported("invalid MMCO")),
413 }
414 if mmco_ops.len() > 128 {
415 return Err(DecodeError::Truncated);
416 }
417 }
418 }
419 let cabac_init_idc = if cabac && !is_i {
423 let v = r.read_ue()?;
424 if v > 2 {
425 return Err(DecodeError::Unsupported("invalid cabac_init_idc"));
426 }
427 v
428 } else {
429 0
430 };
431 let slice_qp_delta = r.read_se()?;
432 let mut deblock = true;
437 let (mut filter_offset_a, mut filter_offset_b) = (0i32, 0i32);
438 if pps.deblocking_filter_control_present_flag {
439 let disable_deblocking_filter_idc = r.read_ue()?;
440 deblock = disable_deblocking_filter_idc != 1;
443 if disable_deblocking_filter_idc != 1 {
444 filter_offset_a = r.read_se()? * 2;
446 filter_offset_b = r.read_se()? * 2;
447 }
448 }
449 let slice_qp = (pps.pic_init_qp + slice_qp_delta).clamp(0, 51) as u8;
450
451 if first_mb_in_slice == 0 && !is_idr && sps.gaps_in_frame_num_allowed {
454 self.insert_frame_num_gaps(
455 frame_num,
456 1u32 << sps.log2_max_frame_num,
457 sps.max_num_ref_frames.max(1) as usize,
458 sps.pic_width_in_mbs * 16,
459 sps.pic_height_in_mbs * 16,
460 );
461 }
462
463 let max_fn = 1u32 << sps.log2_max_frame_num;
466 let (ref_list0, ref_list1) = if is_b {
467 build_ref_list_b(
468 &self.refs, pic_poc, frame_num, max_fn,
469 num_ref_idx_l0, num_ref_idx_l1, &reorder_l0, &reorder_l1,
470 )?
471 } else if is_p {
472 (build_ref_list_p(&self.refs, frame_num, max_fn, num_ref_idx_l0, &reorder_l0)?, Vec::new())
473 } else {
474 (Vec::new(), Vec::new())
475 };
476 if first_mb_in_slice == 0 {
480 if is_idr {
481 self.refs.clear();
482 }
483 let _g_setup = rusty_h264_common::prof::scope(rusty_h264_common::prof::Stage::DecSetup);
494 let mut fd = FrameDecoder::new(
495 sps.pic_width_in_mbs,
496 sps.pic_height_in_mbs,
497 slice_qp,
498 pps.chroma_qp_index_offset,
499 ref_list0,
500 num_ref_idx_l0,
501 pps.constrained_intra_pred_flag,
502 pps.transform_8x8_mode_flag,
503 sps.profile_idc != 66, );
505 if is_b {
506 fd.set_b_context(
507 ref_list1,
508 num_ref_idx_l1,
509 direct_spatial,
510 pic_poc,
511 pps.weighted_bipred_idc,
512 sps.direct_8x8_inference,
513 );
514 }
515 if sps.has_scaling || pps.pic_scaling_matrix_present {
516 let (s4, s8) = resolve_scaling(sps, pps);
517 fd.set_scaling(s4, s8);
518 }
519 if let Some(w) = weights {
520 fd.set_weights(w);
521 }
522 if let Some(prev) = self.cur.take() {
530 if prev.next_mb < prev.total_mb {
531 return Err(DecodeError::Truncated);
532 }
533 }
534 self.cur = Some(PendingPic {
535 fd,
536 frame_num,
537 poc: pic_poc,
538 next_mb: 0,
539 total_mb: sps.pic_width_in_mbs * sps.pic_height_in_mbs,
540 slice_count: 0,
541 deblock,
542 filter_offset_a,
543 filter_offset_b,
544 crop_r: sps.frame_crop_right as usize,
545 crop_b: sps.frame_crop_bottom as usize,
546 max_refs: sps.max_num_ref_frames.max(1) as usize,
547 log2_max_frame_num: sps.log2_max_frame_num,
548 is_reference: nal_ref_idc != 0,
549 idr_long_term,
550 mmco_ops,
551 });
552 } else {
553 let Some(pic) = self.cur.as_mut() else {
555 return Err(DecodeError::Unsupported("slice continues a missing picture"));
556 };
557 pic.fd.begin_slice(slice_qp, ref_list0, num_ref_idx_l0);
558 if is_b {
559 pic.fd.set_b_context(
560 ref_list1,
561 num_ref_idx_l1,
562 direct_spatial,
563 pic.poc,
564 pps.weighted_bipred_idc,
565 sps.direct_8x8_inference,
566 );
567 }
568 if sps.has_scaling || pps.pic_scaling_matrix_present {
569 let (s4, s8) = resolve_scaling(sps, pps);
570 pic.fd.set_scaling(s4, s8);
571 }
572 if let Some(w) = weights {
573 pic.fd.set_weights(w);
574 }
575 pic.deblock = deblock;
577 pic.filter_offset_a = filter_offset_a;
578 pic.filter_offset_b = filter_offset_b;
579 pic.idr_long_term |= idr_long_term;
580 pic.mmco_ops.extend(mmco_ops);
581 }
582
583 let pic = self.cur.as_mut().expect("pending picture set above");
584 let first = first_mb_in_slice.min(pic.total_mb);
585 let next = if cabac {
586 r.align_to_byte().map_err(|_| DecodeError::Truncated)?;
588 let (data, start) = (r.data(), r.bit_pos() / 8);
589 pic.fd
590 .decode_slice_data_cabac(data, start, slice_qp, cabac_init_idc, is_i, is_p, first)
591 } else {
592 pic.fd.decode_slice_data(&mut r, is_p, first)
593 }
594 .map_err(|e| match e {
595 mb16::MbError::Truncated => DecodeError::Truncated,
596 mb16::MbError::Unsupported(s) => DecodeError::Unsupported(s),
597 })?;
598 pic.next_mb = next;
599 pic.slice_count += 1;
600 if std::env::var_os("RH264_DUMP_MB").is_some() {
601 eprintln!(
602 " slice decoded {}/{} MBs{}",
603 next,
604 pic.total_mb,
605 if next < pic.total_mb { " <-- INCOMPLETE" } else { "" }
606 );
607 }
608
609 if pic.next_mb < pic.total_mb {
610 return Ok(None); }
612
613 let pic = self.cur.take().expect("pending picture");
615 let PendingPic {
616 mut fd,
617 frame_num,
618 poc,
619 deblock,
620 filter_offset_a,
621 filter_offset_b,
622 crop_r,
623 crop_b,
624 max_refs,
625 log2_max_frame_num,
626 is_reference,
627 idr_long_term,
628 mmco_ops,
629 ..
630 } = pic;
631 self.last_poc = poc;
632 if deblock && !abl_deblock() {
640 fd.deblock(filter_offset_a, filter_offset_b);
641 }
642 let reference = if is_reference {
645 let _dg = rusty_h264_common::prof::scope(rusty_h264_common::prof::Stage::DpbClone);
646 Some(fd.as_reference())
647 } else {
648 None
650 };
651 let _fg = rusty_h264_common::prof::scope(rusty_h264_common::prof::Stage::Finalize);
652 if let Some(mut reference) = reference {
653 reference.frame_num = frame_num;
654 reference.poc = poc;
655 if std::env::var_os("RH264_DUMP_MB").is_some() {
656 eprintln!("DPB-ADD fn={frame_num} poc={poc}");
657 }
658 if idr_long_term {
659 reference.long_term = true;
660 reference.long_term_idx = 0;
661 }
662 self.prev_ref_frame_num =
664 self.apply_ref_marking(reference, &mmco_ops, frame_num, log2_max_frame_num, max_refs);
665 }
666 Ok(Some(fd.into_frame(crop_r, crop_b)))
667 }
668
669 fn insert_frame_num_gaps(&mut self, frame_num: u32, max_fn: u32, max_refs: usize, w: usize, h: usize) {
675 if max_fn == 0 {
676 return;
677 }
678 let start = (self.prev_ref_frame_num + 1) % max_fn;
679 let gap = (frame_num + max_fn - start) % max_fn;
680 if gap == 0 {
681 return;
682 }
683 let cap = max_refs.max(1);
689 let n = (gap as usize).min(cap);
690 let (cw, ch) = (w, h);
691 let mut expected = (frame_num + max_fn - n as u32) % max_fn;
692 for _ in 0..n {
693 self.refs.insert(
694 0,
695 std::sync::Arc::new(RefFrame {
696 py: vec![128; (cw + 2 * LPAD) * (ch + 2 * LPAD)],
698 pu: vec![128; (cw / 2 + 2 * CPAD) * (ch / 2 + 2 * CPAD)],
699 pv: vec![128; (cw / 2 + 2 * CPAD) * (ch / 2 + 2 * CPAD)],
700 cw,
701 ch,
702 frame_num: expected,
703 poc: 0,
704 mv: Vec::new(),
705 ref_idx: Vec::new(),
706 mv1: Vec::new(),
707 ref_idx1: Vec::new(),
708 ref_poc: Vec::new(),
709 w4: 0,
710 long_term: false,
711 long_term_idx: 0,
712 }),
713 );
714 self.refs.truncate(cap);
715 expected = (expected + 1) % max_fn;
716 }
717 self.prev_ref_frame_num = (frame_num + max_fn - 1) % max_fn;
718 }
719
720 pub fn last_poc(&self) -> i32 {
724 self.last_poc
725 }
726
727 fn compute_poc(
731 &mut self,
732 sps: &Sps,
733 is_idr: bool,
734 nal_ref_idc: u8,
735 frame_num: u32,
736 poc_lsb: u32,
737 delta_bottom: i32,
738 ) -> i32 {
739 match sps.pic_order_cnt_type {
740 0 => {
741 let max_lsb = 1i32 << sps.log2_max_pic_order_cnt_lsb;
742 let (prev_msb, prev_lsb) =
743 if is_idr { (0, 0) } else { (self.poc.prev_msb, self.poc.prev_lsb) };
744 let lsb = poc_lsb as i32;
745 let msb = if lsb < prev_lsb && prev_lsb - lsb >= max_lsb / 2 {
746 prev_msb + max_lsb
747 } else if lsb > prev_lsb && lsb - prev_lsb > max_lsb / 2 {
748 prev_msb - max_lsb
749 } else {
750 prev_msb
751 };
752 let top = msb + lsb;
753 let poc = top.min(top + delta_bottom);
754 if nal_ref_idc != 0 {
755 self.poc.prev_msb = msb;
756 self.poc.prev_lsb = lsb;
757 }
758 poc
759 }
760 2 => {
761 let max_fn = 1i64 << sps.log2_max_frame_num;
762 let offset = if is_idr {
763 0
764 } else if self.poc.prev_frame_num > frame_num {
765 self.poc.prev_frame_num_offset + max_fn
766 } else {
767 self.poc.prev_frame_num_offset
768 };
769 let poc = if is_idr {
770 0
771 } else {
772 2 * (offset + frame_num as i64) - i64::from(nal_ref_idc == 0)
773 };
774 self.poc.prev_frame_num_offset = offset;
775 self.poc.prev_frame_num = frame_num;
776 poc as i32
777 }
778 _ => {
779 self.poc.prev_frame_num = frame_num;
780 frame_num as i32 * 2
781 }
782 }
783 }
784
785 fn apply_ref_marking(
796 &mut self,
797 mut reference: RefFrame,
798 ops: &[Mmco],
799 frame_num: u32,
800 log2_max_frame_num: u32,
801 max_refs: usize,
802 ) -> u32 {
803 let max = 1i64 << log2_max_frame_num;
804 let curr = frame_num as i64;
805 let pic_num = |rf: &RefFrame| -> i64 {
806 if (rf.frame_num as i64) > curr {
807 rf.frame_num as i64 - max
808 } else {
809 rf.frame_num as i64
810 }
811 };
812
813 if ops.is_empty() {
814 let out_fn = reference.frame_num;
818 self.refs.insert(0, std::sync::Arc::new(reference));
819 while self.refs.len() > max_refs {
820 match self.refs.iter().rposition(|r| !r.long_term) {
821 Some(pos) => {
822 self.refs.remove(pos);
823 }
824 None => break,
825 }
826 }
827 return out_fn;
828 }
829
830 for &op in ops {
832 match op {
833 Mmco::Unref(diff) => {
834 let target = curr - (diff as i64 + 1);
835 self.refs.retain(|r| r.long_term || pic_num(r) != target);
836 }
837 Mmco::UnrefLong(ltpn) => {
838 self.refs.retain(|r| !(r.long_term && r.long_term_idx == ltpn));
839 }
840 Mmco::AssignLong(diff, idx) => {
841 let target = curr - (diff as i64 + 1);
842 self.refs.retain(|r| !(r.long_term && r.long_term_idx == idx));
843 for r in self.refs.iter_mut() {
844 if !r.long_term && pic_num(r) == target {
845 let r = std::sync::Arc::make_mut(r);
847 r.long_term = true;
848 r.long_term_idx = idx;
849 }
850 }
851 }
852 Mmco::MaxLong(max_plus1) => {
853 self.refs.retain(|r| !(r.long_term && r.long_term_idx + 1 > max_plus1));
854 }
855 Mmco::Reset => {
856 self.refs.clear();
857 reference.frame_num = 0;
858 }
859 Mmco::CurrentLong(idx) => {
860 self.refs.retain(|r| !(r.long_term && r.long_term_idx == idx));
861 reference.long_term = true;
862 reference.long_term_idx = idx;
863 }
864 }
865 }
866 let out_fn = reference.frame_num;
867 self.refs.insert(0, std::sync::Arc::new(reference));
868 let cap = max_refs.max(16);
870 if self.refs.len() > cap {
871 self.refs.truncate(cap);
872 }
873 out_fn
874 }
875}
876
877fn flush_gop(gop: &mut Vec<(i32, YuvFrame)>, out: &mut Vec<YuvFrame>) {
879 gop.sort_by_key(|(poc, _)| *poc);
880 out.extend(gop.drain(..).map(|(_, f)| f));
881}
882
883fn au_is_idr(au: &[u8]) -> bool {
885 split_annex_b(au)
886 .iter()
887 .any(|n| !n.is_empty() && NalUnitType::from_id(n[0]) == NalUnitType::IdrSlice)
888}
889
890pub fn split_access_units(stream: &[u8]) -> Vec<&[u8]> {
898 let mut codes: Vec<(usize, bool)> = Vec::new();
900 let mut i = 0;
901 while i + 3 <= stream.len() {
902 if stream[i] == 0 && stream[i + 1] == 0 && stream[i + 2] == 1 {
903 let nal_type = NalUnitType::from_id(stream.get(i + 3).copied().unwrap_or(0));
904 let is_vcl = matches!(nal_type, NalUnitType::IdrSlice | NalUnitType::NonIdrSlice);
905 let sc = if i > 0 && stream[i - 1] == 0 { i - 1 } else { i };
907 codes.push((sc, is_vcl));
908 i += 3;
909 } else {
910 i += 1;
911 }
912 }
913 if codes.is_empty() {
914 return vec![stream];
915 }
916 let mut aus = Vec::new();
917 let mut start = codes[0].0;
918 for k in 0..codes.len() {
919 if codes[k].1 {
920 let end = codes.get(k + 1).map_or(stream.len(), |c| c.0);
921 aus.push(&stream[start..end]);
922 start = end;
923 }
924 }
925 aus
926}
927
928fn parse_pred_weight_table(
931 r: &mut BitReader,
932 num_l0: usize,
933 num_l1: usize,
934 is_b: bool,
935) -> Result<WeightTable, DecodeError> {
936 let luma_log2_denom = r.read_ue()? as i32;
937 let chroma_log2_denom = r.read_ue()? as i32;
938 if !(0..=7).contains(&luma_log2_denom) || !(0..=7).contains(&chroma_log2_denom) {
942 return Err(DecodeError::Unsupported("invalid weight denom"));
943 }
944 let mut wt = WeightTable {
945 luma_log2_denom,
946 chroma_log2_denom,
947 ..Default::default()
948 };
949 let lists: &[(usize, usize)] = if is_b {
950 &[(0, num_l0), (1, num_l1)]
951 } else {
952 &[(0, num_l0)]
953 };
954 for &(list, n) in lists {
955 let mut luma = Vec::with_capacity(n);
956 let mut chroma = Vec::with_capacity(n);
957 for _ in 0..n {
958 let (mut lw, mut lo) = (1 << luma_log2_denom, 0);
959 if r.read_bit()? {
960 lw = r.read_se()?;
961 lo = r.read_se()?;
962 }
963 luma.push((lw, lo));
964 let mut ch = [(1 << chroma_log2_denom, 0); 2];
965 if r.read_bit()? {
966 for slot in ch.iter_mut() {
967 *slot = (r.read_se()?, r.read_se()?);
968 }
969 }
970 chroma.push(ch);
971 }
972 wt.luma[list] = luma;
973 wt.chroma[list] = chroma;
974 }
975 Ok(wt)
976}
977
978fn resolve_scaling(sps: &Sps, pps: &Pps) -> ([[i32; 16]; 6], [[i32; 64]; 2]) {
983 use crate::params::{
984 DEFAULT_4X4_INTER, DEFAULT_4X4_INTRA, DEFAULT_8X8_INTER, DEFAULT_8X8_INTRA,
985 };
986 const ZZ4: [usize; 16] = [0, 1, 4, 8, 5, 2, 3, 6, 9, 12, 13, 10, 7, 11, 14, 15];
987 const ZZ8: [usize; 64] = [
989 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,
990 20, 13, 6, 7, 14, 21, 28, 35, 42, 49, 56, 57, 50, 43, 36, 29, 22, 15, 23, 30, 37, 44, 51,
991 58, 59, 52, 45, 38, 31, 39, 46, 53, 60, 61, 54, 47, 55, 62, 63,
992 ];
993 let mut z4 = [[16u8; 16]; 6];
996 for i in 0..6 {
997 z4[i] = if pps.pic_scaling_matrix_present {
998 if pps.scaling_present_4x4[i] {
999 pps.scaling_4x4[i]
1000 } else {
1001 match i {
1002 0 if sps.has_scaling => sps.scaling_4x4[0],
1003 0 => DEFAULT_4X4_INTRA,
1004 3 if sps.has_scaling => sps.scaling_4x4[3],
1005 3 => DEFAULT_4X4_INTER,
1006 _ => z4[i - 1],
1007 }
1008 }
1009 } else {
1010 sps.scaling_4x4[i]
1011 };
1012 }
1013 let mut z8 = [[16u8; 64]; 2];
1014 for (i, list) in z8.iter_mut().enumerate() {
1015 *list = if pps.pic_scaling_matrix_present {
1016 if pps.scaling_present_8x8[i] {
1017 pps.scaling_8x8[i]
1018 } else if sps.has_scaling {
1019 sps.scaling_8x8[i]
1020 } else if i == 0 {
1021 DEFAULT_8X8_INTRA
1022 } else {
1023 DEFAULT_8X8_INTER
1024 }
1025 } else {
1026 sps.scaling_8x8[i]
1027 };
1028 }
1029 let mut out4 = [[16i32; 16]; 6];
1030 for (li, list) in out4.iter_mut().enumerate() {
1031 for k in 0..16 {
1032 list[ZZ4[k]] = z4[li][k] as i32;
1033 }
1034 }
1035 let mut out8 = [[16i32; 64]; 2];
1036 for (li, list) in out8.iter_mut().enumerate() {
1037 for k in 0..64 {
1038 list[ZZ8[k]] = z8[li][k] as i32;
1039 }
1040 }
1041 (out4, out8)
1042}
1043
1044fn parse_ref_pic_list_modification(
1047 r: &mut BitReader,
1048 out: &mut Vec<(u32, u32)>,
1049) -> Result<(), DecodeError> {
1050 loop {
1051 let idc = r.read_ue()?;
1052 if idc == 3 {
1053 break;
1054 }
1055 if idc > 3 {
1056 return Err(DecodeError::Unsupported("invalid ref_pic_list_modification"));
1057 }
1058 let val = r.read_ue()?; out.push((idc, val));
1060 if out.len() > 64 {
1061 return Err(DecodeError::Truncated); }
1063 }
1064 Ok(())
1065}
1066
1067fn build_ref_list_p(
1071 dpb: &[Ref],
1072 curr_frame_num: u32,
1073 max_frame_num: u32,
1074 num_active: usize,
1075 mods: &[(u32, u32)],
1076) -> Result<Vec<Ref>, DecodeError> {
1077 let curr = curr_frame_num as i64;
1078 let max = max_frame_num as i64;
1079 let pic_num = |fnum: u32| -> i64 {
1080 let f = fnum as i64;
1081 if f > curr { f - max } else { f }
1082 };
1083 let mut init: Vec<Ref> = dpb.iter().filter(|r| !r.long_term).cloned().collect();
1084 init.sort_by_key(|rf| core::cmp::Reverse(pic_num(rf.frame_num)));
1085 let mut long: Vec<Ref> = dpb.iter().filter(|r| r.long_term).cloned().collect();
1086 long.sort_by_key(|rf| rf.long_term_idx);
1087 init.extend(long);
1088 apply_list_modification(init, curr_frame_num, max_frame_num, num_active, mods)
1089}
1090
1091#[allow(clippy::too_many_arguments)]
1096fn build_ref_list_b(
1097 dpb: &[Ref],
1098 curr_poc: i32,
1099 curr_frame_num: u32,
1100 max_frame_num: u32,
1101 num0: usize,
1102 num1: usize,
1103 mods0: &[(u32, u32)],
1104 mods1: &[(u32, u32)],
1105) -> Result<(Vec<Ref>, Vec<Ref>), DecodeError> {
1106 let mut less: Vec<Ref> =
1107 dpb.iter().filter(|r| !r.long_term && r.poc < curr_poc).cloned().collect();
1108 let mut greater: Vec<Ref> =
1109 dpb.iter().filter(|r| !r.long_term && r.poc > curr_poc).cloned().collect();
1110 let mut long: Vec<Ref> = dpb.iter().filter(|r| r.long_term).cloned().collect();
1111 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);
1114
1115 let mut init0 = less.clone();
1116 init0.extend(greater.clone());
1117 init0.extend(long.clone());
1118 let mut init1 = greater;
1119 init1.extend(less);
1120 init1.extend(long);
1121
1122 let eq_len = num0.min(num1).min(init0.len()).min(init1.len());
1125 if num1 > 1
1126 && init1.len() > 1
1127 && (0..eq_len).all(|i| same_picture(&init0[i], &init1[i]))
1128 && eq_len == num1.min(init1.len())
1129 && eq_len == num0.min(init0.len())
1130 {
1131 init1.swap(0, 1);
1132 }
1133
1134 let list0 = apply_list_modification(init0, curr_frame_num, max_frame_num, num0, mods0)?;
1135 let list1 = apply_list_modification(init1, curr_frame_num, max_frame_num, num1, mods1)?;
1136 Ok((list0, list1))
1137}
1138
1139fn same_picture(a: &RefFrame, b: &RefFrame) -> bool {
1141 a.long_term == b.long_term
1142 && if a.long_term { a.long_term_idx == b.long_term_idx } else { a.poc == b.poc }
1143}
1144
1145fn apply_list_modification(
1150 init: Vec<Ref>,
1151 curr_frame_num: u32,
1152 max_frame_num: u32,
1153 num_active: usize,
1154 mods: &[(u32, u32)],
1155) -> Result<Vec<Ref>, DecodeError> {
1156 if mods.is_empty() {
1157 let mut init = init;
1158 init.truncate(num_active.max(1));
1159 return Ok(init);
1160 }
1161 let curr = curr_frame_num as i64;
1162 let max = max_frame_num as i64;
1163 let mut list = init.clone();
1164 let mut pic_num_pred = curr;
1165 let mut refidx = 0usize;
1166 for &(idc, val) in mods {
1167 let matches: Box<dyn Fn(&RefFrame) -> bool> = if idc == 2 {
1168 Box::new(move |r: &RefFrame| r.long_term && r.long_term_idx == val)
1169 } else {
1170 let abs_diff = (val as i64) + 1;
1171 let no_wrap = if idc == 0 {
1172 let x = pic_num_pred - abs_diff;
1173 if x < 0 { x + max } else { x }
1174 } else {
1175 let x = pic_num_pred + abs_diff;
1176 if x >= max { x - max } else { x }
1177 };
1178 pic_num_pred = no_wrap;
1179 let target = if no_wrap > curr { no_wrap - max } else { no_wrap };
1180 Box::new(move |r: &RefFrame| {
1181 let pn = if r.frame_num as i64 > curr {
1182 r.frame_num as i64 - max
1183 } else {
1184 r.frame_num as i64
1185 };
1186 !r.long_term && pn == target
1187 })
1188 };
1189 let found = init.iter().find(|r| matches(r)).cloned();
1190 let Some(found) = found else {
1191 if std::env::var_os("RH264_DUMP_MB").is_some() {
1192 let cand: Vec<String> = init
1193 .iter()
1194 .map(|r| {
1195 let pn = if r.frame_num as i64 > curr {
1196 r.frame_num as i64 - max
1197 } else {
1198 r.frame_num as i64
1199 };
1200 format!("(fn={} poc={} lt={} picnum={})", r.frame_num, r.poc, r.long_term, pn)
1201 })
1202 .collect();
1203 eprintln!(
1204 "MODFAIL idc={idc} val={val} curr_frame_num={curr} max={max} init={}",
1205 cand.join(" ")
1206 );
1207 }
1208 return Err(DecodeError::Truncated); };
1210 if refidx > list.len() {
1211 break;
1212 }
1213 list.insert(refidx, found);
1214 if let Some(dup) = list.iter().enumerate().skip(refidx + 1).find(|(_, r)| matches(r)).map(|(i, _)| i) {
1215 list.remove(dup);
1216 }
1217 refidx += 1;
1218 if refidx >= num_active {
1219 break;
1220 }
1221 }
1222 list.truncate(num_active.max(1));
1223 Ok(list)
1224}
1225
1226#[cfg(test)]
1227mod tests {
1228 use super::*;
1229
1230 fn ref_at(poc: i32, fnum: u32) -> Ref {
1231 std::sync::Arc::new(RefFrame {
1232 py: vec![],
1233 pu: vec![],
1234 pv: vec![],
1235 cw: 0,
1236 ch: 0,
1237 frame_num: fnum,
1238 poc,
1239 mv: Vec::new(),
1240 ref_idx: Vec::new(),
1241 mv1: Vec::new(),
1242 ref_idx1: Vec::new(),
1243 ref_poc: Vec::new(),
1244 w4: 0,
1245 long_term: false,
1246 long_term_idx: 0,
1247 })
1248 }
1249
1250 #[test]
1251 fn b_ref_lists_ordered_by_poc() {
1252 let dpb = vec![ref_at(8, 4), ref_at(6, 3), ref_at(2, 1), ref_at(0, 0)];
1254 let (l0, l1) = build_ref_list_b(&dpb, 4, 5, 16, 4, 4, &[], &[]).unwrap();
1255 assert_eq!(l0.iter().map(|r| r.poc).collect::<Vec<_>>(), vec![2, 0, 6, 8]);
1257 assert_eq!(l1.iter().map(|r| r.poc).collect::<Vec<_>>(), vec![6, 8, 2, 0]);
1259 }
1260
1261 #[test]
1262 fn b_ref_list1_swap_when_equal() {
1263 let dpb = vec![ref_at(4, 2), ref_at(2, 1), ref_at(0, 0)];
1266 let (l0, l1) = build_ref_list_b(&dpb, 6, 3, 16, 3, 3, &[], &[]).unwrap();
1267 assert_eq!(l0.iter().map(|r| r.poc).collect::<Vec<_>>(), vec![4, 2, 0]);
1268 assert_eq!(l1.iter().map(|r| r.poc).collect::<Vec<_>>(), vec![2, 4, 0]);
1269 }
1270
1271 #[test]
1272 fn frame_num_gaps_insert_placeholders() {
1273 let mut d = Decoder::new();
1274 d.prev_ref_frame_num = 2;
1275 d.insert_frame_num_gaps(5, 16, 8, 16, 16);
1277 let fns: Vec<u32> = d.refs.iter().map(|r| r.frame_num).collect();
1278 assert_eq!(fns, vec![4, 3], "most-recent placeholder at the front");
1279 assert_eq!(d.prev_ref_frame_num, 4);
1280 assert!(d.refs.iter().all(|r| r.py.iter().all(|&p| p == 128)), "grey fill");
1281 }
1282
1283 #[test]
1284 fn frame_num_gaps_wrap_and_noop() {
1285 let mut d = Decoder::new();
1287 d.prev_ref_frame_num = 14;
1288 d.insert_frame_num_gaps(1, 16, 8, 16, 16);
1289 assert_eq!(d.refs.iter().map(|r| r.frame_num).collect::<Vec<_>>(), vec![0, 15]);
1290 let mut d = Decoder::new();
1292 d.prev_ref_frame_num = 3;
1293 d.insert_frame_num_gaps(4, 16, 8, 16, 16);
1294 assert!(d.refs.is_empty());
1295 }
1296
1297 #[test]
1298 fn missing_param_sets_errors() {
1299 let mut d = Decoder::new();
1300 let nal = rusty_h264_common::NalUnit::new(3, NalUnitType::IdrSlice, vec![0x88, 0x80]);
1303 let err = d.decode(&nal.to_annex_b()).unwrap_err();
1304 assert_eq!(err, DecodeError::MissingParameterSet);
1305 }
1306}