1#![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#[derive(Debug)]
27pub struct Muxer {
28 program_number: u16,
29 pmt_pid: u16,
30 streams: SmallVec<[(ElementaryStream, u8); 4]>, pat_cc: u8,
32 pmt_cc: u8,
33}
34
35impl Muxer {
36 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 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 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); payload.extend_from_slice(section);
111 payload
112}
113
114#[cfg(test)]
115#[path = "mux_tests.rs"]
116mod tests;