nihav-core/src/demuxers/mod.rs

   1 pub use crate::frame::*;
   2 pub use crate::io::byteio::*;
   3
   4 #[derive(Debug,Clone,Copy,PartialEq)]
   5 #[allow(dead_code)]
   6 pub enum DemuxerError {
   7     EOF,
   8     NoSuchInput,
   9     InvalidData,
  10     IOError,
  11     NotImplemented,
  12     MemoryError,
  13     TryAgain,
  14     SeekError,
  15     NotPossible,
  16 }
  17
  18 pub type DemuxerResult<T> = Result<T, DemuxerError>;
  19
  20 pub trait DemuxCore<'a> {
  21     fn open(&mut self, strmgr: &mut StreamManager, seek_idx: &mut SeekIndex) -> DemuxerResult<()>;
  22     fn get_frame(&mut self, strmgr: &mut StreamManager) -> DemuxerResult<NAPacket>;
  23     fn seek(&mut self, time: u64, seek_idx: &SeekIndex) -> DemuxerResult<()>;
  24 }
  25
  26 pub trait NAPacketReader {
  27     fn read_packet(&mut self, str: NAStreamRef, ts: NATimeInfo, keyframe: bool, size: usize) -> DemuxerResult<NAPacket>;
  28     fn fill_packet(&mut self, pkt: &mut NAPacket) -> DemuxerResult<()>;
  29 }
  30
  31 impl<'a> NAPacketReader for ByteReader<'a> {
  32     fn read_packet(&mut self, str: NAStreamRef, ts: NATimeInfo, kf: bool, size: usize) -> DemuxerResult<NAPacket> {
  33         let mut buf: Vec<u8> = Vec::with_capacity(size);
  34         if buf.capacity() < size { return Err(DemuxerError::MemoryError); }
  35         buf.resize(size, 0);
  36         self.read_buf(buf.as_mut_slice())?;
  37         let pkt = NAPacket::new(str, ts, kf, buf);
  38         Ok(pkt)
  39     }
  40     fn fill_packet(&mut self, pkt: &mut NAPacket) -> DemuxerResult<()> {
  41         let mut refbuf = pkt.get_buffer();
  42         let buf = refbuf.as_mut().unwrap();
  43         self.read_buf(buf.as_mut_slice())?;
  44         Ok(())
  45     }
  46 }
  47
  48 #[derive(Default)]
  49 pub struct StreamManager {
  50     streams: Vec<NAStreamRef>,
  51     ignored: Vec<bool>,
  52     no_ign:  bool,
  53 }
  54
  55 impl StreamManager {
  56     pub fn new() -> Self {
  57         StreamManager {
  58             streams: Vec::new(),
  59             ignored: Vec::new(),
  60             no_ign:  true,
  61         }
  62     }
  63     pub fn iter(&self) -> StreamIter { StreamIter::new(&self.streams) }
  64
  65     pub fn add_stream(&mut self, stream: NAStream) -> Option<usize> {
  66         let stream_num = self.streams.len();
  67         let mut str = stream.clone();
  68         str.set_num(stream_num);
  69         self.streams.push(str.into_ref());
  70         self.ignored.push(false);
  71         Some(stream_num)
  72     }
  73     pub fn get_stream(&self, idx: usize) -> Option<NAStreamRef> {
  74         if idx < self.streams.len() {
  75             Some(self.streams[idx].clone())
  76         } else {
  77             None
  78         }
  79     }
  80     pub fn get_stream_by_id(&self, id: u32) -> Option<NAStreamRef> {
  81         for i in 0..self.streams.len() {
  82             if self.streams[i].get_id() == id {
  83                 return Some(self.streams[i].clone());
  84             }
  85         }
  86         None
  87     }
  88     pub fn get_num_streams(&self) -> usize { self.streams.len() }
  89     pub fn is_ignored(&self, idx: usize) -> bool {
  90         if self.no_ign {
  91             true
  92         } else if idx < self.ignored.len() {
  93             self.ignored[idx]
  94         } else {
  95             false
  96         }
  97     }
  98     pub fn is_ignored_id(&self, id: u32) -> bool {
  99         for i in 0..self.streams.len() {
 100             if self.streams[i].get_id() == id {
 101                 return self.ignored[i];
 102             }
 103         }
 104         false
 105     }
 106     pub fn set_ignored(&mut self, idx: usize) {
 107         if idx < self.ignored.len() {
 108             self.ignored[idx] = true;
 109             self.no_ign = false;
 110         }
 111     }
 112     pub fn set_unignored(&mut self, idx: usize) {
 113         if idx < self.ignored.len() {
 114             self.ignored[idx] = false;
 115         }
 116     }
 117 }
 118
 119 pub struct StreamIter<'a> {
 120     streams:    &'a [NAStreamRef],
 121     pos:        usize,
 122 }
 123
 124 impl<'a> StreamIter<'a> {
 125     pub fn new(streams: &'a [NAStreamRef]) -> Self {
 126         StreamIter { streams, pos: 0 }
 127     }
 128 }
 129
 130 impl<'a> Iterator for StreamIter<'a> {
 131     type Item = NAStreamRef;
 132
 133     fn next(&mut self) -> Option<Self::Item> {
 134         if self.pos >= self.streams.len() { return None; }
 135         let ret = self.streams[self.pos].clone();
 136         self.pos += 1;
 137         Some(ret)
 138     }
 139 }
 140
 141 #[derive(Clone,Copy,PartialEq)]
 142 pub enum SeekIndexMode {
 143     None,
 144     Present,
 145     Automatic,
 146 }
 147
 148 impl Default for SeekIndexMode {
 149     fn default() -> Self { SeekIndexMode::None }
 150 }
 151
 152 #[derive(Clone,Copy,Default)]
 153 pub struct SeekEntry {
 154     pub time:   u64, // in milliseconds
 155     pub pts:    u64,
 156     pub pos:    u64,
 157 }
 158
 159 #[derive(Clone)]
 160 pub struct StreamSeekInfo {
 161     pub id:         u32,
 162     pub filled:     bool,
 163     pub entries:    Vec<SeekEntry>,
 164 }
 165
 166 impl StreamSeekInfo {
 167     pub fn new(id: u32) -> Self {
 168         Self {
 169             id,
 170             filled:     false,
 171             entries:    Vec::new(),
 172         }
 173     }
 174     pub fn add_entry(&mut self, entry: SeekEntry) {
 175         self.entries.push(entry);
 176     }
 177     pub fn find_pos(&self, time: u64) -> Option<SeekEntry> {
 178         if !self.entries.is_empty() {
 179 // todo something faster like binary search
 180             let mut cand = None;
 181             for entry in self.entries.iter() {
 182                 if entry.time <= time {
 183                     cand = Some(*entry);
 184                 } else {
 185                     break;
 186                 }
 187             }
 188             cand
 189         } else {
 190             None
 191         }
 192     }
 193 }
 194
 195 #[derive(Clone,Copy,Default)]
 196 pub struct SeekIndexResult {
 197     pub pts:        u64,
 198     pub pos:        u64,
 199     pub str_id:     u32,
 200 }
 201
 202 #[derive(Default)]
 203 pub struct SeekIndex {
 204     pub seek_info:  Vec<StreamSeekInfo>,
 205     pub mode:       SeekIndexMode,
 206     pub skip_index: bool,
 207 }
 208
 209 impl SeekIndex {
 210     pub fn new() -> Self { Self::default() }
 211     pub fn add_stream(&mut self, id: u32) -> usize {
 212         let ret = self.stream_id_to_index(id);
 213         if ret.is_none() {
 214             self.seek_info.push(StreamSeekInfo::new(id));
 215             self.seek_info.len() - 1
 216         } else {
 217             ret.unwrap()
 218         }
 219     }
 220     pub fn stream_id_to_index(&self, id: u32) -> Option<usize> {
 221         for (idx, str) in self.seek_info.iter().enumerate() {
 222             if str.id == id {
 223                 return Some(idx);
 224             }
 225         }
 226         None
 227     }
 228     pub fn get_stream_index(&mut self, id: u32) -> Option<&mut StreamSeekInfo> {
 229         for str in self.seek_info.iter_mut() {
 230             if str.id == id {
 231                 return Some(str);
 232             }
 233         }
 234         None
 235     }
 236     pub fn add_entry(&mut self, id: u32, entry: SeekEntry) {
 237         let mut idx = self.stream_id_to_index(id);
 238         if idx.is_none() {
 239             idx = Some(self.add_stream(id));
 240         }
 241         self.seek_info[idx.unwrap()].add_entry(entry);
 242         self.seek_info[idx.unwrap()].filled = true;
 243     }
 244     pub fn find_pos(&self, time: u64) -> Option<SeekIndexResult> {
 245         let mut cand = None;
 246         for str in self.seek_info.iter() {
 247             if !str.filled { continue; }
 248             let res = str.find_pos(time);
 249             if res.is_none() { continue; }
 250             let res = res.unwrap();
 251             if cand.is_none() {
 252                 cand = Some(SeekIndexResult { pts: res.pts, pos: res.pos, str_id: str.id });
 253             } else if let Some(entry) = cand {
 254                 if res.pos < entry.pos {
 255                     cand = Some(SeekIndexResult { pts: res.pts, pos: res.pos, str_id: str.id });
 256                 }
 257             }
 258         }
 259         cand
 260     }
 261 }
 262
 263 pub struct Demuxer<'a> {
 264     dmx:        Box<dyn DemuxCore<'a> + 'a>,
 265     streams:    StreamManager,
 266     seek_idx:   SeekIndex,
 267 }
 268
 269 impl<'a> Demuxer<'a> {
 270     fn new(dmx: Box<dyn DemuxCore<'a> + 'a>, str: StreamManager, seek_idx: SeekIndex) -> Self {
 271         Demuxer {
 272             dmx,
 273             streams:    str,
 274             seek_idx,
 275         }
 276     }
 277     pub fn get_stream(&self, idx: usize) -> Option<NAStreamRef> {
 278         self.streams.get_stream(idx)
 279     }
 280     pub fn get_stream_by_id(&self, id: u32) -> Option<NAStreamRef> {
 281         self.streams.get_stream_by_id(id)
 282     }
 283     pub fn get_num_streams(&self) -> usize {
 284         self.streams.get_num_streams()
 285     }
 286     pub fn get_streams(&self) -> StreamIter {
 287         self.streams.iter()
 288     }
 289     pub fn is_ignored_stream(&self, idx: usize) -> bool {
 290         self.streams.is_ignored(idx)
 291     }
 292     pub fn set_ignored_stream(&mut self, idx: usize) {
 293         self.streams.set_ignored(idx)
 294     }
 295     pub fn set_unignored_stream(&mut self, idx: usize) {
 296         self.streams.set_unignored(idx)
 297     }
 298
 299     pub fn get_frame(&mut self) -> DemuxerResult<NAPacket> {
 300         loop {
 301             let res = self.dmx.get_frame(&mut self.streams);
 302             if self.streams.no_ign || res.is_err() { return res; }
 303             let res = res.unwrap();
 304             let idx = res.get_stream().get_num();
 305             if !self.is_ignored_stream(idx) {
 306                 return Ok(res);
 307             }
 308         }
 309     }
 310     pub fn seek(&mut self, time: u64) -> DemuxerResult<()> {
 311         if self.seek_idx.skip_index {
 312             return Err(DemuxerError::NotPossible);
 313         }
 314         self.dmx.seek(time, &self.seek_idx)
 315     }
 316     pub fn get_seek_index(&self) -> &SeekIndex {
 317         &self.seek_idx
 318     }
 319 }
 320
 321 impl From<ByteIOError> for DemuxerError {
 322     fn from(_: ByteIOError) -> Self { DemuxerError::IOError }
 323 }
 324
 325 ///The structure used to create demuxers.
 326 pub trait DemuxerCreator {
 327     /// Create new demuxer instance that will use `ByteReader` source as an input.
 328     fn new_demuxer<'a>(&self, br: &'a mut ByteReader<'a>) -> Box<dyn DemuxCore<'a> + 'a>;
 329     /// Get the name of current demuxer creator.
 330     fn get_name(&self) -> &'static str;
 331 }
 332
 333 pub fn create_demuxer<'a>(dmxcr: &DemuxerCreator, br: &'a mut ByteReader<'a>) -> DemuxerResult<Demuxer<'a>> {
 334     let mut dmx = dmxcr.new_demuxer(br);
 335     let mut str = StreamManager::new();
 336     let mut seek_idx = SeekIndex::new();
 337     dmx.open(&mut str, &mut seek_idx)?;
 338     Ok(Demuxer::new(dmx, str, seek_idx))
 339 }
 340
 341 #[derive(Default)]
 342 pub struct RegisteredDemuxers {
 343     dmxs:   Vec<&'static DemuxerCreator>,
 344 }
 345
 346 impl RegisteredDemuxers {
 347     pub fn new() -> Self {
 348         Self { dmxs: Vec::new() }
 349     }
 350     pub fn add_demuxer(&mut self, dmx: &'static DemuxerCreator) {
 351         self.dmxs.push(dmx);
 352     }
 353     pub fn find_demuxer(&self, name: &str) -> Option<&DemuxerCreator> {
 354         for &dmx in self.dmxs.iter() {
 355             if dmx.get_name() == name {
 356                 return Some(dmx);
 357             }
 358         }
 359         None
 360     }
 361 }