From 8480761d25359f1f88c6d6df020127a7a7ea6c74 Mon Sep 17 00:00:00 2001
From: Kostya Shishkov <kostya.shishkov@gmail.com>
Date: Sat, 12 Dec 2020 14:45:18 +0100
Subject: [PATCH] add complex frame reorderer for upcoming H.264

---
 nihav-core/src/reorder.rs      | 67 ++++++++++++++++++++++++++++++++++
 nihav-registry/src/register.rs |  2 +
 2 files changed, 69 insertions(+)

diff --git a/nihav-core/src/reorder.rs b/nihav-core/src/reorder.rs
index 625e759..6ae7e1a 100644
--- a/nihav-core/src/reorder.rs
+++ b/nihav-core/src/reorder.rs
@@ -103,3 +103,70 @@ impl FrameReorderer for IPBReorderer {
     }
 }
 
+/// Frame reorderer for codecs with complex I/P/B frame structure like ITU H.26x.
+#[derive(Default)]
+pub struct ComplexReorderer {
+    last_ref_dts:   Option<u64>,
+    ready_idx:      usize,
+    frames:         Vec<NAFrameRef>,
+}
+
+impl ComplexReorderer {
+    /// Constructs a new instance of `IPBReorderer`.
+    pub fn new() -> Self { Self::default() }
+}
+
+impl FrameReorderer for ComplexReorderer {
+    fn add_frame(&mut self, fref: NAFrameRef) -> bool {
+        if self.frames.len() >= 64 {
+            return false;
+        }
+        let is_ref = fref.frame_type == FrameType::I || fref.frame_type == FrameType::P;
+        if !is_ref {
+            if self.frames.is_empty() || fref.get_dts().is_none() {
+                self.frames.push(fref);
+            } else if let Some(new_dts) = fref.get_dts() {
+                let mut idx = 0;
+                for (i, frm) in self.frames.iter().enumerate() {
+                    idx = i;
+                    if let Some(dts) = frm.get_dts() {
+                        if dts > new_dts {
+                            break;
+                        }
+                    }
+                }
+                self.frames.insert(idx, fref);
+            }
+        } else {
+            for (i, frm) in self.frames.iter().enumerate() {
+                if frm.get_dts() == self.last_ref_dts {
+                    self.ready_idx = i + 1;
+                }
+            }
+            self.last_ref_dts = fref.get_dts();
+            self.frames.push(fref);
+        }
+        true
+    }
+    fn get_frame(&mut self) -> Option<NAFrameRef> {
+        if self.ready_idx > 0 {
+            self.ready_idx -= 1;
+            Some(self.frames.remove(0))
+        } else {
+            None
+        }
+    }
+    fn flush(&mut self) {
+        self.last_ref_dts = None;
+        self.ready_idx = 0;
+        self.frames.truncate(0);
+    }
+    fn get_last_frames(&mut self) -> Option<NAFrameRef> {
+        if !self.frames.is_empty() {
+            Some(self.frames.remove(0))
+        } else {
+            None
+        }
+    }
+}
+
diff --git a/nihav-registry/src/register.rs b/nihav-registry/src/register.rs
index 1ada162..78ccb18 100644
--- a/nihav-registry/src/register.rs
+++ b/nihav-registry/src/register.rs
@@ -41,6 +41,8 @@ pub const CODEC_CAP_REORDER:u32     = 0x0004;
 pub const CODEC_CAP_HYBRID:u32      = 0x0008;
 /// Codec capability flag for codecs with scalability features.
 pub const CODEC_CAP_SCALABLE:u32    = 0x0010;
+/// Codec capability flag for codecs with complex frame reordering.
+pub const CODEC_CAP_COMPLEX_REORDER:u32 = 0x0020;
 
 /// Codec description structure.
 #[derive(Clone)]
-- 
2.30.2