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