From e01dfa5a9e674d6abe3a235de0e1c11f027af866 Mon Sep 17 00:00:00 2001
From: Kostya Shishkov <kostya.shishkov@gmail.com>
Date: Thu, 18 Nov 2021 13:53:21 +0100
Subject: [PATCH] improve raw stream handling

---
 src/demux.rs | 47 ++++++++++++++++++++++++++++++++++++-----------
 1 file changed, 36 insertions(+), 11 deletions(-)

diff --git a/src/demux.rs b/src/demux.rs
index 4ef1f15..084d9e7 100644
--- a/src/demux.rs
+++ b/src/demux.rs
@@ -26,11 +26,24 @@ impl FullRegister {
     }
 }
 
+pub struct RawStreamCtx<'a> {
+    stream: NAStreamRef,
+    pkt:    Box<dyn NAPacketiser + Send>,
+    br:     &'a mut ByteReader<'a>,
+    pts:    u64,
+}
+
+impl<'a> RawStreamCtx<'a> {
+    fn new(stream: NAStreamRef, packetiser: Box<dyn NAPacketiser + Send>, br: &'a mut ByteReader<'a>) -> Self {
+        Self { stream, pkt: packetiser, br, pts: 0 }
+    }
+}
+
 pub enum DemuxerObject<'a> {
     None,
     Normal(Demuxer<'a>),
     Raw(RawDemuxer<'a>, Vec<Option<Box<dyn NAPacketiser + Send>>>, bool),
-    RawStream(Box<dyn NAPacketiser + Send>, NAStreamRef, &'a mut ByteReader<'a>),
+    RawStream(RawStreamCtx<'a>),
 }
 
 impl<'a> DemuxerObject<'a> {
@@ -101,7 +114,7 @@ impl<'a> DemuxerObject<'a> {
             packetiser.add_data(&buf[..size]);
             let stream = packetiser.parse_stream(0).unwrap();
             packetiser.reset();
-            DemuxerObject::RawStream(packetiser, stream, br)
+            DemuxerObject::RawStream(RawStreamCtx::new(stream, packetiser, br))
         } else {
             DemuxerObject::None
         }
@@ -124,14 +137,14 @@ impl<'a> DemuxerObject<'a> {
             DemuxerObject::None => 0,
             DemuxerObject::Normal(ref dmx) => dmx.get_num_streams(),
             DemuxerObject::Raw(ref dmx, _, _) => dmx.get_num_streams(),
-            DemuxerObject::RawStream(_, _, _) => 1,
+            DemuxerObject::RawStream(_) => 1,
         }
     }
     pub fn get_stream(&self, idx: usize) -> Option<NAStreamRef> {
         match *self {
             DemuxerObject::Normal(ref dmx) => dmx.get_stream(idx),
             DemuxerObject::Raw(ref dmx, _, _) => dmx.get_stream(idx),
-            DemuxerObject::RawStream(_, ref stream, _) if idx == 0 => Some(stream.clone()),
+            DemuxerObject::RawStream(ref ctx) if idx == 0 => Some(ctx.stream.clone()),
             _ => None,
         }
     }
@@ -176,22 +189,34 @@ impl<'a> DemuxerObject<'a> {
                     }
                 }
             },
-            DemuxerObject::RawStream(ref mut pkt, ref stream, ref mut br) => {
+            DemuxerObject::RawStream(ref mut ctx) => {
                 let mut buf = [0; 1048576];
                 loop {
-                    match pkt.get_packet(stream.clone()) {
-                        Ok(Some(packet)) => return Ok(packet),
+                    match ctx.pkt.get_packet(ctx.stream.clone()) {
+                        Ok(Some(mut packet)) => {
+                            if packet.get_pts().is_none() && packet.get_duration().is_some() {
+                                packet.ts.pts = Some(ctx.pts);
+                            }
+                            ctx.pts += packet.get_duration().unwrap_or(0);
+                            return Ok(packet);
+                        },
                         Ok(None) => {},
                         Err(DecoderError::ShortData) => {},
                         _ => return Err(DemuxerError::InvalidData),
                     };
-                    match br.read_buf_some(&mut buf) {
+                    match ctx.br.read_buf_some(&mut buf) {
                         Ok(size) => {
-                            pkt.add_data(&buf[..size]);
+                            ctx.pkt.add_data(&buf[..size]);
                         },
                         Err(_) => {
-                            match pkt.get_packet(stream.clone()) {
-                                Ok(Some(packet)) => return Ok(packet),
+                            match ctx.pkt.get_packet(ctx.stream.clone()) {
+                                Ok(Some(mut packet)) => {
+                                    if packet.get_pts().is_none() && packet.get_duration().is_some() {
+                                        packet.ts.pts = Some(ctx.pts);
+                                    }
+                                    ctx.pts += packet.get_duration().unwrap_or(0);
+                                    return Ok(packet);
+                                },
                                 Ok(None) | Err(DecoderError::ShortData) => return Err(DemuxerError::EOF),
                                 _ => return Err(DemuxerError::InvalidData),
                             };
-- 
2.39.5