rmmux: improve the audio packet timestamp handling

3 files changed, 56 insertions, 17 deletions

diff --git a/nihav-realmedia/src/muxers/rmvb/audiostream.rs b/nihav-realmedia/src/muxers/rmvb/audiostream.rs
index fcb3d9c..d6daddd 100644
--- a/nihav-realmedia/src/muxers/rmvb/audiostream.rs
+++ b/nihav-realmedia/src/muxers/rmvb/audiostream.rs
@@ -26,6 +26,7 @@ trait Interleaver {
     fn get_block_size(&self) -> usize;
     fn get_factor(&self) -> usize;
     fn get_frames_per_block(&self) -> usize;
+    fn is_empty(&self) -> bool;
     fn add_packet(&mut self, src: &[u8]) -> bool;
     fn get_packet(&mut self) -> Option<(Vec<u8>, bool)>;
     fn flush(&mut self);
@@ -42,6 +43,7 @@ impl Interleaver for NoInterleaver {
     fn get_block_size(&self) -> usize { self.frame_size }
     fn get_factor(&self) -> usize { 1 }
     fn get_frames_per_block(&self) -> usize { 1 }
+    fn is_empty(&self) -> bool { self.pkt.is_none() }
     fn add_packet(&mut self, src: &[u8]) -> bool {
         if self.pkt.is_none() {
             self.pkt = Some(src.to_vec());
@@ -107,6 +109,7 @@ impl Interleaver for Int4Interleaver {
     fn get_block_size(&self) -> usize { self.block_size }
     fn get_factor(&self) -> usize { self.factor }
     fn get_frames_per_block(&self) -> usize { self.fpb }
+    fn is_empty(&self) -> bool { self.cur_frame == 0 }
     fn add_packet(&mut self, src: &[u8]) -> bool {
         if self.cur_frame == self.factor * self.fpb {
             return false;
@@ -200,6 +203,7 @@ impl Interleaver for GenericInterleaver {
     fn get_block_size(&self) -> usize { self.block_size }
     fn get_factor(&self) -> usize { self.factor }
     fn get_frames_per_block(&self) -> usize { self.fpb }
+    fn is_empty(&self) -> bool { self.cur_frame == 0 }
     fn add_packet(&mut self, src: &[u8]) -> bool {
         if self.cur_frame == self.factor * self.fpb {
             return false;
@@ -266,6 +270,7 @@ impl Interleaver for SiproInterleaver {
     fn get_block_size(&self) -> usize { self.block_size }
     fn get_factor(&self) -> usize { self.factor }
     fn get_frames_per_block(&self) -> usize { 1 }
+    fn is_empty(&self) -> bool { self.wr_pos == 0 }
     fn add_packet(&mut self, src: &[u8]) -> bool {
         if self.wr_pos < self.factor {
             self.buf[self.wr_pos * self.block_size..][..self.block_size].copy_from_slice(src);
@@ -340,6 +345,7 @@ impl Interleaver for AACInterleaver {
     fn get_block_size(&self) -> usize { self.frame_size }
     fn get_factor(&self) -> usize { 1 }
     fn get_frames_per_block(&self) -> usize { 1 }
+    fn is_empty(&self) -> bool { self.data.is_empty() }
     fn add_packet(&mut self, src: &[u8]) -> bool {
         if !self.full {
             self.data.extend_from_slice(src);
@@ -396,6 +402,10 @@ struct AudioStreamWriter {
     header_pos:     u64,
     interleave:     Box<dyn Interleaver>,
     data_size:      usize,
+    first_time:     u32,
+    last_time:      u32,
+    size_in:        usize,
+    size_out:       usize,
 }
 
 impl RMStreamWriter for AudioStreamWriter {
@@ -424,10 +434,18 @@ impl RMStreamWriter for AudioStreamWriter {
         }
         Ok(())
     }
-    fn queue_packet(&mut self, pkt: NAPacket) -> bool {
+    fn queue_packet(&mut self, pkt: NAPacket, ms: u32) -> bool {
         let src = pkt.get_buffer();
         self.data_size += src.len();
         let frame_size = self.interleave.get_frame_size();
+        self.last_time = ms;
+        if self.interleave.is_empty() {
+            self.first_time = ms;
+            self.size_in = 0;
+            self.size_out = 0;
+        } else {
+            self.size_in += src.len() / frame_size;
+        }
         if !self.is_raw {
             let mut ret = false;
             for frame in src.chunks(frame_size) {
@@ -438,7 +456,21 @@ impl RMStreamWriter for AudioStreamWriter {
             self.interleave.add_packet(&src)
         }
     }
-    fn get_packet(&mut self) -> Option<(Vec<u8>, bool)> { self.interleave.get_packet() }
+    fn get_packet(&mut self) -> Option<(Vec<u8>, u32, bool)> {
+        if let Some((pkt, first)) = self.interleave.get_packet() {
+            let time_add = if self.last_time > self.first_time && self.size_in > 0 {
+                    let size = pkt.len();
+                    let time_add = (self.size_out * ((self.last_time - self.first_time) as usize) / self.size_in) as u32;
+                    self.size_out += size / self.interleave.get_frame_size();
+                    time_add
+                } else {
+                    0
+                };
+            Some((pkt, self.first_time + time_add, first))
+        } else {
+            None
+        }
+    }
     fn flush(&mut self) { self.interleave.flush() }
     fn finish(&mut self, bw: &mut ByteWriter) -> MuxerResult<()> {
         let cur_pos = bw.tell();
@@ -617,6 +649,10 @@ pub fn create_audio_stream(stream: &NAStream) -> MuxerResult<Box<dyn RMStreamWri
             interleave: create_interleaver(ileave, fourcc, ainfo)?,
             header_pos: 0,
             data_size:  0,
+            first_time: 0,
+            last_time:  0,
+            size_in:    0,
+            size_out:   0,
         }))
     } else {
         Err(MuxerError::UnsupportedFormat)
diff --git a/nihav-realmedia/src/muxers/rmvb/mod.rs b/nihav-realmedia/src/muxers/rmvb/mod.rs
index 32fbbb9..7d8a7fd 100644
--- a/nihav-realmedia/src/muxers/rmvb/mod.rs
+++ b/nihav-realmedia/src/muxers/rmvb/mod.rs
@@ -34,8 +34,8 @@ impl<'a> RMWriterHelper for ByteWriter<'a> {
 
 pub trait RMStreamWriter {
     fn write_header(&mut self, bw: &mut ByteWriter, astream: &NAStream) -> MuxerResult<()>;
-    fn queue_packet(&mut self, pkt: NAPacket) -> bool;
-    fn get_packet(&mut self) -> Option<(Vec<u8>, bool)>;
+    fn queue_packet(&mut self, pkt: NAPacket, ms: u32) -> bool;
+    fn get_packet(&mut self) -> Option<(Vec<u8>, u32, bool)>;
     fn flush(&mut self);
     fn finish(&mut self, bw: &mut ByteWriter) -> MuxerResult<()>;
 }
@@ -127,21 +127,21 @@ impl RMStream {
             self.time = self.time.max(ms);
             self.cur_time = ms;
         }
-        self.keyframe = pkt.keyframe;
-        self.packetiser.queue_packet(pkt);
+        self.keyframe = pkt.keyframe || self.audio;
+        self.packetiser.queue_packet(pkt, self.cur_time);
         self.write_packets(bw, pkt_no)
     }
     fn write_packets(&mut self, bw: &mut ByteWriter, pkt_no: &mut u32) -> MuxerResult<()> {
-        while let Some((data, first)) = self.packetiser.get_packet() {
+        while let Some((data, ts, first)) = self.packetiser.get_packet() {
             validate!(data.len() < 65000);
             if self.keyframe && first {
-                self.index.push(IndexEntry{ time: self.cur_time, pos: bw.tell(), pkt_no: *pkt_no });
+                self.index.push(IndexEntry{ time: ts, pos: bw.tell(), pkt_no: *pkt_no });
             }
             let is_keyframe = self.keyframe && (!self.audio || first);
             bw.write_u16be(0)?; //version;
             bw.write_u16be((data.len() + 12) as u16)?;
             bw.write_u16be(self.stream_id)?;
-            bw.write_u32be(self.cur_time)?;
+            bw.write_u32be(ts)?;
             bw.write_byte(0)?; //packet group
             bw.write_byte(if is_keyframe { 0x2 } else { 0x0 })?;
             bw.write_buf(&data)?;
@@ -383,8 +383,8 @@ impl<'a> MuxCore<'a> for RAMuxer<'a> {
     }
     fn mux_frame(&mut self, _strmgr: &StreamManager, pkt: NAPacket) -> MuxerResult<()> {
         if let Some(ref mut sw) = self.sw {
-            sw.queue_packet(pkt);
-            while let Some((data, _)) = sw.get_packet() {
+            sw.queue_packet(pkt, 0);
+            while let Some((data, _, _)) = sw.get_packet() {
                 self.bw.write_buf(&data)?;
             }
             Ok(())
@@ -570,7 +570,7 @@ mod test {
             };
         test_remuxing(&dec_config, &enc_config);*/
         test_remuxing_md5(&dec_config, "realmedia", &mux_reg,
-                          [0xe38b36c0, 0x1aedef10, 0x4e418ac4, 0x4ff57f6c]);
+                          [0xcfa1a27b, 0x78314fa7, 0xeb90c31c, 0x7eafeaa8]);
     }
     #[test]
     fn test_rm_muxer_ralf() {
diff --git a/nihav-realmedia/src/muxers/rmvb/videostream.rs b/nihav-realmedia/src/muxers/rmvb/videostream.rs
index 708874d..d1aecce 100644
--- a/nihav-realmedia/src/muxers/rmvb/videostream.rs
+++ b/nihav-realmedia/src/muxers/rmvb/videostream.rs
@@ -17,10 +17,10 @@ impl RMStreamWriter for DummyStreamWriter {
     fn write_header(&mut self, _bw: &mut ByteWriter, _stream: &NAStream) -> MuxerResult<()> {
         Ok(())
     }
-    fn queue_packet(&mut self, _pkt: NAPacket) -> bool {
+    fn queue_packet(&mut self, _pkt: NAPacket, _ms: u32) -> bool {
         true
     }
-    fn get_packet(&mut self) -> Option<(Vec<u8>, bool)> {
+    fn get_packet(&mut self) -> Option<(Vec<u8>, u32, bool)> {
         None
     }
     fn flush(&mut self) { }
@@ -35,6 +35,7 @@ struct VideoStreamWriter {
     nslices:    usize,
     cur_slice:  usize,
     seq_no:     u8,
+    time:       u32,
 }
 
 impl RMStreamWriter for VideoStreamWriter {
@@ -75,7 +76,7 @@ impl RMStreamWriter for VideoStreamWriter {
         bw.seek(SeekFrom::Start(end))?;
         Ok(())
     }
-    fn queue_packet(&mut self, pkt: NAPacket) -> bool {
+    fn queue_packet(&mut self, pkt: NAPacket, ms: u32) -> bool {
         if self.nslices == 0 {
             let src = pkt.get_buffer();
             let nslices = usize::from(src[0]) + 1;
@@ -84,13 +85,14 @@ impl RMStreamWriter for VideoStreamWriter {
                 self.cur_slice = 0;
                 self.buf.resize(src.len(), 0);
                 self.buf.copy_from_slice(&src);
+                self.time = ms;
             }
             true
         } else {
             false
         }
     }
-    fn get_packet(&mut self) -> Option<(Vec<u8>, bool)> {
+    fn get_packet(&mut self) -> Option<(Vec<u8>, u32, bool)> {
         if self.cur_slice < self.nslices {
             let first = self.cur_slice == 0;
             let hdr_size = self.nslices * 8 + 1;
@@ -138,7 +140,7 @@ impl RMStreamWriter for VideoStreamWriter {
                 self.cur_slice = 0;
                 self.seq_no = self.seq_no.wrapping_add(1);
             }
-            Some((ret, first))
+            Some((ret, self.time, first))
         } else {
             None
         }
@@ -161,6 +163,7 @@ pub fn create_video_stream(stream: &NAStream) -> MuxerResult<Box<dyn RMStreamWri
                     nslices:    0,
                     cur_slice:  0,
                     seq_no:     0,
+                    time:       0,
                 }));
         }
     }