Skip to main content

pito_codecs/mp4/
writer.rs

1use crate::media::{Error, Packet, Result};
2use crate::mp4::boxes::BoxBuf;
3use crate::mp4::entry::SampleEntry;
4use crate::mp4::reader::Edit;
5use std::fs::File;
6use std::io::{Read, Seek, SeekFrom, Write};
7use std::path::Path;
8
9const IDENTITY: [i32; 9] = [0x10000, 0, 0, 0, 0x10000, 0, 0, 0, 0x4000_0000];
10const FLUSH_AT: usize = 4 << 20;
11
12#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)]
13pub enum Brand {
14    #[default]
15    Mp4,
16    Mov,
17}
18
19#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)]
20pub struct Options {
21    pub brand: Brand,
22    pub reserve: u64,
23}
24
25pub fn rotation_matrix(degrees: u16) -> [i32; 9] {
26    match degrees % 360 {
27        90 => [0, 0x10000, 0, -0x10000, 0, 0, 0, 0, 0x4000_0000],
28        180 => [-0x10000, 0, 0, 0, -0x10000, 0, 0, 0, 0x4000_0000],
29        270 => [0, -0x10000, 0, 0x10000, 0, 0, 0, 0, 0x4000_0000],
30        _ => IDENTITY,
31    }
32}
33
34#[derive(Clone, Debug)]
35pub struct TrackConfig {
36    pub entry: SampleEntry,
37    pub timescale: u32,
38    pub language: String,
39    pub name: Option<String>,
40    pub matrix: [i32; 9],
41    pub edits: Vec<Edit>,
42    pub media_start: Option<i64>,
43}
44
45impl TrackConfig {
46    pub fn new(entry: SampleEntry, timescale: u32) -> Self {
47        TrackConfig {
48            entry,
49            timescale,
50            language: "und".into(),
51            name: None,
52            matrix: IDENTITY,
53            edits: Vec::new(),
54            media_start: None,
55        }
56    }
57}
58
59#[derive(Clone, Copy, PartialEq, Eq)]
60enum Handler {
61    Video,
62    Audio,
63}
64
65#[derive(Clone, Debug, Default)]
66struct Tables {
67    stts: Vec<(u32, u32)>,
68    ctts: Vec<(u32, i32)>,
69    sizes: Vec<u32>,
70    sync: Vec<u32>,
71    all_sync: bool,
72    chunks: Vec<(u64, u32, u32)>,
73    media_duration: u64,
74    comp_start: Option<i64>,
75    comp_end: i64,
76}
77
78impl Tables {
79    fn reach_end(&mut self) {
80        let end = u64::try_from(self.comp_end).unwrap_or(0);
81        let Some(extra) = end.checked_sub(self.media_duration).filter(|&e| e > 0) else {
82            return;
83        };
84        let Some(last) = self.stts.last_mut() else {
85            return;
86        };
87        let Ok(d) = u32::try_from(last.1 as u64 + extra) else {
88            return;
89        };
90        if last.0 > 1 {
91            last.0 -= 1;
92            self.stts.push((1, d));
93        } else {
94            last.1 = d;
95        }
96        self.media_duration = end;
97    }
98}
99
100fn push_run<T: PartialEq + Copy>(runs: &mut Vec<(u32, T)>, v: T) {
101    if let Some(last) = runs.last_mut()
102        && last.1 == v
103        && last.0 < u32::MAX
104    {
105        last.0 += 1;
106        return;
107    }
108    runs.push((1, v));
109}
110
111struct Pending {
112    data: Vec<u8>,
113    pts: i64,
114    dts: i64,
115    keyframe: bool,
116}
117
118struct WTrack {
119    cfg: TrackConfig,
120    entries: Vec<SampleEntry>,
121    entry: u32,
122    id: u32,
123    handler: Handler,
124    t: Tables,
125    origin: Option<i64>,
126    last_duration: u32,
127    pending: Option<Pending>,
128    spare: Vec<u8>,
129}
130
131impl WTrack {
132    fn new(cfg: TrackConfig, id: u32) -> Result<Self> {
133        if cfg.timescale == 0 {
134            return Err(Error::Invalid("track timescale must be non-zero".into()));
135        }
136        let handler = if cfg.entry.is_video() {
137            Handler::Video
138        } else if cfg.entry.is_audio() {
139            Handler::Audio
140        } else {
141            return Err(Error::Unsupported(
142                "tracks other than audio and video".into(),
143            ));
144        };
145        Ok(WTrack {
146            entries: vec![cfg.entry.clone()],
147            entry: 0,
148            cfg,
149            id,
150            handler,
151            t: Tables {
152                all_sync: true,
153                ..Default::default()
154            },
155            origin: None,
156            last_duration: 0,
157            pending: None,
158            spare: Vec::new(),
159        })
160    }
161
162    fn record(
163        &mut self,
164        size: u32,
165        pts: i64,
166        dts: i64,
167        duration: u32,
168        keyframe: bool,
169    ) -> Result<i64> {
170        let origin = *self.origin.get_or_insert(dts);
171        let offset = pts
172            .checked_sub(dts)
173            .and_then(|d| i32::try_from(d).ok())
174            .ok_or_else(|| Error::Invalid("composition offset out of range".into()))?;
175        let t = &mut self.t;
176        push_run(&mut t.stts, duration);
177        push_run(&mut t.ctts, offset);
178        t.sizes.push(size);
179        if keyframe {
180            t.sync.push(t.sizes.len() as u32);
181        } else {
182            t.all_sync = false;
183        }
184        let start = pts.saturating_sub(origin);
185        t.comp_start = Some(t.comp_start.map_or(start, |s| s.min(start)));
186        t.comp_end = t.comp_end.max(start.saturating_add(duration as i64));
187        t.media_duration = t.media_duration.saturating_add(duration as u64);
188        self.last_duration = duration;
189        Ok(origin)
190    }
191}
192
193fn pack_language(lang: &str) -> u16 {
194    let b = lang.as_bytes();
195    if b.len() == 3 && b.iter().all(|c| c.is_ascii_lowercase()) {
196        (((b[0] - 0x60) as u16) << 10) | (((b[1] - 0x60) as u16) << 5) | ((b[2] - 0x60) as u16)
197    } else {
198        0x55C4
199    }
200}
201
202fn rescale(v: u64, from: u32, to: u32) -> u64 {
203    if from == 0 {
204        return 0;
205    }
206    (v as u128 * to as u128 / from as u128).min(u64::MAX as u128) as u64
207}
208
209fn write_ftyp(b: &mut BoxBuf, brand: Brand) {
210    let m = b.start(b"ftyp");
211    match brand {
212        Brand::Mp4 => {
213            b.bytes(b"isom");
214            b.u32(0x200);
215            for c in [b"isom", b"iso2", b"avc1", b"mp41"] {
216                b.bytes(c);
217            }
218        }
219        Brand::Mov => {
220            b.bytes(b"qt  ");
221            b.u32(0x200);
222            b.bytes(b"qt  ");
223        }
224    }
225    b.end(m);
226}
227
228fn effective_edits(t: &WTrack, movie_ts: u32) -> Vec<Edit> {
229    if !t.cfg.edits.is_empty() {
230        return t.cfg.edits.clone();
231    }
232    let start = t
233        .cfg
234        .media_start
235        .unwrap_or_else(|| t.t.comp_start.unwrap_or(0).max(0));
236    if start == 0 {
237        return Vec::new();
238    }
239    let end = t.t.comp_end.max(t.t.media_duration as i64);
240    let rest = end.saturating_sub(start).max(0) as u64;
241    vec![Edit {
242        segment_duration: rescale(rest, t.cfg.timescale, movie_ts),
243        media_time: start,
244        media_rate: 0x10000,
245    }]
246}
247
248fn presentation_duration(t: &WTrack, movie_ts: u32) -> u64 {
249    let e = effective_edits(t, movie_ts);
250    if e.is_empty() {
251        rescale(t.t.media_duration, t.cfg.timescale, movie_ts)
252    } else {
253        e.iter()
254            .fold(0u64, |a, e| a.saturating_add(e.segment_duration))
255    }
256}
257
258fn write_times(b: &mut BoxBuf, v1: bool, timescale: u32, duration: u64) {
259    b.zeros(if v1 { 16 } else { 8 });
260    b.u32(timescale);
261    write_duration(b, v1, duration);
262}
263
264fn write_duration(b: &mut BoxBuf, v1: bool, duration: u64) {
265    if v1 {
266        b.u64(duration);
267    } else {
268        b.u32(duration as u32);
269    }
270}
271
272fn write_mvhd(b: &mut BoxBuf, timescale: u32, duration: u64, next_id: u32) {
273    let v1 = duration > u32::MAX as u64;
274    let m = b.start_full(b"mvhd", v1 as u8, 0);
275    write_times(b, v1, timescale, duration);
276    b.u32(0x10000);
277    b.u16(0x100);
278    b.zeros(10);
279    for v in IDENTITY {
280        b.i32(v);
281    }
282    b.zeros(24);
283    b.u32(next_id);
284    b.end(m);
285}
286
287fn write_trak(
288    b: &mut BoxBuf,
289    t: &WTrack,
290    movie_ts: u32,
291    shift: u64,
292    co64: bool,
293    fragmented: bool,
294    qt: bool,
295) -> Result<()> {
296    let trak = b.start(b"trak");
297    let dur = if fragmented {
298        0
299    } else {
300        presentation_duration(t, movie_ts)
301    };
302    let v1 = dur > u32::MAX as u64;
303    let m = b.start_full(b"tkhd", v1 as u8, 3);
304    b.zeros(if v1 { 16 } else { 8 });
305    b.u32(t.id);
306    b.u32(0);
307    write_duration(b, v1, dur);
308    b.zeros(8);
309    b.i16(0);
310    b.u16(0);
311    b.i16(if t.handler == Handler::Audio {
312        0x100
313    } else {
314        0
315    });
316    b.u16(0);
317    for v in t.cfg.matrix {
318        b.i32(v);
319    }
320    let (w, h) = match &t.cfg.entry.video {
321        Some(v) => {
322            let w = match v.pixel_aspect {
323                Some((hs, vs)) if hs > 0 && vs > 0 => {
324                    (v.width as u64 * hs as u64 / vs as u64) as u32
325                }
326                _ => v.width as u32,
327            };
328            (w, v.height as u32)
329        }
330        None => (0, 0),
331    };
332    b.u32(w << 16);
333    b.u32(h << 16);
334    b.end(m);
335    let edits = if fragmented {
336        Vec::new()
337    } else {
338        effective_edits(t, movie_ts)
339    };
340    if !edits.is_empty() {
341        let e = b.start(b"edts");
342        let v1 = edits.iter().any(|e| {
343            e.segment_duration > u32::MAX as u64
344                || e.media_time > i32::MAX as i64
345                || e.media_time < i32::MIN as i64
346        });
347        let l = b.start_full(b"elst", v1 as u8, 0);
348        b.u32(edits.len() as u32);
349        for ed in &edits {
350            if v1 {
351                b.u64(ed.segment_duration);
352                b.i64(ed.media_time);
353            } else {
354                b.u32(ed.segment_duration as u32);
355                b.i32(ed.media_time as i32);
356            }
357            b.i32(ed.media_rate);
358        }
359        b.end(l);
360        b.end(e);
361    }
362    let mdia = b.start(b"mdia");
363    let mv1 = t.t.media_duration > u32::MAX as u64;
364    let m = b.start_full(b"mdhd", mv1 as u8, 0);
365    write_times(b, mv1, t.cfg.timescale, t.t.media_duration);
366    b.u16(pack_language(&t.cfg.language));
367    b.u16(0);
368    b.end(m);
369    let (htype, default_name) = match t.handler {
370        Handler::Video => (b"vide", "VideoHandler"),
371        Handler::Audio => (b"soun", "SoundHandler"),
372    };
373    let name = t.cfg.name.as_deref().unwrap_or(default_name).as_bytes();
374    let m = b.start_full(b"hdlr", 0, 0);
375    b.bytes(if qt { b"mhlr" } else { &[0, 0, 0, 0] });
376    b.bytes(htype);
377    b.zeros(12);
378    if qt {
379        let n = name.len().min(255);
380        b.u8(n as u8);
381        b.bytes(&name[..n]);
382    } else {
383        b.bytes(name);
384        b.u8(0);
385    }
386    b.end(m);
387    let minf = b.start(b"minf");
388    match t.handler {
389        Handler::Video => {
390            let m = b.start_full(b"vmhd", 0, 1);
391            b.zeros(8);
392            b.end(m);
393        }
394        Handler::Audio => {
395            let m = b.start_full(b"smhd", 0, 0);
396            b.zeros(4);
397            b.end(m);
398        }
399    }
400    if qt {
401        let m = b.start_full(b"hdlr", 0, 0);
402        b.bytes(b"dhlr");
403        b.bytes(b"url ");
404        b.zeros(12);
405        b.u8(0);
406        b.end(m);
407    }
408    let d = b.start(b"dinf");
409    let r = b.start_full(b"dref", 0, 0);
410    b.u32(1);
411    let u = b.start_full(b"url ", 0, 1);
412    b.end(u);
413    b.end(r);
414    b.end(d);
415    let stbl = b.start(b"stbl");
416    let m = b.start_full(b"stsd", 0, 0);
417    b.u32(t.entries.len() as u32);
418    for e in &t.entries {
419        e.write(b)?;
420    }
421    b.end(m);
422    let tb = &t.t;
423    let m = b.start_full(b"stts", 0, 0);
424    b.u32(tb.stts.len() as u32);
425    for &(c, d) in &tb.stts {
426        b.u32(c);
427        b.u32(d);
428    }
429    b.end(m);
430    if tb.ctts.iter().any(|&(_, o)| o != 0) {
431        let neg = tb.ctts.iter().any(|&(_, o)| o < 0);
432        let m = b.start_full(b"ctts", neg as u8, 0);
433        b.u32(tb.ctts.len() as u32);
434        for &(c, o) in &tb.ctts {
435            b.u32(c);
436            b.i32(o);
437        }
438        b.end(m);
439    }
440    if !tb.all_sync {
441        let m = b.start_full(b"stss", 0, 0);
442        b.u32(tb.sync.len() as u32);
443        for &s in &tb.sync {
444            b.u32(s);
445        }
446        b.end(m);
447    }
448    let m = b.start_full(b"stsc", 0, 0);
449    let mut runs: Vec<(u32, u32, u32)> = Vec::new();
450    for (i, &(_, n, e)) in tb.chunks.iter().enumerate() {
451        if runs.last().is_none_or(|r| r.1 != n || r.2 != e + 1) {
452            runs.push((i as u32 + 1, n, e + 1));
453        }
454    }
455    b.u32(runs.len() as u32);
456    for (first, n, e) in runs {
457        b.u32(first);
458        b.u32(n);
459        b.u32(e);
460    }
461    b.end(m);
462    let m = b.start_full(b"stsz", 0, 0);
463    let first = tb.sizes.first().copied().unwrap_or(0);
464    if first != 0 && tb.sizes.iter().all(|&s| s == first) {
465        b.u32(first);
466        b.u32(tb.sizes.len() as u32);
467    } else {
468        b.u32(0);
469        b.u32(tb.sizes.len() as u32);
470        for &s in &tb.sizes {
471            b.u32(s);
472        }
473    }
474    b.end(m);
475    let co64 = co64
476        || tb
477            .chunks
478            .iter()
479            .any(|&(o, _, _)| o + shift > u32::MAX as u64);
480    let m = b.start_full(if co64 { b"co64" } else { b"stco" }, 0, 0);
481    b.u32(tb.chunks.len() as u32);
482    for &(o, _, _) in &tb.chunks {
483        if co64 {
484            b.u64(o + shift);
485        } else {
486            b.u32((o + shift) as u32);
487        }
488    }
489    b.end(m);
490    b.end(stbl);
491    b.end(minf);
492    b.end(mdia);
493    b.end(trak);
494    Ok(())
495}
496
497pub struct Writer<W: Write + Seek = File> {
498    w: W,
499    brand: Brand,
500    reserve_pos: u64,
501    reserve: u64,
502    tracks: Vec<WTrack>,
503    movie_timescale: u32,
504    pad_pos: u64,
505    mdat_pos: u64,
506    pos: u64,
507    buf: Vec<u8>,
508    last_track: Option<usize>,
509}
510
511impl Writer<File> {
512    pub fn create(path: impl AsRef<Path>) -> Result<Self> {
513        Self::create_with(path, Options::default())
514    }
515
516    pub fn create_with(path: impl AsRef<Path>, options: Options) -> Result<Self> {
517        let f = File::options()
518            .read(true)
519            .write(true)
520            .create(true)
521            .truncate(true)
522            .open(path.as_ref())
523            .map_err(|e| Error::at(path.as_ref(), e))?;
524        Self::new(f, options)
525    }
526}
527
528impl<W: Write + Seek> Writer<W> {
529    pub fn new(mut w: W, options: Options) -> Result<Self> {
530        let start = w.stream_position()?;
531        let mut b = BoxBuf::new();
532        write_ftyp(&mut b, options.brand);
533        let reserve_pos = start + b.len() as u64;
534        let reserve = if options.reserve == 0 {
535            0
536        } else {
537            options.reserve.max(8)
538        };
539        if reserve > 0 {
540            b.u32(
541                u32::try_from(reserve)
542                    .map_err(|_| Error::Invalid("moov reserve too large".into()))?,
543            );
544            b.bytes(b"free");
545            b.zeros(reserve as usize - 8);
546        }
547        let pad_pos = start + b.len() as u64;
548        b.leaf(
549            if options.brand == Brand::Mov {
550                b"wide"
551            } else {
552                b"free"
553            },
554            &[],
555        );
556        let mdat_pos = start + b.len() as u64;
557        b.leaf(b"mdat", &[]);
558        w.write_all(&b.buf)?;
559        let pos = start + b.len() as u64;
560        Ok(Writer {
561            w,
562            brand: options.brand,
563            reserve_pos,
564            reserve,
565            tracks: Vec::new(),
566            movie_timescale: 1000,
567            pad_pos,
568            mdat_pos,
569            pos,
570            buf: Vec::with_capacity(FLUSH_AT + (1 << 20)),
571            last_track: None,
572        })
573    }
574
575    pub fn set_movie_timescale(&mut self, timescale: u32) {
576        if timescale > 0 {
577            self.movie_timescale = timescale;
578        }
579    }
580
581    pub fn add_track(&mut self, cfg: TrackConfig) -> Result<usize> {
582        let t = WTrack::new(cfg, self.tracks.len() as u32 + 1)?;
583        self.tracks.push(t);
584        Ok(self.tracks.len() - 1)
585    }
586
587    pub fn add_entry(&mut self, track: usize, entry: SampleEntry) -> Result<u32> {
588        let t = self.track(track)?;
589        t.entries.push(entry);
590        Ok(t.entries.len() as u32 - 1)
591    }
592
593    pub fn use_entry(&mut self, track: usize, entry: u32) -> Result<()> {
594        self.flush_pending(track, None)?;
595        let t = self.track(track)?;
596        if entry as usize >= t.entries.len() {
597            return Err(Error::Invalid(format!(
598                "no sample entry {entry} in track {track}"
599            )));
600        }
601        t.entry = entry;
602        Ok(())
603    }
604
605    fn track(&mut self, track: usize) -> Result<&mut WTrack> {
606        self.tracks
607            .get_mut(track)
608            .ok_or_else(|| Error::Invalid(format!("no track {track}")))
609    }
610
611    pub fn push(&mut self, track: usize, packet: &Packet) -> Result<()> {
612        self.flush_pending(track, Some(packet.dts))?;
613        let t = self.track(track)?;
614        let mut data = std::mem::take(&mut t.spare);
615        data.clear();
616        data.extend_from_slice(&packet.data);
617        t.pending = Some(Pending {
618            data,
619            pts: packet.pts,
620            dts: packet.dts,
621            keyframe: packet.keyframe,
622        });
623        Ok(())
624    }
625
626    pub fn push_timed(&mut self, track: usize, packet: &Packet, duration: u32) -> Result<()> {
627        self.flush_pending(track, Some(packet.dts))?;
628        self.write_sample(
629            track,
630            &packet.data,
631            packet.pts,
632            packet.dts,
633            duration,
634            packet.keyframe,
635        )
636    }
637
638    pub fn push_sample(
639        &mut self,
640        track: usize,
641        data: &[u8],
642        pts: i64,
643        dts: i64,
644        duration: u32,
645        keyframe: bool,
646    ) -> Result<()> {
647        self.flush_pending(track, Some(dts))?;
648        self.write_sample(track, data, pts, dts, duration, keyframe)
649    }
650
651    fn flush_pending(&mut self, track: usize, next_dts: Option<i64>) -> Result<()> {
652        let t = self.track(track)?;
653        let Some(p) = t.pending.take() else {
654            return Ok(());
655        };
656        let duration = match next_dts {
657            Some(d) => d
658                .checked_sub(p.dts)
659                .and_then(|d| u32::try_from(d).ok())
660                .ok_or_else(|| Error::Invalid("packets must come in decode order".into()))?,
661            None => t.last_duration,
662        };
663        self.write_sample(track, &p.data, p.pts, p.dts, duration, p.keyframe)?;
664        self.tracks[track].spare = p.data;
665        Ok(())
666    }
667
668    fn write_sample(
669        &mut self,
670        track: usize,
671        data: &[u8],
672        pts: i64,
673        dts: i64,
674        duration: u32,
675        keyframe: bool,
676    ) -> Result<()> {
677        let size = u32::try_from(data.len())
678            .map_err(|_| Error::Invalid("sample larger than 4 GiB".into()))?;
679        let offset = self.pos;
680        let contiguous = self.last_track == Some(track);
681        self.last_track = Some(track);
682        let t = self.track(track)?;
683        t.record(size, pts, dts, duration, keyframe)?;
684        let entry = t.entry;
685        match t.t.chunks.last_mut() {
686            Some(c) if contiguous && c.2 == entry => c.1 += 1,
687            _ => t.t.chunks.push((offset, 1, entry)),
688        }
689        self.buf.extend_from_slice(data);
690        self.pos += data.len() as u64;
691        if self.buf.len() >= FLUSH_AT {
692            self.w.write_all(&self.buf)?;
693            self.buf.clear();
694        }
695        Ok(())
696    }
697
698    fn finish_mdat(&mut self) -> Result<()> {
699        for i in 0..self.tracks.len() {
700            self.flush_pending(i, None)?;
701        }
702        for t in &mut self.tracks {
703            t.t.reach_end();
704        }
705        self.w.write_all(&self.buf)?;
706        self.buf = Vec::new();
707        let size = self.pos - self.mdat_pos;
708        if size <= u32::MAX as u64 {
709            self.w.seek(SeekFrom::Start(self.mdat_pos))?;
710            self.w.write_all(&(size as u32).to_be_bytes())?;
711        } else {
712            let mut h = [0u8; 16];
713            h[..4].copy_from_slice(&1u32.to_be_bytes());
714            h[4..8].copy_from_slice(b"mdat");
715            h[8..].copy_from_slice(&(self.pos - self.pad_pos).to_be_bytes());
716            self.w.seek(SeekFrom::Start(self.pad_pos))?;
717            self.w.write_all(&h)?;
718        }
719        self.w.seek(SeekFrom::Start(self.pos))?;
720        Ok(())
721    }
722
723    fn build_moov(&self, shift: u64, co64: bool) -> Result<Vec<u8>> {
724        let ts = self.movie_timescale;
725        let mut b = BoxBuf::new();
726        let m = b.start(b"moov");
727        let duration = self
728            .tracks
729            .iter()
730            .map(|t| presentation_duration(t, ts))
731            .max()
732            .unwrap_or(0);
733        write_mvhd(&mut b, ts, duration, self.tracks.len() as u32 + 1);
734        for t in &self.tracks {
735            write_trak(&mut b, t, ts, shift, co64, false, self.brand == Brand::Mov)?;
736        }
737        b.end(m);
738        Ok(b.buf)
739    }
740
741    pub(crate) fn position(&self) -> u64 {
742        self.pos
743    }
744
745    pub(crate) fn moov_len(&mut self) -> Result<u64> {
746        self.finish_mdat()?;
747        Ok(self.build_moov(0, false)?.len() as u64)
748    }
749
750    pub fn finish(mut self) -> Result<W> {
751        self.finish_mdat()?;
752        let moov = self.build_moov(0, false)?;
753        let room = self.reserve;
754        if room > 0 && (moov.len() as u64 == room || moov.len() as u64 + 8 <= room) {
755            self.w.seek(SeekFrom::Start(self.reserve_pos))?;
756            self.w.write_all(&moov)?;
757            let rest = room - moov.len() as u64;
758            if rest > 0 {
759                self.w.write_all(&(rest as u32).to_be_bytes())?;
760                self.w.write_all(b"free")?;
761            }
762            self.w.seek(SeekFrom::Start(self.pos))?;
763        } else {
764            self.w.write_all(&moov)?;
765        }
766        self.w.flush()?;
767        Ok(self.w)
768    }
769}
770
771impl<W: Read + Write + Seek> Writer<W> {
772    pub fn finish_faststart(self) -> Result<W> {
773        self.finish_faststart_with(|_, _| true)
774    }
775
776    pub fn finish_faststart_with(mut self, mut moved: impl FnMut(u64, u64) -> bool) -> Result<W> {
777        self.finish_mdat()?;
778        let mut moov = self.build_moov(0, false)?;
779        let mut co64 = false;
780        loop {
781            let next = self.build_moov(moov.len() as u64, co64)?;
782            if next.len() == moov.len() {
783                moov = next;
784                break;
785            }
786            co64 = true;
787            moov = next;
788        }
789        let shift = moov.len() as u64;
790        let start = self.pad_pos;
791        let end = self.pos;
792        let mut buf = vec![0u8; 8 << 20];
793        let mut cur = end;
794        while cur > start {
795            let n = (cur - start).min(buf.len() as u64) as usize;
796            cur -= n as u64;
797            self.w.seek(SeekFrom::Start(cur))?;
798            self.w.read_exact(&mut buf[..n])?;
799            self.w.seek(SeekFrom::Start(cur + shift))?;
800            self.w.write_all(&buf[..n])?;
801            if !moved(end - cur, end - start) {
802                return Err(Error::Cancelled);
803            }
804        }
805        self.w.seek(SeekFrom::Start(start))?;
806        self.w.write_all(&moov)?;
807        self.w.seek(SeekFrom::Start(end + shift))?;
808        self.w.flush()?;
809        Ok(self.w)
810    }
811}
812
813struct FragSample {
814    data: Vec<u8>,
815    duration: u32,
816    offset: i32,
817    keyframe: bool,
818}
819
820pub struct FragmentedWriter<W: Write = File> {
821    w: W,
822    tracks: Vec<WTrack>,
823    pending: Vec<Vec<FragSample>>,
824    next_dts: Vec<u64>,
825    fragment: f64,
826    seq: u32,
827    init_written: bool,
828}
829
830impl FragmentedWriter<File> {
831    pub fn create(path: impl AsRef<Path>, fragment_seconds: f64) -> Result<Self> {
832        let path = path.as_ref();
833        let f = File::create(path).map_err(|e| Error::at(path, e))?;
834        Ok(Self::new(f, fragment_seconds))
835    }
836}
837
838impl<W: Write> FragmentedWriter<W> {
839    pub fn new(w: W, fragment_seconds: f64) -> Self {
840        FragmentedWriter {
841            w,
842            tracks: Vec::new(),
843            pending: Vec::new(),
844            next_dts: Vec::new(),
845            fragment: fragment_seconds.max(0.0),
846            seq: 0,
847            init_written: false,
848        }
849    }
850
851    pub fn add_track(&mut self, cfg: TrackConfig) -> Result<usize> {
852        if self.init_written {
853            return Err(Error::Invalid(
854                "tracks must be added before the first fragment".into(),
855            ));
856        }
857        let t = WTrack::new(cfg, self.tracks.len() as u32 + 1)?;
858        self.tracks.push(t);
859        self.pending.push(Vec::new());
860        self.next_dts.push(0);
861        Ok(self.tracks.len() - 1)
862    }
863
864    fn write_init(&mut self) -> Result<()> {
865        let ts = 1000;
866        let mut b = BoxBuf::new();
867        write_ftyp(&mut b, Brand::Mp4);
868        let m = b.start(b"moov");
869        write_mvhd(&mut b, ts, 0, self.tracks.len() as u32 + 1);
870        for t in &self.tracks {
871            write_trak(&mut b, t, ts, 0, false, true, false)?;
872        }
873        let mv = b.start(b"mvex");
874        for t in &self.tracks {
875            let x = b.start_full(b"trex", 0, 0);
876            b.u32(t.id);
877            b.u32(1);
878            b.zeros(12);
879            b.end(x);
880        }
881        b.end(mv);
882        b.end(m);
883        self.w.write_all(&b.buf)?;
884        self.init_written = true;
885        Ok(())
886    }
887
888    pub fn push(&mut self, track: usize, packet: &Packet) -> Result<()> {
889        let t = self
890            .tracks
891            .get_mut(track)
892            .ok_or_else(|| Error::Invalid(format!("no track {track}")))?;
893        if let Some(p) = t.pending.take() {
894            let duration = packet
895                .dts
896                .checked_sub(p.dts)
897                .and_then(|d| u32::try_from(d).ok())
898                .ok_or_else(|| Error::Invalid("packets must come in decode order".into()))?;
899            self.queue(track, p, duration)?;
900        }
901        let t = &self.tracks[track];
902        let pending_ticks: u64 = self.pending[track].iter().map(|s| s.duration as u64).sum();
903        if packet.keyframe
904            && t.handler == Handler::Video
905            && pending_ticks as f64 >= self.fragment * t.cfg.timescale as f64
906            && !self.pending[track].is_empty()
907        {
908            self.flush_fragment()?;
909        }
910        self.tracks[track].pending = Some(Pending {
911            data: packet.data.clone(),
912            pts: packet.pts,
913            dts: packet.dts,
914            keyframe: packet.keyframe,
915        });
916        Ok(())
917    }
918
919    fn queue(&mut self, track: usize, p: Pending, duration: u32) -> Result<()> {
920        let t = &mut self.tracks[track];
921        let offset = p
922            .pts
923            .checked_sub(p.dts)
924            .and_then(|d| i32::try_from(d).ok())
925            .ok_or_else(|| Error::Invalid("composition offset out of range".into()))?;
926        t.last_duration = duration;
927        self.pending[track].push(FragSample {
928            data: p.data,
929            duration,
930            offset,
931            keyframe: p.keyframe,
932        });
933        Ok(())
934    }
935
936    fn build_moof(&self, data_offsets: &[i32]) -> Vec<u8> {
937        let mut b = BoxBuf::new();
938        let m = b.start(b"moof");
939        let h = b.start_full(b"mfhd", 0, 0);
940        b.u32(self.seq);
941        b.end(h);
942        for (i, t) in self.tracks.iter().enumerate() {
943            let p = &self.pending[i];
944            if p.is_empty() {
945                continue;
946            }
947            let tf = b.start(b"traf");
948            let h = b.start_full(b"tfhd", 0, 0x20000);
949            b.u32(t.id);
950            b.end(h);
951            let h = b.start_full(b"tfdt", 1, 0);
952            b.u64(self.next_dts[i]);
953            b.end(h);
954            let neg = p.iter().any(|s| s.offset < 0);
955            let h = b.start_full(b"trun", neg as u8, 0x1 | 0x100 | 0x200 | 0x400 | 0x800);
956            b.u32(p.len() as u32);
957            b.i32(data_offsets[i]);
958            for s in p {
959                b.u32(s.duration);
960                b.u32(s.data.len() as u32);
961                b.u32(if s.keyframe { 0x0200_0000 } else { 0x0101_0000 });
962                b.i32(s.offset);
963            }
964            b.end(h);
965            b.end(tf);
966        }
967        b.end(m);
968        b.buf
969    }
970
971    pub fn flush_fragment(&mut self) -> Result<()> {
972        if !self.init_written {
973            self.write_init()?;
974        }
975        if self.pending.iter().all(Vec::is_empty) {
976            return Ok(());
977        }
978        self.seq += 1;
979        let probe = self.build_moof(&vec![0; self.tracks.len()]);
980        let data_len: u64 = self
981            .pending
982            .iter()
983            .flatten()
984            .map(|s| s.data.len() as u64)
985            .sum();
986        let large = data_len + 8 > u32::MAX as u64;
987        let mut offs = Vec::with_capacity(self.tracks.len());
988        let mut acc = probe.len() as u64 + if large { 16 } else { 8 };
989        for p in &self.pending {
990            offs.push(i32::try_from(acc).map_err(|_| Error::Invalid("fragment too large".into()))?);
991            acc += p.iter().map(|s| s.data.len() as u64).sum::<u64>();
992        }
993        let moof = self.build_moof(&offs);
994        self.w.write_all(&moof)?;
995        if large {
996            self.w.write_all(&1u32.to_be_bytes())?;
997            self.w.write_all(b"mdat")?;
998            self.w.write_all(&(data_len + 16).to_be_bytes())?;
999        } else {
1000            self.w.write_all(&((data_len + 8) as u32).to_be_bytes())?;
1001            self.w.write_all(b"mdat")?;
1002        }
1003        for (i, p) in self.pending.iter_mut().enumerate() {
1004            for s in p.drain(..) {
1005                self.w.write_all(&s.data)?;
1006                self.next_dts[i] += s.duration as u64;
1007            }
1008        }
1009        Ok(())
1010    }
1011
1012    pub fn finish(mut self) -> Result<W> {
1013        for i in 0..self.tracks.len() {
1014            if let Some(p) = self.tracks[i].pending.take() {
1015                let d = self.tracks[i].last_duration;
1016                self.queue(i, p, d)?;
1017            }
1018        }
1019        self.flush_fragment()?;
1020        self.w.flush()?;
1021        Ok(self.w)
1022    }
1023}
1024
1025impl TrackConfig {
1026    pub fn from_track(track: &crate::mp4::Track) -> Result<Self> {
1027        let entry = track
1028            .entry()
1029            .cloned()
1030            .ok_or_else(|| Error::Invalid(format!("track {} has no sample entry", track.id)))?;
1031        Ok(TrackConfig {
1032            entry,
1033            timescale: track.timescale,
1034            language: track.language.clone(),
1035            name: (!track.name.is_empty()).then(|| track.name.clone()),
1036            matrix: track.matrix,
1037            edits: Vec::new(),
1038            media_start: None,
1039        })
1040    }
1041}