pito_codecs/mp4/
remux.rs

1use crate::media::{Error, Result};
2use crate::mp4::reader::{Edit, Reader, Track, TrackKind};
3use crate::mp4::source::ByteSource;
4use crate::mp4::truncated;
5use crate::mp4::writer::{Brand, Options, TrackConfig, Writer};
6use std::io::{self, Seek, SeekFrom, Write};
7use std::ops::Range;
8use std::path::Path;
9
10const RUN: u64 = 8 << 20;
11const INTERLEAVE: f64 = 0.5;
12
13#[derive(Clone, Copy, Debug, PartialEq, Eq)]
14pub enum Snap {
15    Inward,
16    Outward,
17}
18
19#[derive(Clone, Copy, Debug, PartialEq)]
20pub struct Span {
21    pub start: f64,
22    pub end: f64,
23}
24
25#[derive(Clone, Debug, PartialEq)]
26pub struct Cut {
27    pub span: Span,
28    pub tracks: Vec<(usize, Range<usize>)>,
29}
30
31pub(crate) trait Timeline {
32    fn len(&self) -> usize;
33    fn keyframes(&self) -> &[u32];
34    fn pts(&self, i: usize) -> i64;
35    fn duration(&self, i: usize) -> i64;
36    fn ticks(&self, seconds: f64) -> i64;
37    fn seconds(&self, ticks: i64) -> f64;
38
39    fn end(&self) -> i64 {
40        (0..self.len())
41            .map(|i| self.pts(i).saturating_add(self.duration(i)))
42            .max()
43            .unwrap_or(0)
44    }
45}
46
47impl Timeline for Track {
48    fn len(&self) -> usize {
49        self.samples.len()
50    }
51    fn keyframes(&self) -> &[u32] {
52        Track::keyframes(self)
53    }
54    fn pts(&self, i: usize) -> i64 {
55        self.presentation(i).unwrap_or(0)
56    }
57    fn duration(&self, i: usize) -> i64 {
58        self.samples[i].duration as i64
59    }
60    fn ticks(&self, seconds: f64) -> i64 {
61        Track::ticks(self, seconds)
62    }
63    fn seconds(&self, ticks: i64) -> f64 {
64        Track::seconds(self, ticks)
65    }
66}
67
68fn video_range<T: Timeline>(
69    t: &T,
70    from: f64,
71    to: f64,
72    snap: Snap,
73) -> Result<(Range<usize>, i64, i64)> {
74    let keys: Vec<usize> = t.keyframes().iter().map(|&k| k as usize).collect();
75    if keys.is_empty() {
76        return Err(Error::Invalid("video track has no keyframes".into()));
77    }
78    let n = t.len();
79    let gop_min: Vec<i64> = keys
80        .iter()
81        .enumerate()
82        .map(|(j, &k)| {
83            let end = keys.get(j + 1).copied().unwrap_or(n);
84            (k..end).map(|i| t.pts(i)).min().unwrap_or(i64::MAX)
85        })
86        .collect();
87    let (from_t, to_t) = (t.ticks(from), t.ticks(to));
88    let end_t = t.end();
89    match snap {
90        Snap::Outward => {
91            let ja = keys.iter().rposition(|&k| t.pts(k) <= from_t).unwrap_or(0);
92            let b = (ja + 1..keys.len())
93                .find(|&j| gop_min[j] >= to_t)
94                .map_or(n, |j| keys[j]);
95            let a = keys[ja];
96            let copied_end = (a..b)
97                .map(|i| t.pts(i).saturating_add(t.duration(i)))
98                .max()
99                .unwrap_or(0);
100            Ok((a..b, from_t.max(t.pts(a)), to_t.min(copied_end)))
101        }
102        Snap::Inward => {
103            let ja = keys
104                .iter()
105                .position(|&k| t.pts(k) >= from_t)
106                .ok_or_else(|| Error::Invalid(format!("no keyframe after {from:.3} s")))?;
107            let start = t.pts(keys[ja]);
108            if to_t >= end_t {
109                return Ok((keys[ja]..n, start, end_t));
110            }
111            let jb = (ja + 1..keys.len())
112                .rev()
113                .find(|&j| t.pts(keys[j]) <= to_t)
114                .ok_or_else(|| {
115                    Error::Invalid(format!(
116                        "no keyframe-to-keyframe stretch inside {from:.3}..{to:.3} s"
117                    ))
118                })?;
119            Ok((keys[ja]..keys[jb], start, t.pts(keys[jb])))
120        }
121    }
122}
123
124fn audio_range<T: Timeline>(t: &T, start: f64, end: f64) -> Range<usize> {
125    let (s, e) = (t.ticks(start), t.ticks(end));
126    let n = t.len();
127    let first = (0..n)
128        .position(|i| t.pts(i).saturating_add(t.duration(i)) > s)
129        .unwrap_or(n);
130    let first = if first < n && t.pts(first) < s {
131        first.saturating_sub(1)
132    } else {
133        first
134    };
135    let last = (first..n).find(|&i| t.pts(i) >= e).unwrap_or(n);
136    first..last
137}
138
139pub(crate) fn plan_tracks<T: Timeline>(
140    total: f64,
141    tracks: &[T],
142    video: Option<usize>,
143    audio: Vec<usize>,
144    (from, to): (f64, f64),
145    snap: Snap,
146) -> Result<Cut> {
147    let (from, to) = (from.max(0.0), to.min(total.max(from)));
148    if to.is_nan() || from.is_nan() || to <= from {
149        return Err(Error::Invalid(format!("empty range {from:.3}..{to:.3} s")));
150    }
151    let mut cut = Vec::new();
152    let span = match video {
153        Some(v) => {
154            let t = &tracks[v];
155            let (range, start, end) = video_range(t, from, to, snap)?;
156            cut.push((v, range));
157            Span {
158                start: t.seconds(start),
159                end: t.seconds(end),
160            }
161        }
162        None => Span {
163            start: from,
164            end: to,
165        },
166    };
167    for i in audio {
168        cut.push((i, audio_range(&tracks[i], span.start, span.end)));
169    }
170    Ok(Cut { span, tracks: cut })
171}
172
173pub fn plan<S: ByteSource>(reader: &Reader<S>, from: f64, to: f64, snap: Snap) -> Result<Cut> {
174    plan_tracks(
175        reader.duration_seconds(),
176        &reader.tracks,
177        reader.video(),
178        reader.audio(),
179        (from, to),
180        snap,
181    )
182}
183
184pub fn copy_samples<S: ByteSource, W: Write + Seek>(
185    reader: &Reader<S>,
186    track: usize,
187    range: Range<usize>,
188    writer: &mut Writer<W>,
189    dst: usize,
190    shift: i64,
191) -> Result<()> {
192    let mut buf = Vec::new();
193    copy_run(reader, track, range, writer, dst, shift, &mut buf, false)
194}
195
196#[allow(clippy::too_many_arguments)]
197fn copy_run<S: ByteSource, W: Write + Seek>(
198    reader: &Reader<S>,
199    track: usize,
200    range: Range<usize>,
201    writer: &mut Writer<W>,
202    dst: usize,
203    shift: i64,
204    buf: &mut Vec<u8>,
205    dry: bool,
206) -> Result<()> {
207    let t = reader
208        .tracks
209        .get(track)
210        .ok_or_else(|| Error::Invalid(format!("no track {track}")))?;
211    let samples = t
212        .samples
213        .get(range.clone())
214        .ok_or_else(|| Error::Invalid(format!("no samples {range:?} in track {track}")))?;
215    let len = reader.source().len();
216    let mut entry = None;
217    let mut i = 0;
218    while i < samples.len() {
219        let start = samples[i].offset;
220        let mut end = start;
221        let mut j = i;
222        while j < samples.len()
223            && samples[j].offset == end
224            && (j == i || end - start + samples[j].size as u64 <= RUN)
225        {
226            end = end
227                .checked_add(samples[j].size as u64)
228                .filter(|&e| e <= len)
229                .ok_or_else(|| truncated("sample data"))?;
230            j += 1;
231        }
232        buf.resize((end - start) as usize, 0);
233        if !dry {
234            reader.read_span(start, buf)?;
235        }
236        let mut at = 0usize;
237        for s in &samples[i..j] {
238            if entry != Some(s.entry) {
239                writer.use_entry(dst, s.entry)?;
240                entry = Some(s.entry);
241            }
242            let size = s.size as usize;
243            writer.push_sample(
244                dst,
245                &buf[at..at + size],
246                s.pts.saturating_add(shift),
247                s.dts.saturating_add(shift),
248                s.duration,
249                s.keyframe,
250            )?;
251            at += size;
252        }
253        i = j;
254    }
255    Ok(())
256}
257
258struct Sink {
259    pos: u64,
260    end: u64,
261}
262
263impl Write for Sink {
264    fn write(&mut self, data: &[u8]) -> io::Result<usize> {
265        self.pos += data.len() as u64;
266        self.end = self.end.max(self.pos);
267        Ok(data.len())
268    }
269
270    fn flush(&mut self) -> io::Result<()> {
271        Ok(())
272    }
273}
274
275impl Seek for Sink {
276    fn seek(&mut self, to: SeekFrom) -> io::Result<u64> {
277        self.pos = match to {
278            SeekFrom::Start(p) => p,
279            SeekFrom::Current(d) => self.pos.saturating_add_signed(d),
280            SeekFrom::End(d) => self.end.saturating_add_signed(d),
281        };
282        Ok(self.pos)
283    }
284}
285
286fn feed<S: ByteSource, W: Write + Seek>(
287    reader: &Reader<S>,
288    cut: &Cut,
289    writer: &mut Writer<W>,
290    dry: bool,
291) -> Result<()> {
292    let segment = ((cut.span.end - cut.span.start) * 1000.0).round().max(0.0) as u64;
293    let mut lanes = Vec::new();
294    for (ti, range) in &cut.tracks {
295        let t = &reader.tracks[*ti];
296        if range.is_empty() || !matches!(t.kind, TrackKind::Video | TrackKind::Audio) {
297            continue;
298        }
299        let mut cfg = TrackConfig::from_track(t)?;
300        let origin = t.samples[range.start].dts;
301        let start = t
302            .ticks(cut.span.start)
303            .saturating_sub(t.edit_offset)
304            .saturating_sub(origin);
305        cfg.edits = vec![Edit {
306            segment_duration: segment,
307            media_time: start.max(0),
308            media_rate: 0x10000,
309        }];
310        let dst = writer.add_track(cfg)?;
311        for e in t.entries.iter().skip(1) {
312            writer.add_entry(dst, e.clone())?;
313        }
314        lanes.push((*ti, range.clone(), dst, origin, range.start));
315    }
316    let mut buf = Vec::new();
317    let mut window = INTERLEAVE;
318    loop {
319        let mut left = false;
320        let mut ahead = f64::INFINITY;
321        for (ti, range, dst, origin, next) in &mut lanes {
322            let t = &reader.tracks[*ti];
323            let stop = (*next..range.end)
324                .find(|&i| t.seconds(t.samples[i].dts.saturating_sub(*origin)) >= window)
325                .unwrap_or(range.end);
326            if stop > *next {
327                copy_run(
328                    reader,
329                    *ti,
330                    *next..stop,
331                    writer,
332                    *dst,
333                    -*origin,
334                    &mut buf,
335                    dry,
336                )?;
337                *next = stop;
338            }
339            if *next < range.end {
340                left = true;
341                ahead = ahead.min(t.seconds(t.samples[*next].dts.saturating_sub(*origin)));
342            }
343        }
344        if !left {
345            return Ok(());
346        }
347        window = (window + INTERLEAVE).max(ahead);
348    }
349}
350
351pub fn remux(
352    src: impl AsRef<Path>,
353    dst: impl AsRef<Path>,
354    from: f64,
355    to: f64,
356    snap: Snap,
357) -> Result<Span> {
358    let reader = Reader::open(src)?;
359    let cut = plan(&reader, from, to, snap)?;
360    let brand = if reader.quicktime {
361        Brand::Mov
362    } else {
363        Brand::Mp4
364    };
365    let mut dry = Writer::new(Sink { pos: 0, end: 0 }, Options { brand, reserve: 0 })?;
366    feed(&reader, &cut, &mut dry, true)?;
367    let moov = dry.moov_len()?;
368    let reserve = moov
369        + 8
370        + if dry.position() > u32::MAX as u64 / 2 {
371            moov
372        } else {
373            0
374        };
375    let mut writer = Writer::create_with(dst, Options { brand, reserve })?;
376    feed(&reader, &cut, &mut writer, false)?;
377    writer.finish()?;
378    Ok(cut.span)
379}
pito_codecs/mp4/remux.rs

pito_codecs/mp4/
remux.rs
