mpeg_ts_core/
demux.rs

1//! MPEG-TS demux: PAT/PMT tracking + per-PID PES reassembly from 188-byte TS
2//! packets.
3//!
4//! PSI (PAT/PMT) sections spanning more than one TS packet are not
5//! reassembled (v1 scope, crate-local ADR-0001) — this crate's own `Muxer`
6//! never produces one (a single program with a handful of streams always fits
7//! in one packet), but an arbitrary third-party multi-program stream with a
8//! very large PMT would not parse correctly here.
9
10#![forbid(unsafe_code)]
11
12use std::collections::{HashMap, VecDeque};
13
14use bytes::Bytes;
15use smallvec::SmallVec;
16
17use crate::error::Error;
18use crate::packet::{PACKET_LEN, parse_ts_packet};
19use crate::pes::parse_pes_header;
20use crate::psi::{parse_pat_section, parse_pmt_section};
21use crate::types::{AccessUnit, ElementaryStream};
22
23const PAT_PID: u16 = 0;
24
25#[derive(Debug, Default)]
26struct PesAccumulator {
27    data: Vec<u8>,
28    random_access: bool,
29}
30
31/// Reads TS packets from pushed byte chunks, tracks PAT/PMT, and reassembles
32/// per-PID PES packets into [`AccessUnit`]s.
33#[derive(Debug, Default)]
34pub struct Demuxer {
35    buf: Vec<u8>,
36    pmt_pid: Option<u16>,
37    streams: SmallVec<[ElementaryStream; 4]>,
38    accumulators: HashMap<u16, PesAccumulator>,
39    pending: VecDeque<AccessUnit>,
40}
41
42impl Demuxer {
43    /// New, empty demux session.
44    #[must_use]
45    pub fn new() -> Self {
46        Self::default()
47    }
48
49    /// Append incoming bytes (need not be 188-byte aligned across calls).
50    pub fn push_bytes(&mut self, data: &[u8]) {
51        self.buf.extend_from_slice(data);
52    }
53
54    /// Elementary streams from the most recently parsed PMT (empty until then).
55    #[must_use]
56    pub fn streams(&self) -> &[ElementaryStream] {
57        &self.streams
58    }
59
60    /// Pop the next fully reassembled access unit, or `Ok(None)` if not enough
61    /// bytes are buffered yet.
62    ///
63    /// An access unit is only confirmed complete once the *next* PES packet on
64    /// the same PID starts (or [`Demuxer::finish`] is called) — this is
65    /// inherent to how PES packetization signals its own boundaries, not a
66    /// limitation specific to this crate.
67    pub fn poll_access_unit(&mut self) -> Result<Option<AccessUnit>, Error> {
68        loop {
69            if let Some(unit) = self.pending.pop_front() {
70                return Ok(Some(unit));
71            }
72            if !self.consume_one_packet()? {
73                return Ok(None);
74            }
75        }
76    }
77
78    /// Force-emit whatever is still accumulating per PID as a final access
79    /// unit each — call once at the end of a stream to avoid losing the very
80    /// last access unit per PID (see [`Demuxer::poll_access_unit`]'s doc).
81    pub fn finish(&mut self) -> Vec<AccessUnit> {
82        let mut out = Vec::new();
83        for (pid, acc) in self.accumulators.drain() {
84            if acc.data.is_empty() {
85                continue;
86            }
87            if let Ok(unit) = finish_pes(pid, &acc) {
88                out.push(unit);
89            }
90        }
91        out
92    }
93
94    fn consume_one_packet(&mut self) -> Result<bool, Error> {
95        if self.buf.len() < PACKET_LEN {
96            return Ok(false);
97        }
98        let packet_owned = self.buf[..PACKET_LEN].to_vec();
99        let parsed = parse_ts_packet(&packet_owned)?;
100        let pid = parsed.pid;
101
102        if pid == PAT_PID {
103            if parsed.pusi {
104                let pat = parse_pat_section(parsed.payload)?;
105                self.pmt_pid = Some(pat.pmt_pid);
106            }
107        } else if Some(pid) == self.pmt_pid {
108            if parsed.pusi {
109                let pmt = parse_pmt_section(parsed.payload)?;
110                self.streams = pmt.streams;
111            }
112        } else if self.streams.iter().any(|s| s.pid == pid) {
113            self.feed_pes(pid, parsed.pusi, parsed.random_access, parsed.payload)?;
114        }
115
116        self.buf.drain(0..PACKET_LEN);
117        Ok(true)
118    }
119
120    fn feed_pes(
121        &mut self,
122        pid: u16,
123        pusi: bool,
124        random_access: bool,
125        payload: &[u8],
126    ) -> Result<(), Error> {
127        if pusi {
128            if let Some(prev) = self.accumulators.get(&pid) {
129                if !prev.data.is_empty() {
130                    let unit = finish_pes(pid, prev)?;
131                    self.pending.push_back(unit);
132                }
133            }
134            self.accumulators.insert(
135                pid,
136                PesAccumulator {
137                    data: payload.to_vec(),
138                    random_access,
139                },
140            );
141        } else if let Some(acc) = self.accumulators.get_mut(&pid) {
142            acc.data.extend_from_slice(payload);
143        }
144        Ok(())
145    }
146}
147
148fn finish_pes(pid: u16, acc: &PesAccumulator) -> Result<AccessUnit, Error> {
149    let header = parse_pes_header(&acc.data)?;
150    Ok(AccessUnit {
151        pid,
152        data: Bytes::copy_from_slice(&acc.data[header.header_len..]),
153        pts_90k: header.pts_90k,
154        dts_90k: header.dts_90k,
155        random_access: acc.random_access,
156    })
157}
158
159#[cfg(test)]
160#[path = "demux_tests.rs"]
161mod tests;
mpeg_ts_core/demux.rs

mpeg_ts_core/
demux.rs
