use crate::ts::payload::{Bytes, Null, Pat, Pes, Pmt};
use crate::ts::{AdaptationField, Pid, TsHeader, TsPacket, TsPayload};
use crate::{ErrorKind, Result};
use std::collections::HashMap;
use std::io::Read;
pub trait ReadTsPacket {
fn read_ts_packet(&mut self) -> Result<Option<TsPacket>>;
}
#[derive(Debug)]
pub struct TsPacketReader<R> {
stream: R,
pids: HashMap<Pid, PidKind>,
}
impl<R: Read> TsPacketReader<R> {
pub fn new(stream: R) -> Self {
TsPacketReader {
stream,
pids: HashMap::new(),
}
}
pub fn stream(&self) -> &R {
&self.stream
}
pub fn into_stream(self) -> R {
self.stream
}
}
impl<R: Read> ReadTsPacket for TsPacketReader<R> {
fn read_ts_packet(&mut self) -> Result<Option<TsPacket>> {
let mut reader = self.stream.by_ref().take(TsPacket::SIZE as u64);
let mut peek = [0; 1];
let eos = track_io!(reader.read(&mut peek))? == 0;
if eos {
return Ok(None);
}
let (header, adaptation_field_control, payload_unit_start_indicator) =
track!(TsHeader::read_from(peek.chain(&mut reader)))?;
let adaptation_field = if adaptation_field_control.has_adaptation_field() {
track!(AdaptationField::read_from(&mut reader))?
} else {
None
};
let payload = if adaptation_field_control.has_payload() {
let payload = match header.pid.as_u16() {
Pid::PAT => {
let pat = track!(Pat::read_from(&mut reader))?;
for pa in &pat.table {
self.pids.insert(pa.program_map_pid, PidKind::Pmt);
}
TsPayload::Pat(pat)
}
Pid::NULL => {
let null = track!(Null::read_from(&mut reader))?;
TsPayload::Null(null)
}
0x01..=0x1F | 0x1FFB => {
let bytes = track!(Bytes::read_from(&mut reader))?;
TsPayload::Raw(bytes)
}
_ => {
let kind = track_assert_some!(
self.pids.get(&header.pid).cloned(),
ErrorKind::InvalidInput,
"Unknown PID: header={:?}",
header
);
match kind {
PidKind::Pmt => {
let pmt = track!(Pmt::read_from(&mut reader))?;
for es in &pmt.es_info {
self.pids.insert(es.elementary_pid, PidKind::Pes);
}
TsPayload::Pmt(pmt)
}
PidKind::Pes => {
if payload_unit_start_indicator {
let pes = track!(Pes::read_from(&mut reader))?;
TsPayload::Pes(pes)
} else {
let bytes = track!(Bytes::read_from(&mut reader))?;
TsPayload::Raw(bytes)
}
}
}
}
};
Some(payload)
} else {
None
};
track_assert_eq!(reader.limit(), 0, ErrorKind::InvalidInput);
Ok(Some(TsPacket {
header,
adaptation_field,
payload,
}))
}
}
#[derive(Debug, Clone)]
enum PidKind {
Pmt,
Pes,
}