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 }
  15
  16 pub type DemuxerResult<T> = Result<T, DemuxerError>;
  17
  18 pub trait DemuxCore<'a> {
  19     fn open(&mut self, strmgr: &mut StreamManager) -> DemuxerResult<()>;
  20     fn get_frame(&mut self, strmgr: &mut StreamManager) -> DemuxerResult<NAPacket>;
  21     fn seek(&mut self, time: u64) -> DemuxerResult<()>;
  22 }
  23
  24 pub trait NAPacketReader {
  25     fn read_packet(&mut self, str: NAStreamRef, ts: NATimeInfo, keyframe: bool, size: usize) -> DemuxerResult<NAPacket>;
  26     fn fill_packet(&mut self, pkt: &mut NAPacket) -> DemuxerResult<()>;
  27 }
  28
  29 impl<'a> NAPacketReader for ByteReader<'a> {
  30     fn read_packet(&mut self, str: NAStreamRef, ts: NATimeInfo, kf: bool, size: usize) -> DemuxerResult<NAPacket> {
  31         let mut buf: Vec<u8> = Vec::with_capacity(size);
  32         if buf.capacity() < size { return Err(DemuxerError::MemoryError); }
  33         buf.resize(size, 0);
  34         self.read_buf(buf.as_mut_slice())?;
  35         let pkt = NAPacket::new(str, ts, kf, buf);
  36         Ok(pkt)
  37     }
  38     fn fill_packet(&mut self, pkt: &mut NAPacket) -> DemuxerResult<()> {
  39         let mut refbuf = pkt.get_buffer();
  40         let buf = refbuf.as_mut().unwrap();
  41         self.read_buf(buf.as_mut_slice())?;
  42         Ok(())
  43     }
  44 }
  45
  46 #[derive(Default)]
  47 pub struct StreamManager {
  48     streams: Vec<NAStreamRef>,
  49     ignored: Vec<bool>,
  50     no_ign:  bool,
  51 }
  52
  53 impl StreamManager {
  54     pub fn new() -> Self {
  55         StreamManager {
  56             streams: Vec::new(),
  57             ignored: Vec::new(),
  58             no_ign:  true,
  59         }
  60     }
  61     pub fn iter(&self) -> StreamIter { StreamIter::new(&self.streams) }
  62
  63     pub fn add_stream(&mut self, stream: NAStream) -> Option<usize> {
  64         let stream_num = self.streams.len();
  65         let mut str = stream.clone();
  66         str.set_num(stream_num);
  67         self.streams.push(str.into_ref());
  68         self.ignored.push(false);
  69         Some(stream_num)
  70     }
  71     pub fn get_stream(&self, idx: usize) -> Option<NAStreamRef> {
  72         if idx < self.streams.len() {
  73             Some(self.streams[idx].clone())
  74         } else {
  75             None
  76         }
  77     }
  78     pub fn get_stream_by_id(&self, id: u32) -> Option<NAStreamRef> {
  79         for i in 0..self.streams.len() {
  80             if self.streams[i].get_id() == id {
  81                 return Some(self.streams[i].clone());
  82             }
  83         }
  84         None
  85     }
  86     pub fn get_num_streams(&self) -> usize { self.streams.len() }
  87     pub fn is_ignored(&self, idx: usize) -> bool {
  88         if self.no_ign {
  89             true
  90         } else if idx < self.ignored.len() {
  91             self.ignored[idx]
  92         } else {
  93             false
  94         }
  95     }
  96     pub fn is_ignored_id(&self, id: u32) -> bool {
  97         for i in 0..self.streams.len() {
  98             if self.streams[i].get_id() == id {
  99                 return self.ignored[i];
 100             }
 101         }
 102         false
 103     }
 104     pub fn set_ignored(&mut self, idx: usize) {
 105         if idx < self.ignored.len() {
 106             self.ignored[idx] = true;
 107             self.no_ign = false;
 108         }
 109     }
 110     pub fn set_unignored(&mut self, idx: usize) {
 111         if idx < self.ignored.len() {
 112             self.ignored[idx] = false;
 113         }
 114     }
 115 }
 116
 117 pub struct StreamIter<'a> {
 118     streams:    &'a [NAStreamRef],
 119     pos:        usize,
 120 }
 121
 122 impl<'a> StreamIter<'a> {
 123     pub fn new(streams: &'a [NAStreamRef]) -> Self {
 124         StreamIter { streams, pos: 0 }
 125     }
 126 }
 127
 128 impl<'a> Iterator for StreamIter<'a> {
 129     type Item = NAStreamRef;
 130
 131     fn next(&mut self) -> Option<Self::Item> {
 132         if self.pos >= self.streams.len() { return None; }
 133         let ret = self.streams[self.pos].clone();
 134         self.pos += 1;
 135         Some(ret)
 136     }
 137 }
 138
 139 pub struct Demuxer<'a> {
 140     dmx:        Box<dyn DemuxCore<'a> + 'a>,
 141     streams:    StreamManager,
 142 }
 143
 144 impl<'a> Demuxer<'a> {
 145     fn new(dmx: Box<dyn DemuxCore<'a> + 'a>, str: StreamManager) -> Self {
 146         Demuxer {
 147             dmx,
 148             streams:    str,
 149         }
 150     }
 151     pub fn get_stream(&self, idx: usize) -> Option<NAStreamRef> {
 152         self.streams.get_stream(idx)
 153     }
 154     pub fn get_stream_by_id(&self, id: u32) -> Option<NAStreamRef> {
 155         self.streams.get_stream_by_id(id)
 156     }
 157     pub fn get_num_streams(&self) -> usize {
 158         self.streams.get_num_streams()
 159     }
 160     pub fn get_streams(&self) -> StreamIter {
 161         self.streams.iter()
 162     }
 163     pub fn is_ignored_stream(&self, idx: usize) -> bool {
 164         self.streams.is_ignored(idx)
 165     }
 166     pub fn set_ignored_stream(&mut self, idx: usize) {
 167         self.streams.set_ignored(idx)
 168     }
 169     pub fn set_unignored_stream(&mut self, idx: usize) {
 170         self.streams.set_unignored(idx)
 171     }
 172
 173     pub fn get_frame(&mut self) -> DemuxerResult<NAPacket> {
 174         loop {
 175             let res = self.dmx.get_frame(&mut self.streams);
 176             if self.streams.no_ign || res.is_err() { return res; }
 177             let res = res.unwrap();
 178             let idx = res.get_stream().get_num();
 179             if !self.is_ignored_stream(idx) {
 180                 return Ok(res);
 181             }
 182         }
 183     }
 184     pub fn seek(&mut self, time: u64) -> DemuxerResult<()> {
 185         self.dmx.seek(time)
 186     }
 187 }
 188
 189 impl From<ByteIOError> for DemuxerError {
 190     fn from(_: ByteIOError) -> Self { DemuxerError::IOError }
 191 }
 192
 193 ///The structure used to create demuxers.
 194 pub trait DemuxerCreator {
 195     /// Create new demuxer instance that will use `ByteReader` source as an input.
 196     fn new_demuxer<'a>(&self, br: &'a mut ByteReader<'a>) -> Box<dyn DemuxCore<'a> + 'a>;
 197     /// Get the name of current demuxer creator.
 198     fn get_name(&self) -> &'static str;
 199 }
 200
 201 pub fn create_demuxer<'a>(dmxcr: &DemuxerCreator, br: &'a mut ByteReader<'a>) -> DemuxerResult<Demuxer<'a>> {
 202     let mut dmx = dmxcr.new_demuxer(br);
 203     let mut str = StreamManager::new();
 204     dmx.open(&mut str)?;
 205     Ok(Demuxer::new(dmx, str))
 206 }
 207
 208 #[derive(Default)]
 209 pub struct RegisteredDemuxers {
 210     dmxs:   Vec<&'static DemuxerCreator>,
 211 }
 212
 213 impl RegisteredDemuxers {
 214     pub fn new() -> Self {
 215         Self { dmxs: Vec::new() }
 216     }
 217     pub fn add_demuxer(&mut self, dmx: &'static DemuxerCreator) {
 218         self.dmxs.push(dmx);
 219     }
 220     pub fn find_demuxer(&self, name: &str) -> Option<&DemuxerCreator> {
 221         for &dmx in self.dmxs.iter() {
 222             if dmx.get_name() == name {
 223                 return Some(dmx);
 224             }
 225         }
 226         None
 227     }
 228 }