Skip to main content

mpeg_ts_core/
mux.rs

1//! MPEG-TS mux: PAT/PMT + per-PID PES packetization over 188-byte TS packets.
2
3#![forbid(unsafe_code)]
4
5use smallvec::SmallVec;
6
7use crate::error::Error;
8use crate::packet::write_ts_packets;
9use crate::pes::{build_pes_header, stream_id_for};
10use crate::psi::{build_pat_section, build_pmt_section};
11use crate::types::{ElementaryStream, StreamType};
12
13const PAT_PID: u16 = 0;
14
15const fn validate_pid(pid: u16) -> Result<(), Error> {
16    if pid > 0x1FFF || pid == 0 || pid == 1 {
17        Err(Error::InvalidPid(pid))
18    } else {
19        Ok(())
20    }
21}
22
23/// Writes a single-program MPEG-TS stream: `write_pat_pmt` (call at the start,
24/// and periodically — real players expect PAT/PMT to repeat) plus
25/// `write_access_unit` per elementary-stream access unit.
26#[derive(Debug)]
27pub struct Muxer {
28    program_number: u16,
29    pmt_pid: u16,
30    streams: SmallVec<[(ElementaryStream, u8); 4]>, // (stream, continuity_counter)
31    pat_cc: u8,
32    pmt_cc: u8,
33}
34
35impl Muxer {
36    /// Start a mux session for one program. `pmt_pid` and every stream's `pid`
37    /// must be in `2..=0x1FFF` (`0`/`1` are reserved for PAT/CAT).
38    pub fn new(
39        program_number: u16,
40        pmt_pid: u16,
41        streams: &[ElementaryStream],
42    ) -> Result<Self, Error> {
43        validate_pid(pmt_pid)?;
44        for stream in streams {
45            validate_pid(stream.pid)?;
46        }
47        Ok(Self {
48            program_number,
49            pmt_pid,
50            streams: streams.iter().map(|s| (*s, 0u8)).collect(),
51            pat_cc: 0,
52            pmt_cc: 0,
53        })
54    }
55
56    /// Write one PAT packet + one PMT packet describing this program's streams.
57    pub fn write_pat_pmt(&mut self, out: &mut Vec<u8>) {
58        let pat = build_pat_section(1, self.program_number, self.pmt_pid);
59        write_ts_packets(
60            out,
61            PAT_PID,
62            &mut self.pat_cc,
63            &pusi_payload(&pat),
64            true,
65            false,
66        );
67
68        let elementary: SmallVec<[ElementaryStream; 4]> =
69            self.streams.iter().map(|(s, _)| *s).collect();
70        let pmt = build_pmt_section(self.program_number, &elementary);
71        write_ts_packets(
72            out,
73            self.pmt_pid,
74            &mut self.pmt_cc,
75            &pusi_payload(&pmt),
76            true,
77            false,
78        );
79    }
80
81    /// Packetize one access unit (already-encoded elementary-stream data, e.g.
82    /// Annex-B H.264 or an ADTS AAC frame) into a PES packet and split it across
83    /// TS packets on `pid`.
84    pub fn write_access_unit(
85        &mut self,
86        pid: u16,
87        data: &[u8],
88        pts_90k: u64,
89        dts_90k: Option<u64>,
90        random_access: bool,
91        out: &mut Vec<u8>,
92    ) -> Result<(), Error> {
93        let (stream, cc) = self
94            .streams
95            .iter_mut()
96            .find(|(s, _)| s.pid == pid)
97            .ok_or(Error::UnknownPid(pid))?;
98        let is_video = matches!(stream.stream_type, StreamType::H264 | StreamType::Hevc);
99
100        let mut payload = build_pes_header(stream_id_for(is_video), data.len(), pts_90k, dts_90k);
101        payload.extend_from_slice(data);
102        write_ts_packets(out, pid, cc, &payload, true, random_access);
103        Ok(())
104    }
105}
106
107fn pusi_payload(section: &[u8]) -> Vec<u8> {
108    let mut payload = Vec::with_capacity(1 + section.len());
109    payload.push(0); // pointer_field = 0 (section starts immediately)
110    payload.extend_from_slice(section);
111    payload
112}
113
114#[cfg(test)]
115#[path = "mux_tests.rs"]
116mod tests;