X-Git-Url: https://git.nihav.org/?p=nihav-tool.git;a=blobdiff_plain;f=src%2Fdemux.rs;h=78743bead1d75ec6c247f71b39355f587c82cbb2;hp=4ef1f158e7f7cf4b193d03938544651830bbbc35;hb=1c5de8c5a8b86e254e9941b49d476a861abb7d64;hpb=55d3e99fc4841c63564e337e5b7ec40559f1d19a diff --git a/src/demux.rs b/src/demux.rs index 4ef1f15..78743be 100644 --- a/src/demux.rs +++ b/src/demux.rs @@ -26,11 +26,75 @@ impl FullRegister { } } +pub struct RawStreamCtx<'a> { + stream: NAStreamRef, + pkt: Box, + br: &'a mut ByteReader<'a>, + pts: u64, + seek: SeekIndex, +} + +impl<'a> RawStreamCtx<'a> { + fn new(stream: NAStreamRef, packetiser: Box, br: &'a mut ByteReader<'a>) -> Self { + let mut seek = SeekIndex::new(); + seek.add_stream(0); + Self { stream, pkt: packetiser, br, pts: 0, seek } + } + fn account_for_packet(&mut self, packet: &mut NAPacket) { + let pos = self.br.tell() - (self.pkt.bytes_left() as u64); + if packet.get_pts().is_none() && packet.get_duration().is_some() { + packet.ts.pts = Some(self.pts); + } + if packet.is_keyframe() { + let pts = packet.get_pts().unwrap_or(self.pts); + let time = NATimeInfo::ts_to_time(pts, 1000, self.stream.tb_num, self.stream.tb_den); + let in_range = if let Some(last) = self.seek.seek_info[0].entries.last() { + last.pts >= pts + } else { + false + }; + if !in_range { + self.seek.add_entry(0, SeekEntry { time, pts, pos }); + } + } + self.pts += packet.get_duration().unwrap_or(0); + } + fn get_frame(&mut self) -> DemuxerResult { + let mut buf = [0; 1048576]; + loop { + match self.pkt.get_packet(self.stream.clone()) { + Ok(Some(mut packet)) => { + self.account_for_packet(&mut packet); + return Ok(packet); + }, + Ok(None) => {}, + Err(DecoderError::ShortData) => {}, + _ => return Err(DemuxerError::InvalidData), + }; + match self.br.read_buf_some(&mut buf) { + Ok(size) => { + self.pkt.add_data(&buf[..size]); + }, + Err(_) => { + match self.pkt.get_packet(self.stream.clone()) { + Ok(Some(mut packet)) => { + self.account_for_packet(&mut packet); + return Ok(packet); + }, + Ok(None) | Err(DecoderError::ShortData) => return Err(DemuxerError::EOF), + _ => return Err(DemuxerError::InvalidData), + }; + }, + }; + } + } +} + pub enum DemuxerObject<'a> { None, Normal(Demuxer<'a>), Raw(RawDemuxer<'a>, Vec>>, bool), - RawStream(Box, NAStreamRef, &'a mut ByteReader<'a>), + RawStream(RawStreamCtx<'a>), } impl<'a> DemuxerObject<'a> { @@ -101,7 +165,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 } @@ -116,6 +180,9 @@ impl<'a> DemuxerObject<'a> { match *self { DemuxerObject::Normal(ref dmx) => dmx.get_duration(), DemuxerObject::Raw(ref dmx, _, _) => dmx.get_duration(), + DemuxerObject::RawStream(ref ctx) => { + NATimeInfo::ts_to_time(ctx.stream.duration, 1000, ctx.stream.tb_num, ctx.stream.tb_den) + }, _ => 0, } } @@ -124,14 +191,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 { 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,29 +243,7 @@ impl<'a> DemuxerObject<'a> { } } }, - DemuxerObject::RawStream(ref mut pkt, ref stream, ref mut br) => { - let mut buf = [0; 1048576]; - loop { - match pkt.get_packet(stream.clone()) { - Ok(Some(packet)) => return Ok(packet), - Ok(None) => {}, - Err(DecoderError::ShortData) => {}, - _ => return Err(DemuxerError::InvalidData), - }; - match br.read_buf_some(&mut buf) { - Ok(size) => { - pkt.add_data(&buf[..size]); - }, - Err(_) => { - match pkt.get_packet(stream.clone()) { - Ok(Some(packet)) => return Ok(packet), - Ok(None) | Err(DecoderError::ShortData) => return Err(DemuxerError::EOF), - _ => return Err(DemuxerError::InvalidData), - }; - }, - }; - } - }, + DemuxerObject::RawStream(ref mut ctx) => ctx.get_frame(), _ => unreachable!(), } } @@ -206,6 +251,45 @@ impl<'a> DemuxerObject<'a> { match *self { DemuxerObject::Normal(ref mut dmx) => dmx.seek(seek_time), DemuxerObject::Raw(ref mut dmx, _, _) => dmx.seek(seek_time), + DemuxerObject::RawStream(ref mut ctx) => { + if seek_time == NATimePoint::None { + return Err(DemuxerError::SeekError); + } + if let Some(last) = ctx.seek.seek_info[0].entries.last() { + let in_index = match seek_time { + NATimePoint::None => unreachable!(), + NATimePoint::PTS(pts) => last.pts >= pts, + NATimePoint::Milliseconds(ms) => last.time >= ms, + }; + if in_index { + if let Some(result) = ctx.seek.find_pos(seek_time) { + ctx.br.seek(SeekFrom::Start(result.pos))?; + ctx.pts = result.pts; + ctx.pkt.reset(); + return Ok(()); + } + } + } + if let Some(last) = ctx.seek.seek_info[0].entries.last() { + ctx.br.seek(SeekFrom::Start(last.pos))?; + ctx.pts = last.pts; + ctx.pkt.reset(); + } + let mut key_pts = 0; + while let Ok(pkt) = ctx.get_frame() { + if !pkt.ts.less_than(seek_time) && !pkt.ts.equal(seek_time) { + break; + } + if pkt.is_keyframe() { + key_pts = pkt.get_pts().unwrap_or(0); + } + } + let result = ctx.seek.find_pos(NATimePoint::PTS(key_pts)).unwrap(); + ctx.br.seek(SeekFrom::Start(result.pos))?; + ctx.pts = result.pts; + ctx.pkt.reset(); + Ok(()) + }, _ => Err(DemuxerError::NotImplemented), } }