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 }
 207
 208 impl SeekIndex {
 209     pub fn new() -> Self { Self::default() }
 210     pub fn add_stream(&mut self, id: u32) -> usize {
 211         let ret = self.stream_id_to_index(id);
 212         if ret.is_none() {
 213             self.seek_info.push(StreamSeekInfo::new(id));
 214             self.seek_info.len() - 1
 215         } else {
 216             ret.unwrap()
 217         }
 218     }
 219     pub fn stream_id_to_index(&self, id: u32) -> Option<usize> {
 220         for (idx, str) in self.seek_info.iter().enumerate() {
 221             if str.id == id {
 222                 return Some(idx);
 223             }
 224         }
 225         None
 226     }
 227     pub fn get_stream_index(&mut self, id: u32) -> Option<&mut StreamSeekInfo> {
 228         for str in self.seek_info.iter_mut() {
 229             if str.id == id {
 230                 return Some(str);
 231             }
 232         }
 233         None
 234     }
 235     pub fn add_entry(&mut self, id: u32, entry: SeekEntry) {
 236         let mut idx = self.stream_id_to_index(id);
 237         if idx.is_none() {
 238             idx = Some(self.add_stream(id));
 239         }
 240         self.seek_info[idx.unwrap()].add_entry(entry);
 241         self.seek_info[idx.unwrap()].filled = true;
 242     }
 243     pub fn find_pos(&self, time: u64) -> Option<SeekIndexResult> {
 244         let mut cand = None;
 245         for str in self.seek_info.iter() {
 246             if !str.filled { continue; }
 247             let res = str.find_pos(time);
 248             if res.is_none() { continue; }
 249             let res = res.unwrap();
 250             if cand.is_none() {
 251                 cand = Some(SeekIndexResult { pts: res.pts, pos: res.pos, str_id: str.id });
 252             } else if let Some(entry) = cand {
 253                 if res.pos < entry.pos {
 254                     cand = Some(SeekIndexResult { pts: res.pts, pos: res.pos, str_id: str.id });
 255                 }
 256             }
 257         }
 258         cand
 259     }
 260 }
 261
 262 pub struct Demuxer<'a> {
 263     dmx:        Box<dyn DemuxCore<'a> + 'a>,
 264     streams:    StreamManager,
 265     seek_idx:   SeekIndex,
 266 }
 267
 268 impl<'a> Demuxer<'a> {
 269     fn new(dmx: Box<dyn DemuxCore<'a> + 'a>, str: StreamManager, seek_idx: SeekIndex) -> Self {
 270         Demuxer {
 271             dmx,
 272             streams:    str,
 273             seek_idx,
 274         }
 275     }
 276     pub fn get_stream(&self, idx: usize) -> Option<NAStreamRef> {
 277         self.streams.get_stream(idx)
 278     }
 279     pub fn get_stream_by_id(&self, id: u32) -> Option<NAStreamRef> {
 280         self.streams.get_stream_by_id(id)
 281     }
 282     pub fn get_num_streams(&self) -> usize {
 283         self.streams.get_num_streams()
 284     }
 285     pub fn get_streams(&self) -> StreamIter {
 286         self.streams.iter()
 287     }
 288     pub fn is_ignored_stream(&self, idx: usize) -> bool {
 289         self.streams.is_ignored(idx)
 290     }
 291     pub fn set_ignored_stream(&mut self, idx: usize) {
 292         self.streams.set_ignored(idx)
 293     }
 294     pub fn set_unignored_stream(&mut self, idx: usize) {
 295         self.streams.set_unignored(idx)
 296     }
 297
 298     pub fn get_frame(&mut self) -> DemuxerResult<NAPacket> {
 299         loop {
 300             let res = self.dmx.get_frame(&mut self.streams);
 301             if self.streams.no_ign || res.is_err() { return res; }
 302             let res = res.unwrap();
 303             let idx = res.get_stream().get_num();
 304             if !self.is_ignored_stream(idx) {
 305                 return Ok(res);
 306             }
 307         }
 308     }
 309     pub fn seek(&mut self, time: u64) -> DemuxerResult<()> {
 310         self.dmx.seek(time, &self.seek_idx)
 311     }
 312     pub fn get_seek_index(&self) -> &SeekIndex {
 313         &self.seek_idx
 314     }
 315 }
 316
 317 impl From<ByteIOError> for DemuxerError {
 318     fn from(_: ByteIOError) -> Self { DemuxerError::IOError }
 319 }
 320
 321 ///The structure used to create demuxers.
 322 pub trait DemuxerCreator {
 323     /// Create new demuxer instance that will use `ByteReader` source as an input.
 324     fn new_demuxer<'a>(&self, br: &'a mut ByteReader<'a>) -> Box<dyn DemuxCore<'a> + 'a>;
 325     /// Get the name of current demuxer creator.
 326     fn get_name(&self) -> &'static str;
 327 }
 328
 329 pub fn create_demuxer<'a>(dmxcr: &DemuxerCreator, br: &'a mut ByteReader<'a>) -> DemuxerResult<Demuxer<'a>> {
 330     let mut dmx = dmxcr.new_demuxer(br);
 331     let mut str = StreamManager::new();
 332     let mut seek_idx = SeekIndex::new();
 333     dmx.open(&mut str, &mut seek_idx)?;
 334     Ok(Demuxer::new(dmx, str, seek_idx))
 335 }
 336
 337 #[derive(Default)]
 338 pub struct RegisteredDemuxers {
 339     dmxs:   Vec<&'static DemuxerCreator>,
 340 }
 341
 342 impl RegisteredDemuxers {
 343     pub fn new() -> Self {
 344         Self { dmxs: Vec::new() }
 345     }
 346     pub fn add_demuxer(&mut self, dmx: &'static DemuxerCreator) {
 347         self.dmxs.push(dmx);
 348     }
 349     pub fn find_demuxer(&self, name: &str) -> Option<&DemuxerCreator> {
 350         for &dmx in self.dmxs.iter() {
 351             if dmx.get_name() == name {
 352                 return Some(dmx);
 353             }
 354         }
 355         None
 356     }
 357 }