h264: fix sliding window frame reference dropping

[nihav.git] / nihav-itu / src / codecs / h264 / pic_ref.rs

diff --git a/nihav-itu/src/codecs/h264/pic_ref.rs b/nihav-itu/src/codecs/h264/pic_ref.rs
index d75d8893dd19249838757a14b8bbb4a3d2cf5bd3..0d69292c96d82695d15e080ae68a75fe7c7ee051 100644 (file)
--- a/nihav-itu/src/codecs/h264/pic_ref.rs
+++ b/nihav-itu/src/codecs/h264/pic_ref.rs
@@ -50,12 +50,14 @@ pub struct FrameRefs {
     pub ref_list0:  Vec<Option<PictureInfo>>,
     pub ref_list1:  Vec<Option<PictureInfo>>,
     pub long_term:  Vec<Option<PictureInfo>>,
+    pub cur_id:         u32,
 
     prev_poc_msb:       u32,
     prev_poc_lsb:       u16,
     prev_ref_poc_lsb:   u16,
     prev_frame_num:     u16,
     frame_num_offset:   u32,
+    max_frame_num:      i32,
 }
 
 impl FrameRefs {
@@ -65,15 +67,18 @@ impl FrameRefs {
             ref_list0:  Vec::with_capacity(3),
             ref_list1:  Vec::with_capacity(3),
             long_term:  Vec::new(),
+            cur_id:     0,
 
             prev_poc_msb:       0,
             prev_poc_lsb:       0,
             prev_ref_poc_lsb:   0,
             prev_frame_num:     0,
             frame_num_offset:   0,
+            max_frame_num:      0,
         }
     }
     pub fn calc_picture_num(&mut self, slice_hdr: &SliceHeader, is_idr: bool, ref_id: u8, sps: &SeqParameterSet) -> u32 {
+        self.max_frame_num = 1 << sps.log2_max_frame_num;
         match sps.pic_order_cnt_type {
             0 => {
                 if is_idr {
@@ -200,8 +205,8 @@ impl FrameRefs {
                     self.long_term.resize(arg1 as usize, None);
                 },
                 5 => {
-                    self.ref_pics.truncate(0);
-                    self.long_term.truncate(0);
+                    self.ref_pics.clear();
+                    self.long_term.clear();
                 },
                 6 => {
                     // assign an long term index to current pic - done elsewhere
@@ -212,13 +217,14 @@ impl FrameRefs {
         Ok(())
     }
     pub fn clear_refs(&mut self) {
-        self.ref_pics.truncate(0);
-        self.long_term.truncate(0);
+        self.ref_pics.clear();
+        self.long_term.clear();
     }
     #[allow(clippy::cognitive_complexity)]
     pub fn select_refs(&mut self, sps: &SeqParameterSet, slice_hdr: &SliceHeader, cur_id: u32) {
-        self.ref_list0.truncate(0);
-        self.ref_list1.truncate(0);
+        self.cur_id = cur_id;
+        self.ref_list0.clear();
+        self.ref_list1.clear();
         let pic_num_mask = if sps.log2_max_frame_num == 16 {
                 0xFFFF
             } else {
@@ -328,7 +334,20 @@ impl FrameRefs {
     }
     pub fn add_short_term(&mut self, cpic: PictureInfo, num_ref_frames: usize) {
         if !self.ref_pics.is_empty() && self.ref_pics.len() >= num_ref_frames {
-            self.ref_pics.remove(0);
+            let base_id = i32::from(cpic.id);
+            let mut min_id  = base_id;
+            let mut min_idx = 0;
+            for (i, pic) in self.ref_pics.iter().enumerate() {
+                let mut pic_id = i32::from(pic.id);
+                if pic_id > base_id {
+                    pic_id -= self.max_frame_num;
+                }
+                if pic_id < min_id {
+                    min_id = pic_id;
+                    min_idx = i;
+                }
+            }
+            self.ref_pics.remove(min_idx);
         }
         if self.ref_pics.is_empty() || self.ref_pics.last().unwrap().full_id < cpic.full_id {
             self.ref_pics.push(cpic);
@@ -407,6 +426,35 @@ impl FrameRefs {
             };
         [ref0, ref1]
     }
+    pub fn cmp_refs(&self, ref1: [PicRef; 2], ref2: [PicRef; 2]) -> bool {
+        if ref1 != ref2 {
+            self.cmp_ref(ref1[0], ref2[0], 0) && self.cmp_ref(ref1[1], ref2[1], 1)
+        } else {
+            true
+        }
+    }
+    fn cmp_ref(&self, ref1: PicRef, ref2: PicRef, list: u8) -> bool {
+        if ref1 == ref2 {
+            true
+        } else {
+            let idx0 = ref1.index();
+            let idx1 = ref2.index();
+            if idx0 == idx1 {
+                return true;
+            }
+            let src = if list == 0 { &self.ref_list0 } else { &self.ref_list1 };
+            if idx0 >= src.len() || idx1 >= src.len() {
+//panic!("wrong refs");
+                return false;
+            }
+            if let (Some(ref pic0), Some(ref pic1)) = (&src[idx0], &src[idx1]) {
+                pic0.full_id == pic1.full_id
+            } else {
+//panic!("missing pics");
+                false
+            }
+        }
+    }
 }
 
 fn form_ref_list(ref_list: &mut Vec<Option<PictureInfo>>, ref_pics: &[PictureInfo], long_term: &[Option<PictureInfo>], reord_info: &ReorderingInfo, cur_id: u16, pic_num_mask: u16) {