Skip to main content

iso_bmff/mux/
mod.rs

1//! Typestate MP4 muxer (fMP4 primary) — session only; boxes live in `isobmff`.
2
3#![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/// Track registration phase — call [`Muxer::begin`] before packets.
18#[derive(Debug, Clone, Copy, Default)]
19pub struct Open;
20
21/// Streaming phase — packet/byte API.
22#[derive(Debug, Clone, Copy, Default)]
23pub struct Live;
24
25/// Default samples per fragment.
26pub 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    /// Per-sample decode timestamps (media timescale) — durations are derived
38    /// from consecutive `dts` deltas at flush time, see [`Muxer::push_packet`].
39    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/// Sans-IO fragmented MP4 muxer.
48#[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    /// Empty muxer in track-registration state.
62    #[must_use]
63    pub fn new() -> Self {
64        Self::with_fragment_batch(DEFAULT_FRAGMENT_BATCH)
65    }
66
67    /// Muxer with custom fragment batch size.
68    #[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    /// Registered tracks so far.
83    #[must_use]
84    pub fn tracks(&self) -> &[Track] {
85        &self.tracks
86    }
87
88    /// Register a track. `id` must be unique.
89    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    /// Lock tracks and enter the live streaming state.
100    #[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    /// Registered tracks (may gain extradata from Annex-B / ADTS).
123    #[must_use]
124    pub fn tracks(&self) -> &[Track] {
125        &self.tracks
126    }
127
128    /// Push a compressed Sample (H.264 Annex-B auto-converted to AVCC).
129    ///
130    /// Sample durations are computed from consecutive `dts` deltas inside each
131    /// fragment (standard muxer convention), so `Sample::duration` is optional:
132    /// it is only trusted for the **last** sample of a fragment, and when it is
133    /// zero the last sample's duration is estimated from the previous sample's
134    /// delta (a lone-sample fragment defaults to one media tick). `dts` must be
135    /// monotonically non-decreasing per track; out-of-order `dts` degrades to
136    /// a 1-tick duration.
137    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                // clone: passthrough codecs keep caller's Bytes; no in-place transform
157                (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    /// Flush all pending fragments.
243    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    /// Append available output bytes into `out`.
253    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        // Sample durations are the dts deltas between consecutive samples
276        // (standard muxer convention). The caller-provided duration is only
277        // consulted for the LAST sample of the fragment; when it is zero we
278        // estimate it from the previous sample's delta, and a lone-sample
279        // fragment defaults to one media tick.
280        let n = pending.durations.len();
281        for i in 0..n.saturating_sub(1) {
282            // Non-monotonic dts (caller out of order) clamps to a 1-tick
283            // duration instead of a zero/u32::MAX sample: `saturating_sub`
284            // saturates at i64::MIN, so clamp the delta explicitly.
285            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;