pito_codecs/mkv/
remux.rs

1use std::io::{Seek, Write};
2use std::path::Path;
3
4use super::reader::Reader;
5use super::track::Track;
6use super::writer::{Options, TrackConfig, Writer};
7use crate::media::{Error, Result};
8use crate::mp4::remux::{Timeline, plan_tracks};
9use crate::mp4::{ByteSource, Cut, Snap, Span, TrackKind};
10
11const WINDOW: u64 = 8 << 20;
12
13impl Timeline for Track {
14    fn len(&self) -> usize {
15        self.samples.len()
16    }
17    fn keyframes(&self) -> &[u32] {
18        Track::keyframes(self)
19    }
20    fn pts(&self, i: usize) -> i64 {
21        self.samples[i].pts
22    }
23    fn duration(&self, i: usize) -> i64 {
24        self.samples[i].duration as i64
25    }
26    fn ticks(&self, seconds: f64) -> i64 {
27        Track::ticks(self, seconds)
28    }
29    fn seconds(&self, ticks: i64) -> f64 {
30        Track::seconds(self, ticks)
31    }
32    fn end(&self) -> i64 {
33        Track::end(self)
34    }
35}
36
37pub fn plan<S: ByteSource>(reader: &Reader<S>, from: f64, to: f64, snap: Snap) -> Result<Cut> {
38    plan_tracks(
39        reader.duration_seconds(),
40        &reader.tracks,
41        reader.video(),
42        reader.audio(),
43        (from, to),
44        snap,
45    )
46}
47
48pub fn copy<S: ByteSource, W: Write + Seek>(
49    reader: &Reader<S>,
50    cut: &Cut,
51    writer: &mut Writer<W>,
52) -> Result<()> {
53    let mut map = Vec::new();
54    let mut origin_ns = i64::MAX;
55    for (src, range) in &cut.tracks {
56        let t = reader
57            .tracks
58            .get(*src)
59            .ok_or_else(|| Error::Invalid(format!("no track {src}")))?;
60        let samples = t
61            .samples
62            .get(range.clone())
63            .ok_or_else(|| Error::Invalid(format!("no samples {range:?} in track {src}")))?;
64        map.push(writer.add_track(TrackConfig::from_track(t))?);
65        if let Some(first) = samples.iter().map(|s| s.pts).min() {
66            origin_ns = origin_ns.min(t.nanos(first).saturating_add(t.codec_delay as i64));
67        }
68    }
69    if origin_ns == i64::MAX {
70        origin_ns = 0;
71    }
72    let mut order: Vec<(u64, usize, usize)> = cut
73        .tracks
74        .iter()
75        .enumerate()
76        .flat_map(|(k, (src, range))| {
77            let t = &reader.tracks[*src];
78            range.clone().map(move |i| (t.samples[i].offset, k, i))
79        })
80        .collect();
81    order.sort_unstable();
82    let src = reader.source();
83    let len = reader.len();
84    let mut window = Vec::new();
85    let mut window_start = 0u64;
86    let mut frame = Vec::new();
87    for (offset, k, i) in order {
88        let t = &reader.tracks[cut.tracks[k].0];
89        let s = t.samples[i];
90        let shift =
91            origin_ns as i128 * t.time_base.den as i128 / (t.time_base.num as i128 * 1_000_000_000);
92        let pts = s.pts.saturating_sub(shift as i64);
93        let prefix = t.stripped_header().unwrap_or(&[]);
94        let end = offset
95            .checked_add(s.size as u64)
96            .filter(|&e| e <= len)
97            .ok_or_else(|| Error::Invalid("truncated Matroska data".into()))?;
98        let data: &[u8] = if prefix.is_empty() && s.size as u64 <= WINDOW {
99            if offset < window_start || end > window_start + window.len() as u64 {
100                let n = (len.saturating_sub(offset)).min(WINDOW.max(s.size as u64));
101                window.resize(n as usize, 0);
102                src.read_at(offset, &mut window)?;
103                window_start = offset;
104            }
105            let at = (offset - window_start) as usize;
106            &window[at..at + s.size as usize]
107        } else {
108            super::reader::read_frame(src, len, t, offset, s.size, &mut frame)?;
109            &frame
110        };
111        writer.push_sample(map[k], data, pts, None, s.keyframe)?;
112    }
113    Ok(())
114}
115
116pub fn remux(
117    src: impl AsRef<Path>,
118    dst: impl AsRef<Path>,
119    from: f64,
120    to: f64,
121    snap: Snap,
122) -> Result<Span> {
123    let reader = Reader::open(src)?;
124    let cut = plan(&reader, from, to, snap)?;
125    let options = Options {
126        webm: reader.is_webm(),
127        ..Options::default()
128    };
129    let mut writer = Writer::create_with(dst, options)?;
130    copy(&reader, &cut, &mut writer)?;
131    writer.finish()?;
132    Ok(cut.span)
133}
134
135pub fn remux_mp4(src: impl AsRef<Path>, dst: impl AsRef<Path>) -> Result<()> {
136    let reader = crate::mp4::Reader::open(src)?;
137    let mut writer = Writer::create(dst)?;
138    let mut tracks = Vec::new();
139    for (i, t) in reader.tracks.iter().enumerate() {
140        if matches!(t.kind, TrackKind::Video | TrackKind::Audio) {
141            tracks.push((i, writer.add_track(TrackConfig::from_mp4(t)?)?));
142        }
143    }
144    let mut order: Vec<(u64, usize, usize)> = tracks
145        .iter()
146        .enumerate()
147        .flat_map(|(k, &(i, _))| {
148            reader.tracks[i]
149                .samples
150                .iter()
151                .enumerate()
152                .map(move |(n, s)| (s.offset, k, n))
153        })
154        .collect();
155    order.sort_unstable();
156    let shift: Vec<i64> = tracks
157        .iter()
158        .map(|&(i, _)| {
159            let t = &reader.tracks[i];
160            (0..t.samples.len())
161                .filter_map(|n| t.presentation(n))
162                .min()
163                .unwrap_or(0)
164                .min(0)
165        })
166        .collect();
167    let mut data = Vec::new();
168    for (_, k, n) in order {
169        let (i, dst) = tracks[k];
170        let t = &reader.tracks[i];
171        reader.read_into(i, n, &mut data)?;
172        let pts = t.presentation(n).unwrap_or(0).saturating_sub(shift[k]);
173        writer.push_sample(dst, &data, pts, None, t.samples[n].keyframe)?;
174    }
175    writer.finish()?;
176    Ok(())
177}
pito_codecs/mkv/remux.rs

pito_codecs/mkv/
remux.rs
