1#![forbid(unsafe_code)]
4
5#[cfg(feature = "audio")]
6use crate::bitstream::strip_adts;
7#[cfg(feature = "video")]
8use crate::bitstream::to_avcc;
9use crate::codec_features::check_codec;
10use crate::error::Error;
11use crate::isobmff::{write_fragment, write_ftyp, write_moov};
12use crate::types::{Codec, Sample, Track};
13use crate::{INLINE_SAMPLES, INLINE_TRACKS};
14use smallvec::SmallVec;
15use std::marker::PhantomData;
16
17#[derive(Debug, Clone, Copy, Default)]
19pub struct Open;
20
21#[derive(Debug, Clone, Copy, Default)]
23pub struct Live;
24
25pub const DEFAULT_FRAGMENT_BATCH: usize = 30;
27
28const _: () = assert!(
29 DEFAULT_FRAGMENT_BATCH <= INLINE_SAMPLES,
30 "INLINE_SAMPLES must cover the default fragment batch"
31);
32
33#[derive(Debug)]
34struct Pending {
35 track_id: u32,
36 base_dts: u64,
37 dts: SmallVec<[i64; INLINE_SAMPLES]>,
40 durations: SmallVec<[u32; INLINE_SAMPLES]>,
41 sizes: SmallVec<[u32; INLINE_SAMPLES]>,
42 flags: SmallVec<[u32; INLINE_SAMPLES]>,
43 ctos: SmallVec<[i32; INLINE_SAMPLES]>,
44 payload: Vec<u8>,
45}
46
47#[derive(Debug)]
49pub struct Muxer<S = Open> {
50 tracks: SmallVec<[Track; INLINE_TRACKS]>,
51 output: Vec<u8>,
52 output_consumed: usize,
53 header_written: bool,
54 sequence: u32,
55 batch: usize,
56 pending: SmallVec<[Pending; INLINE_TRACKS]>,
57 _state: PhantomData<S>,
58}
59
60impl Muxer<Open> {
61 #[must_use]
63 pub fn new() -> Self {
64 Self::with_fragment_batch(DEFAULT_FRAGMENT_BATCH)
65 }
66
67 #[must_use]
69 pub fn with_fragment_batch(batch: usize) -> Self {
70 Self {
71 tracks: SmallVec::new(),
72 output: Vec::with_capacity(64 * 1024),
73 output_consumed: 0,
74 header_written: false,
75 sequence: 0,
76 batch: batch.max(1),
77 pending: SmallVec::new(),
78 _state: PhantomData,
79 }
80 }
81
82 #[must_use]
84 pub fn tracks(&self) -> &[Track] {
85 &self.tracks
86 }
87
88 pub fn add_track(&mut self, track: Track) -> Result<u32, Error> {
90 check_codec(track.codec)?;
91 if self.tracks.iter().any(|t| t.id == track.id) {
92 return Err(Error::InvalidTrack);
93 }
94 let id = track.id;
95 self.tracks.push(track);
96 Ok(id)
97 }
98
99 #[must_use]
101 pub fn begin(self) -> Muxer<Live> {
102 Muxer {
103 tracks: self.tracks,
104 output: self.output,
105 output_consumed: self.output_consumed,
106 header_written: self.header_written,
107 sequence: self.sequence,
108 batch: self.batch,
109 pending: self.pending,
110 _state: PhantomData,
111 }
112 }
113}
114
115impl Default for Muxer<Open> {
116 fn default() -> Self {
117 Self::new()
118 }
119}
120
121impl Muxer<Live> {
122 #[must_use]
124 pub fn tracks(&self) -> &[Track] {
125 &self.tracks
126 }
127
128 pub fn push_packet(&mut self, sample: &Sample) -> Result<(), Error> {
138 let idx = self
139 .tracks
140 .iter()
141 .position(|t| t.id == sample.stream_id)
142 .ok_or(Error::InvalidPacket)?;
143
144 let track_codec = self.tracks[idx].codec;
145 check_codec(track_codec)?;
146
147 let (payload, extra) = match track_codec {
148 #[cfg(feature = "video")]
149 Codec::H264 => {
150 let o = to_avcc(&sample.payload);
151 (o.payload, o.avcc)
152 }
153 #[cfg(feature = "audio")]
154 Codec::Aac => strip_adts(&sample.payload),
155 _ => {
156 (sample.payload.clone(), None)
158 }
159 };
160 if let Some(e) = extra {
161 if self.tracks[idx].extra_data.is_empty() {
162 self.tracks[idx].extra_data = e;
163 }
164 }
165
166 if !self.header_written {
167 write_ftyp(&mut self.output, &self.tracks);
168 write_moov(&mut self.output, &self.tracks);
169 self.header_written = true;
170 }
171
172 let isobmff_id = sample.stream_id.saturating_add(1);
173 let size = u32::try_from(payload.len()).unwrap_or(u32::MAX);
174 let dur = u32::try_from(sample.duration.min(u64::from(u32::MAX))).unwrap_or(u32::MAX);
175 let flags = if sample.is_keyframe {
176 0x0200_0000
177 } else {
178 0x0101_0000
179 };
180 let cto = i32::try_from(
181 sample
182 .pts
183 .saturating_sub(sample.dts)
184 .clamp(i64::from(i32::MIN), i64::from(i32::MAX)),
185 )
186 .unwrap_or(0);
187
188 let batch = self.batch;
189 let key_flush = sample.is_keyframe
190 && self.tracks[idx].codec != Codec::Aac
191 && self
192 .pending
193 .iter()
194 .any(|p| p.track_id == isobmff_id && !p.durations.is_empty());
195 if key_flush
196 || self
197 .pending
198 .iter()
199 .any(|p| p.track_id == isobmff_id && p.durations.len() >= batch)
200 {
201 self.flush_track(isobmff_id);
202 }
203
204 let base_dts = u64::try_from(sample.dts.max(0)).unwrap_or(0);
205 if let Some(p) = self.pending.iter_mut().find(|p| p.track_id == isobmff_id) {
206 p.dts.push(sample.dts);
207 p.durations.push(dur);
208 p.sizes.push(size);
209 p.flags.push(flags);
210 p.ctos.push(cto);
211 p.payload.extend_from_slice(&payload);
212 } else {
213 let mut pending = Pending {
214 track_id: isobmff_id,
215 base_dts,
216 dts: SmallVec::with_capacity(batch),
217 durations: SmallVec::with_capacity(batch),
218 sizes: SmallVec::with_capacity(batch),
219 flags: SmallVec::with_capacity(batch),
220 ctos: SmallVec::with_capacity(batch),
221 payload: Vec::with_capacity(payload.len().saturating_mul(batch)),
222 };
223 pending.dts.push(sample.dts);
224 pending.durations.push(dur);
225 pending.sizes.push(size);
226 pending.flags.push(flags);
227 pending.ctos.push(cto);
228 pending.payload.extend_from_slice(&payload);
229 self.pending.push(pending);
230 }
231
232 if self
233 .pending
234 .iter()
235 .any(|p| p.track_id == isobmff_id && p.durations.len() >= batch)
236 {
237 self.flush_track(isobmff_id);
238 }
239 Ok(())
240 }
241
242 pub fn flush(&mut self) {
244 let mut ids: SmallVec<[u32; INLINE_TRACKS]> =
245 self.pending.iter().map(|p| p.track_id).collect();
246 ids.sort_unstable();
247 for id in ids {
248 self.flush_track(id);
249 }
250 }
251
252 pub fn poll_bytes(&mut self, out: &mut Vec<u8>) -> usize {
254 let available = self.output.len().saturating_sub(self.output_consumed);
255 if available == 0 {
256 return 0;
257 }
258 out.extend_from_slice(&self.output[self.output_consumed..]);
259 self.output_consumed = self.output.len();
260 if self.output_consumed >= 64 * 1024 {
261 self.output.drain(..self.output_consumed);
262 self.output_consumed = 0;
263 }
264 available
265 }
266
267 fn flush_track(&mut self, track_id: u32) {
268 let Some(pos) = self.pending.iter().position(|p| p.track_id == track_id) else {
269 return;
270 };
271 let mut pending = self.pending.swap_remove(pos);
272 if pending.durations.is_empty() {
273 return;
274 }
275 let n = pending.durations.len();
281 for i in 0..n.saturating_sub(1) {
282 let delta = pending.dts[i + 1].saturating_sub(pending.dts[i]).max(1);
286 pending.durations[i] = u32::try_from(delta).unwrap_or(u32::MAX);
287 }
288 if pending.durations[n - 1] == 0 {
289 pending.durations[n - 1] = if n >= 2 { pending.durations[n - 2] } else { 1 };
290 }
291 self.sequence = self.sequence.saturating_add(1);
292 write_fragment(
293 &mut self.output,
294 self.sequence,
295 pending.track_id,
296 pending.base_dts,
297 &pending.durations,
298 &pending.sizes,
299 &pending.flags,
300 &pending.ctos,
301 &pending.payload,
302 );
303 }
304}
305
306#[cfg(test)]
307#[path = "mux_tests.rs"]
308mod tests;