nihav-core/src/reorder.rs

   1 //! Output frame reordering.
   2 //!
   3 //! NihAV decoders output frames in the same order as they are put in.
   4 //! In result if you want to have frames in display order you might need some frame reorderer.
   5 //! This module provides such functionality depending on codec type: audio codecs and video codecs without B-frames do not need any reorderer and can use `NoReorderer` if the common interface is required. Codecs with B-frames should use `IPBReorderer`. For codecs with very complex reordering rules like H.264 or H.256 `PictureIDReorderer` will be added eventually.
   6 //!
   7 //! You can find out required reorderer by quering codec properties using `nihav_core::register` module.
   8 use std::mem::swap;
   9 use std::collections::VecDeque;
  10 pub use crate::frame::{FrameType, NAFrameRef};
  11
  12 /// A trait for frame reorderer.
  13 pub trait FrameReorderer {
  14     /// Stores a newly decoded frame.
  15     fn add_frame(&mut self, fref: NAFrameRef) -> bool;
  16     /// Gets the next frame to be displayed (or `None` if that is not possible).
  17     fn get_frame(&mut self) -> Option<NAFrameRef>;
  18     /// Clears all stored frames.
  19     fn flush(&mut self);
  20     /// Retrieves the last frames stored by the reorderer.
  21     fn get_last_frames(&mut self) -> Option<NAFrameRef>;
  22 }
  23
  24 /// Zero reorderer.
  25 pub struct NoReorderer {
  26     fref:   Option<NAFrameRef>,
  27 }
  28
  29 impl NoReorderer {
  30     /// Constructs a new instance of `NoReorderer`.
  31     pub fn new() -> Self {
  32         Self { fref: None }
  33     }
  34 }
  35
  36 impl Default for NoReorderer {
  37     fn default() -> Self {
  38         Self::new()
  39     }
  40 }
  41
  42 impl FrameReorderer for NoReorderer {
  43     fn add_frame(&mut self, fref: NAFrameRef) -> bool {
  44         if self.fref.is_none() {
  45             self.fref = Some(fref);
  46             true
  47         } else {
  48             false
  49         }
  50     }
  51     fn get_frame(&mut self) -> Option<NAFrameRef> {
  52         let mut ret = None;
  53         swap(&mut ret, &mut self.fref);
  54         ret
  55     }
  56     fn flush(&mut self) { self.fref = None; }
  57     fn get_last_frames(&mut self) -> Option<NAFrameRef> { None }
  58 }
  59
  60 /// Frame reorderer for codecs with I/P/B frames.
  61 #[derive(Default)]
  62 pub struct IPBReorderer {
  63     rframe:     Option<NAFrameRef>,
  64     bframe:     Option<NAFrameRef>,
  65 }
  66
  67 impl IPBReorderer {
  68     /// Constructs a new instance of `IPBReorderer`.
  69     pub fn new() -> Self { Self::default() }
  70 }
  71
  72 impl FrameReorderer for IPBReorderer {
  73     fn add_frame(&mut self, fref: NAFrameRef) -> bool {
  74         if self.rframe.is_some() && self.bframe.is_some() { return false; }
  75         let is_b = fref.get_frame_type() == FrameType::B;
  76         if is_b && self.bframe.is_some() { return false; }
  77         if is_b {
  78             self.bframe = Some(fref);
  79         } else {
  80             std::mem::swap(&mut self.bframe, &mut self.rframe);
  81             self.rframe = Some(fref);
  82         }
  83         true
  84     }
  85     fn get_frame(&mut self) -> Option<NAFrameRef> {
  86         let mut ret = None;
  87         if self.bframe.is_some() {
  88             std::mem::swap(&mut ret, &mut self.bframe);
  89         }
  90         ret
  91     }
  92     fn flush(&mut self) {
  93         self.rframe = None;
  94         self.bframe = None;
  95     }
  96     fn get_last_frames(&mut self) -> Option<NAFrameRef> {
  97         let mut ret = None;
  98         if self.bframe.is_some() {
  99             std::mem::swap(&mut ret, &mut self.bframe);
 100         } else if self.rframe.is_some() {
 101             std::mem::swap(&mut ret, &mut self.rframe);
 102         }
 103         ret
 104     }
 105 }
 106
 107 /// Frame reorderer for codecs with complex I/P/B frame structure like ITU H.26x.
 108 #[derive(Default)]
 109 pub struct ComplexReorderer {
 110     last_ref_dts:   Option<u64>,
 111     ready_idx:      usize,
 112     frames:         Vec<NAFrameRef>,
 113 }
 114
 115 impl ComplexReorderer {
 116     /// Constructs a new instance of `IPBReorderer`.
 117     pub fn new() -> Self { Self::default() }
 118 }
 119
 120 impl FrameReorderer for ComplexReorderer {
 121     fn add_frame(&mut self, fref: NAFrameRef) -> bool {
 122         if self.frames.len() >= 64 {
 123             return false;
 124         }
 125         let is_ref = fref.frame_type == FrameType::I || fref.frame_type == FrameType::P;
 126         if !is_ref {
 127             if self.frames.is_empty() || fref.get_dts().is_none() {
 128                 self.frames.push(fref);
 129             } else if let Some(new_dts) = fref.get_dts() {
 130                 let mut idx = 0;
 131                 for (i, frm) in self.frames.iter().enumerate() {
 132                     idx = i;
 133                     if let Some(dts) = frm.get_dts() {
 134                         if dts > new_dts {
 135                             break;
 136                         }
 137                     }
 138                 }
 139                 self.frames.insert(idx, fref);
 140             }
 141         } else {
 142             for (i, frm) in self.frames.iter().enumerate() {
 143                 if frm.get_dts() == self.last_ref_dts {
 144                     self.ready_idx = i + 1;
 145                 }
 146             }
 147             self.last_ref_dts = fref.get_dts();
 148             self.frames.push(fref);
 149         }
 150         true
 151     }
 152     fn get_frame(&mut self) -> Option<NAFrameRef> {
 153         if self.ready_idx > 0 {
 154             self.ready_idx -= 1;
 155             Some(self.frames.remove(0))
 156         } else {
 157             None
 158         }
 159     }
 160     fn flush(&mut self) {
 161         self.last_ref_dts = None;
 162         self.ready_idx = 0;
 163         self.frames.clear();
 164     }
 165     fn get_last_frames(&mut self) -> Option<NAFrameRef> {
 166         if !self.frames.is_empty() {
 167             Some(self.frames.remove(0))
 168         } else {
 169             None
 170         }
 171     }
 172 }
 173
 174 /// A generic reorderer for a multi-threaded decoder.
 175 #[derive(Default)]
 176 pub struct MTFrameReorderer {
 177     ids:            Vec<u32>,
 178     frames:         VecDeque<(u32, NAFrameRef)>,
 179     cur_id:         u32,
 180     flush_mode:     bool,
 181     output_to:      Option<u32>,
 182     last_ts:        Option<u64>,
 183 }
 184
 185 impl MTFrameReorderer {
 186     /// Constructs a new instance of `MTFrameReorderer`.
 187     pub fn new() -> Self { Self::default() }
 188     /// Registers the fact that a new frame is queued for decoding and returns an internal ID for it.
 189     pub fn register_frame(&mut self) -> u32 {
 190         if self.flush_mode {
 191             self.flush_mode = false;
 192         }
 193         let ret = self.cur_id;
 194         self.cur_id += 1;
 195         self.ids.push(ret);
 196         ret
 197     }
 198     /// Puts a newly decoded frame into the internal queue.
 199     pub fn add_frame(&mut self, frm: NAFrameRef, id: u32) {
 200         //let ftype = frm.get_frame_type();
 201         let frm_id = if let Some(ts) = frm.ts.dts { ts } else { u64::from(id) };
 202         let mut idx = 0;
 203         for (_, frm) in self.frames.iter() {
 204             let cur_id = if let Some(ts) = frm.ts.dts { ts } else { frm.id as u64 };
 205             if frm_id < cur_id {
 206                 break;
 207             }
 208             idx += 1;
 209         }
 210         self.frames.insert(idx, (id, frm));
 211         /*if self.frames.len() > 48 {
 212             for (id, frm) in self.frames.iter() { print!(" {}{}({})", frm.get_frame_type(), frm.get_dts().unwrap_or(0), *id); } println!();
 213             print!("reg IDs:"); for &id in self.ids.iter() { print!(" {}", id); } println!();
 214             panic!("too many frames in the queue");
 215         }*/
 216     }
 217     /// Removes the registered frame (e.g. in case of a decoding error).
 218     pub fn drop_frame(&mut self, id: u32) {
 219         self.ids.retain(|&el| el != id);
 220     }
 221     fn get_first_frame(&mut self) -> Option<NAFrameRef> {
 222         let (id, frm) = self.frames.pop_front().unwrap();
 223         self.drop_frame(id);
 224         self.last_ts = frm.get_dts();
 225         Some(frm)
 226     }
 227     /// Gets the next frame to be displayed (or `None` if that is not possible).
 228     #[allow(clippy::collapsible_if)]
 229     #[allow(clippy::collapsible_else_if)]
 230     pub fn get_frame(&mut self) -> Option<NAFrameRef> {
 231         // check if we have consequent timestamps that we can output
 232         if !self.frames.is_empty() {
 233             if let Some(dts) = self.frames[0].1.get_dts() {
 234                 let last_ts = self.last_ts.unwrap_or(0);
 235                 if self.last_ts.is_none() || (dts == last_ts + 1) {
 236                     self.output_to = None;
 237                     return self.get_first_frame();
 238                 }
 239             }
 240         }
 241         if !self.flush_mode {
 242             'out_loop: loop {
 243                 if let Some(last_id) = self.output_to {
 244                     if self.frames[0].0 != last_id {
 245                         return self.get_first_frame();
 246                     } else {
 247                         self.output_to = None;
 248                     }
 249                 }
 250                 for (pos, (id, frm)) in self.frames.iter().enumerate() {
 251                     if frm.is_keyframe() || (self.frames.len() > 32 && matches!(frm.get_frame_type(), FrameType::I | FrameType::P)) {
 252                         let kf_id = *id;
 253                         self.ids.sort();
 254                         if pos == 0 && kf_id == self.ids[0] {
 255                             return self.get_first_frame();
 256                         }
 257                         let end = self.ids.iter().position(|&id| id == kf_id).unwrap();
 258                         for ref_id in self.ids[..end].iter() {
 259                             if !self.frames.iter().any(|(id, _)| id == ref_id) {
 260                                 return None;
 261                             }
 262                         }
 263                         self.output_to = if pos < self.frames.len() - 1 {
 264                                 Some(self.frames[pos + 1].0)
 265                             } else {
 266                                 Some(kf_id)
 267                             };
 268                         continue 'out_loop;
 269                     }
 270                 }
 271                 return None;
 272             }
 273         } else {
 274             if !self.frames.is_empty() {
 275                 Some(self.frames.pop_front().unwrap().1)
 276             } else {
 277                 None
 278             }
 279         }
 280     }
 281     /// Retrieves the last frames stored by the reorderer.
 282     pub fn get_last_frames(&mut self) -> Option<NAFrameRef> {
 283         self.flush_mode = true;
 284         self.get_frame()
 285     }
 286     /// Clears all stored frames.
 287     pub fn flush(&mut self) {
 288         self.flush_mode = false;
 289         self.frames.clear();
 290         self.ids.clear();
 291         self.output_to = None;
 292         self.last_ts = None;
 293     }
 294 }