mkv: implement lacing

diff --git a/nihav-commonfmt/src/demuxers/mkv/mod.rs b/nihav-commonfmt/src/demuxers/mkv/mod.rs
index 60f2ad5c7282964a3e9adf023eceee38b7fa1f75..da034f8e9e8625f1658b1a51e167824152a22734 100644 (file)
--- a/nihav-commonfmt/src/demuxers/mkv/mod.rs
+++ b/nihav-commonfmt/src/demuxers/mkv/mod.rs
@@ -1,3 +1,4 @@
+use std::collections::VecDeque;
 use nihav_core::demuxers::*;
 
 #[allow(clippy::neg_cmp_op_on_partial_ord)]
@@ -66,6 +67,7 @@ impl std::convert::From<CueEntry> for MSeekInfo {
 struct MatroskaDemuxer<'a> {
     src:            &'a mut dyn ByteIO,
     segments:       Vec<Segment>,
+    queue:          VecDeque<NAPacket>,
 
     cur_seg:        usize,
     seg_end:        Option<u64>,
@@ -135,6 +137,9 @@ impl<'a> DemuxCore<'a> for MatroskaDemuxer<'a> {
     }
 
     fn get_frame(&mut self, strmgr: &mut StreamManager) -> DemuxerResult<NAPacket> {
+        if let Some(pkt) = self.queue.pop_front() {
+            return Ok(pkt);
+        }
         'main_loop: loop {
             if let Some(clu_end) = self.clu_end {
                 while self.src.tell() < clu_end {
@@ -173,16 +178,73 @@ impl<'a> DemuxCore<'a> for MatroskaDemuxer<'a> {
                             if trk_id > U32_MAX { return Err(DemuxerError::NotImplemented); }
                             let trk_id = trk_id as u32;
                             validate!(self.src.tell() < end);
-                            if lacing != 0 {
-                                if let Some(trk) = self.segments[self.cur_seg].tracks.iter().find(|x| x.id == trk_id) {
-                                    validate!(trk.lacing);
-                                }
-                                return Err(DemuxerError::NotImplemented);
-                            }
                             if let Some(stream) = strmgr.get_stream_by_id(trk_id) {
-                                let pkt_size = (end - self.src.tell()) as usize;
-                                let ts = NATimeInfo::new(Some(full_ts), None, None, stream.tb_num, stream.tb_den);
-                                return self.src.read_packet(stream, ts, keyframe, pkt_size);
+                                if lacing != 0 {
+                                    if let Some(trk) = self.segments[self.cur_seg].tracks.iter().find(|x| x.id == trk_id) {
+                                        validate!(trk.lacing);
+                                    }
+                                    let nframes = usize::from(self.src.read_byte()?) + 1;
+                                    validate!(nframes > 1 && self.src.tell() < end);
+                                    let mut sizes = [0; 256];
+                                    let mut tot_size = 0;
+                                    match lacing {
+                                        1 => { // Xiph lacing
+                                            for dsize in sizes[..nframes - 1].iter_mut() {
+                                                let mut b = self.src.read_byte()?;
+                                                let mut size = usize::from(b);
+                                                while b == 255 {
+                                                    b = self.src.read_byte()?;
+                                                    size += usize::from(b);
+                                                }
+                                                *dsize = size;
+                                                tot_size += size;
+                                            }
+                                            validate!(self.src.tell() + (tot_size as u64) < end);
+                                            sizes[nframes - 1] = (end - self.src.tell()) as usize - tot_size;
+                                        },
+                                        3 => { // EBML lacing
+                                            let first_size = self.src.read_vint()?;
+                                            let mut prev = first_size as i64;
+                                            let fsize = first_size as usize;
+                                            sizes[0] = fsize;
+                                            tot_size += fsize;
+                                            for dsize in sizes[..nframes - 1].iter_mut().skip(1) {
+                                                let diff = self.src.read_vint_s()?;
+                                                let size = prev + diff;
+                                                validate!((1..=(U32_MAX as i64)).contains(&size));
+                                                prev = size;
+                                                let size = size as usize;
+                                                *dsize = size;
+                                                tot_size += size;
+                                            }
+                                            validate!(self.src.tell() + (tot_size as u64) < end);
+                                            sizes[nframes - 1] = (end - self.src.tell()) as usize - tot_size;
+                                        },
+                                        2 => { // fixed-size lacing
+                                            let rest_size = (end - self.src.tell()) as usize;
+                                            validate!((rest_size > 0) && (rest_size % nframes == 0));
+                                            let pkt_size = rest_size / nframes;
+                                            for el in sizes.iter_mut().take(nframes) {
+                                                *el = pkt_size;
+                                            }
+                                        },
+                                        _ => unreachable!()
+                                    }
+                                    let mut ts = NATimeInfo::new(Some(full_ts), None, None, stream.tb_num, stream.tb_den);
+                                    let mut keyframe = keyframe;
+                                    for &pkt_size in sizes[..nframes].iter() {
+                                        let pkt = self.src.read_packet(stream.clone(), ts, keyframe, pkt_size)?;
+                                        keyframe = false;
+                                        ts.pts = None;
+                                        self.queue.push_back(pkt);
+                                    }
+                                    validate!(self.src.tell() == end);
+                                    return Ok(self.queue.pop_front().unwrap());
+                                } else {
+                                    let pkt_size = (end - self.src.tell()) as usize;
+                                    let ts = NATimeInfo::new(Some(full_ts), None, None, stream.tb_num, stream.tb_den);
+                                    return self.src.read_packet(stream, ts, keyframe, pkt_size);
+                                }
                             }
                         },
                         _ => {},
@@ -248,6 +310,7 @@ impl<'a> DemuxCore<'a> for MatroskaDemuxer<'a> {
             self.clu_end = None;
             self.clu_ts  = None;
             self.cur_seg = self.segments.len();
+            self.queue.clear();
 
             let tgt_seg = &self.segments[seg_no];
 
@@ -332,6 +395,7 @@ impl<'a> MatroskaDemuxer<'a> {
             seg_end:        None,
             clu_end:        None,
             clu_ts:         None,
+            queue:          VecDeque::new(),
         }
     }
 }

diff --git a/nihav-commonfmt/src/demuxers/mkv/parse.rs b/nihav-commonfmt/src/demuxers/mkv/parse.rs
index cda0615f9bcceea834bce4fbe02e7ca768ba794f..bf8d86d1bda0feffb9f132dc6f4a96115d0e6d46 100644 (file)
--- a/nihav-commonfmt/src/demuxers/mkv/parse.rs
+++ b/nihav-commonfmt/src/demuxers/mkv/parse.rs
@@ -18,6 +18,7 @@ fn gcd(mut a: u32, mut b: u32) -> u32 {
 pub trait ReadEBML {
     fn read_ebml_id(&mut self) -> DemuxerResult<u32>;
     fn read_vint(&mut self) -> DemuxerResult<u64>;
+    fn read_vint_s(&mut self) -> DemuxerResult<i64>;
     fn read_uint(&mut self, size: u64) -> DemuxerResult<u64>;
     fn read_float(&mut self, size: u64) -> DemuxerResult<f64>;
 }
@@ -86,6 +87,40 @@ impl<T: ?Sized + ByteIO> ReadEBML for T {
             },
         }
     }
+    fn read_vint_s(&mut self) -> DemuxerResult<i64> {
+        let b       = self.read_byte()?;
+        match b {
+            0 => Err(DemuxerError::InvalidData),
+            0x80..=0xFF => { // 1 + 7 bits
+                Ok(i64::from(b & 0x7F) - (1 << 6) + 1)
+            },
+            0x40..=0x7F => { // 01 + 14 bits
+                let low = self.read_byte()?;
+                Ok(((u64::from(b & 0x3F) << 8) | u64::from(low)) as i64 - (1 << 13) + 1)
+            },
+            0x20..=0x3F => { // 001 + 21 bit
+                let low = self.read_u16be()?;
+                Ok(((u64::from(b & 0x1F) << 16) | u64::from(low)) as i64 - (1 << 20) + 1)
+            },
+            0x10..=0x1F => { // 0001 + 28 bits
+                let low = self.read_u24be()?;
+                Ok(((u64::from(b & 0xF) << 24) | u64::from(low)) as i64 - (1 << 27) + 1)
+            },
+            0x08..=0x0F => { // 0000 1 + 37 bits
+                let low = self.read_u32be()?;
+                Ok(((u64::from(b & 0x7) << 32) | u64::from(low)) as i64 - (1 << 34) + 1)
+            },
+            _ => {
+                let len = b.leading_zeros();
+                let mut val = if len < 7 { u64::from((b << (len + 1)) >> (len + 1)) } else { 0 };
+                for _ in 0..len {
+                    val <<= 8;
+                    val |= u64::from(self.read_byte()?);
+                }
+                Ok(val as i64 - (1 << (len * 7 - 1)) + 1)
+            },
+        }
+    }
     fn read_uint(&mut self, size: u64) -> DemuxerResult<u64> {
         match size {
             1 => Ok(u64::from(self.read_byte()?)),