1use std::fs::File;
2use std::io::{Seek, SeekFrom, Write};
3use std::path::Path;
4
5use super::ebml::{write_id, write_size, write_unknown_size};
6use super::ids::*;
7use super::track::Track;
8use crate::media::{Error, Packet, Result, TimeBase};
9use crate::mp4::{CodecConfig, TrackKind};
10
11const SEEK_HEAD_SPACE: usize = 128;
12const SCALE: i64 = 1_000_000;
13const KEYFRAME_CLUSTER_NS: i64 = 500_000_000;
14
15#[derive(Clone, Copy, Debug, PartialEq)]
16pub struct Options {
17 pub webm: bool,
18 pub cluster_seconds: f64,
19}
20
21impl Default for Options {
22 fn default() -> Self {
23 Self {
24 webm: false,
25 cluster_seconds: 5.0,
26 }
27 }
28}
29
30impl Options {
31 pub fn webm() -> Self {
32 Self {
33 webm: true,
34 ..Self::default()
35 }
36 }
37}
38
39#[derive(Clone, Debug, PartialEq)]
40pub struct TrackConfig {
41 pub kind: TrackKind,
42 pub codec_id: String,
43 pub codec_private: Vec<u8>,
44 pub time_base: TimeBase,
45 pub width: u32,
46 pub height: u32,
47 pub rotation: u16,
48 pub sample_rate: u32,
49 pub channels: u32,
50 pub default_duration: Option<u64>,
51 pub codec_delay: u64,
52 pub seek_pre_roll: u64,
53 pub language: String,
54 pub name: Option<String>,
55}
56
57impl TrackConfig {
58 pub fn new(kind: TrackKind, codec_id: &str, time_base: TimeBase) -> Self {
59 Self {
60 kind,
61 codec_id: codec_id.into(),
62 codec_private: Vec::new(),
63 time_base,
64 width: 0,
65 height: 0,
66 rotation: 0,
67 sample_rate: 0,
68 channels: 0,
69 default_duration: None,
70 codec_delay: 0,
71 seek_pre_roll: 0,
72 language: "und".into(),
73 name: None,
74 }
75 }
76
77 fn video(
78 codec_id: &str,
79 private: Vec<u8>,
80 width: u32,
81 height: u32,
82 time_base: TimeBase,
83 ) -> Self {
84 Self {
85 codec_private: private,
86 width,
87 height,
88 ..Self::new(TrackKind::Video, codec_id, time_base)
89 }
90 }
91
92 pub fn h264(avcc: Vec<u8>, width: u32, height: u32, time_base: TimeBase) -> Self {
93 Self::video("V_MPEG4/ISO/AVC", avcc, width, height, time_base)
94 }
95
96 pub fn hevc(hvcc: Vec<u8>, width: u32, height: u32, time_base: TimeBase) -> Self {
97 Self::video("V_MPEGH/ISO/HEVC", hvcc, width, height, time_base)
98 }
99
100 pub fn vp9(width: u32, height: u32, time_base: TimeBase) -> Self {
101 Self::video("V_VP9", Vec::new(), width, height, time_base)
102 }
103
104 pub fn opus(head: Vec<u8>) -> Self {
105 let channels = head.get(9).copied().unwrap_or(2) as u32;
106 let pre_skip = head
107 .get(10..12)
108 .map_or(0, |b| u16::from_le_bytes([b[0], b[1]]) as u64);
109 Self {
110 codec_private: head,
111 sample_rate: 48000,
112 channels,
113 codec_delay: pre_skip * 1_000_000_000 / 48000,
114 seek_pre_roll: 80_000_000,
115 ..Self::new(TrackKind::Audio, "A_OPUS", TimeBase::new(1, 48000))
116 }
117 }
118
119 pub fn aac(asc: Vec<u8>, sample_rate: u32, channels: u32) -> Self {
120 Self {
121 codec_private: asc,
122 sample_rate,
123 channels,
124 ..Self::new(
125 TrackKind::Audio,
126 "A_AAC",
127 TimeBase::new(1, sample_rate.max(1)),
128 )
129 }
130 }
131
132 pub fn from_track(track: &Track) -> Self {
133 Self {
134 kind: track.kind,
135 codec_id: track.codec_id.clone(),
136 codec_private: track.extradata().unwrap_or_default(),
137 time_base: track.time_base,
138 width: track.width,
139 height: track.height,
140 rotation: track.rotation(),
141 sample_rate: track.sample_rate(),
142 channels: track.channels(),
143 default_duration: track.default_duration,
144 codec_delay: track.codec_delay,
145 seek_pre_roll: track.seek_pre_roll,
146 language: track.language.clone(),
147 name: (!track.name.is_empty()).then(|| track.name.clone()),
148 }
149 }
150
151 pub fn from_mp4(track: &crate::mp4::Track) -> Result<Self> {
152 let time_base = TimeBase::new(1, track.timescale.max(1));
153 let entry = track
154 .entry()
155 .ok_or_else(|| Error::Invalid(format!("track {} has no sample entry", track.id)))?;
156 let (w, h) = entry.video.map_or((track.width, track.height), |v| {
157 (v.width as u32, v.height as u32)
158 });
159 let mut cfg = match &entry.codec {
160 CodecConfig::Avc(c) => Self::h264(c.to_bytes(), w, h, time_base),
161 CodecConfig::Hevc(c) => Self::hevc(c.to_bytes(), w, h, time_base),
162 CodecConfig::Aac(c) => {
163 let mut cfg = Self::aac(c.asc.clone(), track.sample_rate(), track.channels());
164 cfg.time_base = time_base;
165 cfg
166 }
167 other => {
168 return Err(Error::Unsupported(format!(
169 "{} in Matroska from MP4",
170 other.name()
171 )));
172 }
173 };
174 cfg.rotation = track.rotation();
175 cfg.language = track.language.clone();
176 cfg.name = (!track.name.is_empty()).then(|| track.name.clone());
177 Ok(cfg)
178 }
179
180 fn nanos(&self, ticks: i64) -> i64 {
181 (ticks as i128 * self.time_base.num as i128 * 1_000_000_000
182 / self.time_base.den.max(1) as i128) as i64
183 }
184}
185
186fn round_ticks(ns: i64) -> i64 {
187 if ns >= 0 {
188 ns.saturating_add(SCALE / 2) / SCALE
189 } else {
190 ns.saturating_sub(SCALE / 2) / SCALE
191 }
192}
193
194fn webm_codec(id: &str) -> bool {
195 matches!(
196 id,
197 "V_VP8" | "V_VP9" | "V_AV1" | "A_OPUS" | "A_VORBIS" | "S_TEXT/WEBVTT"
198 ) || id.starts_with("D_WEBVTT/")
199}
200
201pub(crate) fn el(out: &mut Vec<u8>, id: u32, data: &[u8]) {
202 write_id(out, id);
203 write_size(out, data.len() as u64, 1);
204 out.extend_from_slice(data);
205}
206
207pub(crate) fn el_uint(out: &mut Vec<u8>, id: u32, v: u64) {
208 let n = (8 - v.leading_zeros() as usize / 8).max(1);
209 el(out, id, &v.to_be_bytes()[8 - n..]);
210}
211
212pub(crate) fn el_float(out: &mut Vec<u8>, id: u32, v: f64) {
213 el(out, id, &v.to_be_bytes());
214}
215
216pub(crate) fn el_str(out: &mut Vec<u8>, id: u32, s: &str) {
217 el(out, id, s.as_bytes());
218}
219
220fn void(out: &mut Vec<u8>, total: usize) {
221 let sl = if total - 2 <= 126 { 1 } else { 8 };
222 write_id(out, VOID);
223 write_size(out, (total - 1 - sl) as u64, sl);
224 out.resize(out.len() + total - 1 - sl, 0);
225}
226
227fn track_entry(t: &TrackConfig, number: u64) -> Vec<u8> {
228 let mut e = Vec::new();
229 el_uint(&mut e, TRACK_NUMBER, number);
230 el_uint(&mut e, TRACK_UID, number);
231 let kind = match t.kind {
232 TrackKind::Video => 1,
233 TrackKind::Audio => 2,
234 TrackKind::Subtitle => 0x11,
235 TrackKind::Other => 0x21,
236 };
237 el_uint(&mut e, TRACK_TYPE, kind);
238 el_uint(&mut e, FLAG_LACING, 0);
239 el_str(&mut e, LANGUAGE, &t.language);
240 if let Some(n) = &t.name {
241 el_str(&mut e, NAME, n);
242 }
243 el_str(&mut e, CODEC_ID, &t.codec_id);
244 if !t.codec_private.is_empty() {
245 el(&mut e, CODEC_PRIVATE, &t.codec_private);
246 }
247 if let Some(d) = t.default_duration {
248 el_uint(&mut e, DEFAULT_DURATION, d);
249 }
250 if t.codec_delay > 0 {
251 el_uint(&mut e, CODEC_DELAY, t.codec_delay);
252 }
253 if t.seek_pre_roll > 0 {
254 el_uint(&mut e, SEEK_PRE_ROLL, t.seek_pre_roll);
255 }
256 if t.kind == TrackKind::Video {
257 let mut v = Vec::new();
258 el_uint(&mut v, PIXEL_WIDTH, t.width as u64);
259 el_uint(&mut v, PIXEL_HEIGHT, t.height as u64);
260 if !t.rotation.is_multiple_of(360) {
261 let mut p = Vec::new();
262 el_uint(&mut p, PROJECTION_TYPE, 0);
263 el_float(&mut p, PROJECTION_POSE_ROLL, -((t.rotation % 360) as f64));
264 el(&mut v, PROJECTION, &p);
265 }
266 el(&mut e, VIDEO, &v);
267 }
268 if t.kind == TrackKind::Audio {
269 let mut a = Vec::new();
270 el_float(&mut a, SAMPLING_FREQUENCY, t.sample_rate as f64);
271 el_uint(&mut a, CHANNELS, t.channels.max(1) as u64);
272 el(&mut e, AUDIO, &a);
273 }
274 e
275}
276
277struct WTrack {
278 cfg: TrackConfig,
279 last_ns: Option<i64>,
280 delta_ns: i64,
281}
282
283struct Layout {
284 segment_size: u64,
285 segment_start: u64,
286 duration: u64,
287 info: u64,
288 tracks: u64,
289}
290
291pub struct Writer<W: Write + Seek = File> {
292 w: W,
293 options: Options,
294 tracks: Vec<WTrack>,
295 layout: Option<Layout>,
296 has_video: bool,
297 cluster: Vec<u8>,
298 cluster_ts: Option<i64>,
299 cluster_pos: u64,
300 cues: Vec<(u64, u64, u64)>,
301 end_ns: i64,
302}
303
304impl Writer<File> {
305 pub fn create(path: impl AsRef<Path>) -> Result<Self> {
306 Self::create_with(path, Options::default())
307 }
308
309 pub fn create_with(path: impl AsRef<Path>, options: Options) -> Result<Self> {
310 let path = path.as_ref();
311 let f = File::create(path).map_err(|e| Error::at(path, e))?;
312 Ok(Self::new(f, options))
313 }
314}
315
316impl<W: Write + Seek> Writer<W> {
317 pub fn new(w: W, options: Options) -> Self {
318 Self {
319 w,
320 options,
321 tracks: Vec::new(),
322 layout: None,
323 has_video: false,
324 cluster: Vec::with_capacity(1 << 20),
325 cluster_ts: None,
326 cluster_pos: 0,
327 cues: Vec::new(),
328 end_ns: 0,
329 }
330 }
331
332 pub fn add_track(&mut self, cfg: TrackConfig) -> Result<usize> {
333 if self.layout.is_some() {
334 return Err(Error::Invalid(
335 "tracks must be added before the first packet".into(),
336 ));
337 }
338 if self.options.webm && !webm_codec(&cfg.codec_id) {
339 return Err(Error::Unsupported(format!(
340 "WebM carries VP8, VP9, AV1, Opus and Vorbis, not {}; write Matroska",
341 super::track::codec_name(&cfg.codec_id)
342 )));
343 }
344 if cfg.time_base.num == 0 || cfg.time_base.den == 0 {
345 return Err(Error::Invalid("track time base is zero".into()));
346 }
347 if self.tracks.len() >= 126 {
348 return Err(Error::Unsupported("more than 126 tracks".into()));
349 }
350 self.has_video |= cfg.kind == TrackKind::Video;
351 self.tracks.push(WTrack {
352 cfg,
353 last_ns: None,
354 delta_ns: 0,
355 });
356 Ok(self.tracks.len() - 1)
357 }
358
359 fn start(&mut self) -> Result<()> {
360 if self.layout.is_some() {
361 return Ok(());
362 }
363 if self.tracks.is_empty() {
364 return Err(Error::Invalid("no tracks".into()));
365 }
366 let mut out = Vec::new();
367 let mut h = Vec::new();
368 el_uint(&mut h, EBML_VERSION, 1);
369 el_uint(&mut h, EBML_READ_VERSION, 1);
370 el_uint(&mut h, EBML_MAX_ID_LENGTH, 4);
371 el_uint(&mut h, EBML_MAX_SIZE_LENGTH, 8);
372 el_str(
373 &mut h,
374 DOC_TYPE,
375 if self.options.webm {
376 "webm"
377 } else {
378 "matroska"
379 },
380 );
381 el_uint(&mut h, DOC_TYPE_VERSION, 4);
382 el_uint(&mut h, DOC_TYPE_READ_VERSION, 2);
383 el(&mut out, EBML, &h);
384 write_id(&mut out, SEGMENT);
385 let base = self.w.stream_position()?;
386 let segment_size = base + out.len() as u64;
387 write_unknown_size(&mut out, 8);
388 let segment_start = base + out.len() as u64;
389 void(&mut out, SEEK_HEAD_SPACE);
390 let info_pos = base + out.len() as u64;
391 let mut info = Vec::new();
392 el_uint(&mut info, TIMESTAMP_SCALE, SCALE as u64);
393 el_str(&mut info, MUXING_APP, "pito-codecs");
394 el_str(&mut info, WRITING_APP, "pito-codecs");
395 let duration_in_info = info.len() + 3;
396 el_float(&mut info, DURATION, 0.0);
397 write_id(&mut out, INFO);
398 write_size(&mut out, info.len() as u64, 1);
399 let duration = base + out.len() as u64 + duration_in_info as u64;
400 out.extend_from_slice(&info);
401 let tracks_pos = base + out.len() as u64;
402 let mut tr = Vec::new();
403 for (i, t) in self.tracks.iter().enumerate() {
404 el(&mut tr, TRACK_ENTRY, &track_entry(&t.cfg, i as u64 + 1));
405 }
406 el(&mut out, TRACKS, &tr);
407 self.w.write_all(&out)?;
408 self.layout = Some(Layout {
409 segment_size,
410 segment_start,
411 duration,
412 info: info_pos,
413 tracks: tracks_pos,
414 });
415 Ok(())
416 }
417
418 pub fn push(&mut self, track: usize, packet: &Packet) -> Result<()> {
419 self.push_sample(track, &packet.data, packet.pts, None, packet.keyframe)
420 }
421
422 pub fn push_timed(&mut self, track: usize, packet: &Packet, duration: u32) -> Result<()> {
423 self.push_sample(
424 track,
425 &packet.data,
426 packet.pts,
427 Some(duration),
428 packet.keyframe,
429 )
430 }
431
432 pub fn push_sample(
433 &mut self,
434 track: usize,
435 data: &[u8],
436 pts: i64,
437 duration: Option<u32>,
438 keyframe: bool,
439 ) -> Result<()> {
440 self.start()?;
441 let t = self
442 .tracks
443 .get_mut(track)
444 .ok_or_else(|| Error::Invalid(format!("no track {track}")))?;
445 let pts_ns = t.cfg.nanos(pts);
446 let ts = round_ticks(pts_ns) + round_ticks(t.cfg.codec_delay as i64);
447 if ts < 0 {
448 return Err(Error::Invalid(format!(
449 "negative timestamp {pts} in track {track}"
450 )));
451 }
452 let duration_ns = duration.map(|d| t.cfg.nanos(d as i64));
453 if let Some(last) = t.last_ns
454 && pts_ns > last
455 {
456 t.delta_ns = pts_ns.saturating_sub(last);
457 }
458 t.last_ns = Some(t.last_ns.map_or(pts_ns, |l| l.max(pts_ns)));
459 let length = duration_ns
460 .or(t.cfg.default_duration.map(|d| d as i64))
461 .unwrap_or(t.delta_ns);
462 self.end_ns = self.end_ns.max(pts_ns.saturating_add(length));
463 let video = t.cfg.kind == TrackKind::Video;
464 let explicit = duration_ns.filter(|&d| Some(d as u64) != t.cfg.default_duration);
465 let start_new = match self.cluster_ts {
466 None => true,
467 Some(cts) => {
468 let age = (ts - cts).saturating_mul(SCALE);
469 let limit = (self.options.cluster_seconds * 1e9) as i64;
470 let rel = ts - cts;
471 !(-32768..=32767).contains(&rel)
472 || (video && keyframe && age >= KEYFRAME_CLUSTER_NS)
473 || (!self.has_video && age >= limit)
474 }
475 };
476 if start_new {
477 self.flush_cluster()?;
478 self.cluster_ts = Some(ts);
479 self.cluster_pos = self.w.stream_position()?;
480 el_uint(&mut self.cluster, TIMESTAMP, ts as u64);
481 if !self.has_video {
482 self.cue(ts, track)?;
483 }
484 }
485 if video && keyframe {
486 self.cue(ts, track)?;
487 }
488 let rel = (ts - self.cluster_ts.unwrap_or(ts)) as i16;
489 let number = track as u8 + 1;
490 match explicit {
491 None => {
492 write_id(&mut self.cluster, SIMPLE_BLOCK);
493 write_size(&mut self.cluster, data.len() as u64 + 4, 1);
494 self.cluster.push(0x80 | number);
495 self.cluster.extend_from_slice(&rel.to_be_bytes());
496 self.cluster.push(if keyframe { 0x80 } else { 0 });
497 self.cluster.extend_from_slice(data);
498 }
499 Some(d) => {
500 let mut tail = Vec::with_capacity(16);
501 el_uint(
502 &mut tail,
503 BLOCK_DURATION,
504 (d.saturating_add(SCALE / 2) / SCALE) as u64,
505 );
506 if !keyframe {
507 el(&mut tail, REFERENCE_BLOCK, &[0xFF]);
508 }
509 let block_len = data.len() as u64 + 4;
510 let mut head = Vec::with_capacity(16);
511 write_id(&mut head, BLOCK);
512 write_size(&mut head, block_len, 1);
513 write_id(&mut self.cluster, BLOCK_GROUP);
514 write_size(
515 &mut self.cluster,
516 (head.len() as u64 + block_len) + tail.len() as u64,
517 1,
518 );
519 self.cluster.extend_from_slice(&head);
520 self.cluster.push(0x80 | number);
521 self.cluster.extend_from_slice(&rel.to_be_bytes());
522 self.cluster.push(0);
523 self.cluster.extend_from_slice(data);
524 self.cluster.extend_from_slice(&tail);
525 }
526 }
527 Ok(())
528 }
529
530 fn cue(&mut self, ts: i64, track: usize) -> Result<()> {
531 let Some(l) = &self.layout else {
532 return Ok(());
533 };
534 let pos = self.cluster_pos - l.segment_start;
535 let number = track as u64 + 1;
536 if self.cues.last().is_none_or(|c| c.2 != pos || c.1 != number) {
537 self.cues.push((ts as u64, number, pos));
538 }
539 Ok(())
540 }
541
542 fn flush_cluster(&mut self) -> Result<()> {
543 if self.cluster_ts.is_some() {
544 let mut head = Vec::with_capacity(12);
545 write_id(&mut head, CLUSTER);
546 write_size(&mut head, self.cluster.len() as u64, 8);
547 self.w.write_all(&head)?;
548 self.w.write_all(&self.cluster)?;
549 }
550 self.cluster.clear();
551 self.cluster_ts = None;
552 Ok(())
553 }
554
555 pub fn finish(mut self) -> Result<W> {
556 self.start()?;
557 self.flush_cluster()?;
558 let Some(l) = self.layout.take() else {
559 return Err(Error::Invalid("no tracks".into()));
560 };
561 let cues_pos = self.w.stream_position()?;
562 let mut cues = Vec::new();
563 for &(time, track, pos) in &self.cues {
564 let mut ctp = Vec::new();
565 el_uint(&mut ctp, CUE_TRACK, track);
566 el_uint(&mut ctp, CUE_CLUSTER_POSITION, pos);
567 let mut cp = Vec::new();
568 el_uint(&mut cp, CUE_TIME, time);
569 el(&mut cp, CUE_TRACK_POSITIONS, &ctp);
570 el(&mut cues, CUE_POINT, &cp);
571 }
572 if !cues.is_empty() {
573 let mut out = Vec::new();
574 el(&mut out, CUES, &cues);
575 self.w.write_all(&out)?;
576 }
577 let end = self.w.stream_position()?;
578 let mut sh = Vec::new();
579 let mut entries = vec![(INFO, l.info), (TRACKS, l.tracks)];
580 if !cues.is_empty() {
581 entries.push((CUES, cues_pos));
582 }
583 for (id, pos) in entries {
584 let mut s = Vec::new();
585 let mut idb = Vec::new();
586 write_id(&mut idb, id);
587 el(&mut s, SEEK_ID, &idb);
588 el_uint(&mut s, SEEK_POSITION, pos - l.segment_start);
589 el(&mut sh, SEEK, &s);
590 }
591 let mut out = Vec::new();
592 el(&mut out, SEEK_HEAD, &sh);
593 let rest = SEEK_HEAD_SPACE - out.len();
594 void(&mut out, rest);
595 self.w.seek(SeekFrom::Start(l.segment_start))?;
596 self.w.write_all(&out)?;
597 let duration = self.end_ns.max(0) as f64 / SCALE as f64;
598 self.w.seek(SeekFrom::Start(l.duration))?;
599 self.w.write_all(&duration.to_be_bytes())?;
600 let mut size = Vec::new();
601 write_size(&mut size, end - l.segment_start, 8);
602 self.w.seek(SeekFrom::Start(l.segment_size))?;
603 self.w.write_all(&size)?;
604 self.w.seek(SeekFrom::Start(end))?;
605 self.w.flush()?;
606 Ok(self.w)
607 }
608}