1use std::{
2 fmt::{self, Debug},
3 ops::{Range, RangeBounds},
4 time::Duration,
5};
6
7use crate::{
8 atom::{
9 atom_ref::{unwrap_atom_data, AtomRef, AtomRefMut},
10 elst::EditEntry,
11 stsd::{
12 AudioSampleEntry, BtrtExtension, DecoderSpecificInfo, EsdsExtension, SampleEntry,
13 SampleEntryData, SampleEntryType, StsdExtension,
14 },
15 tkhd::TKHD,
16 tref::TREF,
17 util::{scaled_duration_range, unscaled_duration},
18 EdtsAtomRef, EdtsAtomRefMut, MdiaAtomRef, MdiaAtomRefMut, TrackHeaderAtom,
19 TrackReferenceAtom, EDTS, MDIA,
20 },
21 Atom, AtomData, FourCC,
22};
23
24pub const TRAK: FourCC = FourCC::new(b"trak");
25
26#[derive(Clone, Copy)]
27pub struct TrakAtomRef<'a>(AtomRef<'a>);
28
29impl fmt::Debug for TrakAtomRef<'_> {
30 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
31 f.debug_struct("TrakAtomRef")
32 .field("track_id", &self.header().unwrap().track_id)
33 .finish()
34 }
35}
36
37impl<'a> TrakAtomRef<'a> {
38 pub(crate) fn new(atom: &'a Atom) -> Self {
39 Self(AtomRef(Some(atom)))
40 }
41
42 pub fn children(&self) -> impl Iterator<Item = &'a Atom> {
43 self.0.children()
44 }
45
46 pub fn header(&self) -> Option<&'a TrackHeaderAtom> {
48 let atom = self.0.find_child(TKHD)?;
49 match atom.data.as_ref()? {
50 AtomData::TrackHeader(data) => Some(data),
51 _ => None,
52 }
53 }
54
55 pub fn edit_list_container(&self) -> EdtsAtomRef<'a> {
56 EdtsAtomRef(AtomRef(self.0.find_child(EDTS)))
57 }
58
59 pub fn media(&self) -> MdiaAtomRef<'a> {
61 MdiaAtomRef(AtomRef(self.0.find_child(MDIA)))
62 }
63
64 pub fn track_id(&self) -> Option<u32> {
65 let tkhd = self.header()?;
66 Some(tkhd.track_id)
67 }
68
69 pub fn size(&self) -> usize {
71 self.media()
72 .media_information()
73 .sample_table()
74 .sample_size()
75 .map_or(0, |s| {
76 if s.entry_sizes.is_empty() {
77 s.sample_size * s.sample_count
78 } else {
79 s.entry_sizes.iter().sum::<u32>()
80 }
81 }) as usize
82 }
83
84 pub fn bitrate(&self) -> Option<u32> {
88 let duration_secds = self
89 .media()
90 .header()
91 .map(|mdhd| (mdhd.duration as f64) / f64::from(mdhd.timescale))?;
92
93 self.media()
94 .media_information()
95 .sample_table()
96 .sample_size()
97 .map(|s| {
98 let num_bits = s
99 .entry_sizes
100 .iter()
101 .map(|s| *s as usize)
102 .sum::<usize>()
103 .saturating_mul(8);
104
105 let bitrate = (num_bits as f64) / duration_secds;
106 bitrate.round() as u32
107 })
108 }
109}
110
111#[derive(Debug)]
112pub struct TrakAtomRefMut<'a>(pub(crate) AtomRefMut<'a>);
113
114impl<'a> TrakAtomRefMut<'a> {
115 pub(crate) fn new(atom: &'a mut Atom) -> Self {
116 Self(AtomRefMut(atom))
117 }
118
119 pub fn as_ref(&self) -> TrakAtomRef<'_> {
120 TrakAtomRef(self.0.as_ref())
121 }
122
123 pub fn into_ref(self) -> TrakAtomRef<'a> {
124 TrakAtomRef(self.0.into_ref())
125 }
126
127 pub fn children(&mut self) -> impl Iterator<Item = &'_ mut Atom> {
128 self.0.children()
129 }
130
131 pub fn header(&mut self) -> &mut TrackHeaderAtom {
133 unwrap_atom_data!(
134 self.0
135 .find_or_insert_child(TKHD)
136 .insert_data(AtomData::TrackHeader(TrackHeaderAtom::default()))
137 .call(),
138 AtomData::TrackHeader,
139 )
140 }
141
142 pub fn media(&mut self) -> MdiaAtomRefMut<'_> {
144 MdiaAtomRefMut(
145 self.0
146 .find_or_insert_child(MDIA)
147 .insert_after(vec![TREF, EDTS, TKHD])
148 .call(),
149 )
150 }
151
152 pub fn into_media(self) -> Option<MdiaAtomRefMut<'a>> {
154 let atom = self.0.into_child(MDIA)?;
155 Some(MdiaAtomRefMut(AtomRefMut(atom)))
156 }
157
158 pub fn edit_list_container(&mut self) -> EdtsAtomRefMut<'_> {
160 EdtsAtomRefMut(
161 self.0
162 .find_or_insert_child(EDTS)
163 .insert_after(vec![TREF, TKHD])
164 .call(),
165 )
166 }
167
168 pub fn track_reference(&mut self) -> &mut TrackReferenceAtom {
170 unwrap_atom_data!(
171 self.0
172 .find_or_insert_child(TREF)
173 .insert_after(vec![TKHD])
174 .insert_data(AtomData::TrackReference(TrackReferenceAtom::default()))
175 .call(),
176 AtomData::TrackReference,
177 )
178 }
179
180 pub fn update_audio_bitrate(&mut self, bitrate: u32) {
184 let mut mdia = self.media();
185 let mut minf = mdia.media_information();
186 let mut stbl = minf.sample_table();
187 let stsd = stbl.sample_description();
188
189 let entry = stsd.find_or_create_entry(
190 |entry| matches!(entry.data, SampleEntryData::Audio(_)),
191 || SampleEntry {
192 entry_type: SampleEntryType::Mp4a,
193 data_reference_index: 0,
194 data: SampleEntryData::Audio(AudioSampleEntry::default()),
195 },
196 );
197
198 entry.entry_type = SampleEntryType::Mp4a;
199
200 if let SampleEntryData::Audio(audio) = &mut entry.data {
201 let mut sample_frequency = None;
202 audio
203 .extensions
204 .retain(|ext| matches!(ext, StsdExtension::Esds(_)));
205 let esds = audio.find_or_create_extension(
206 |ext| matches!(ext, StsdExtension::Esds(_)),
207 || StsdExtension::Esds(EsdsExtension::default()),
208 );
209 if let StsdExtension::Esds(esds) = esds {
210 let cfg = esds
211 .es_descriptor
212 .decoder_config_descriptor
213 .get_or_insert_default();
214 cfg.avg_bitrate = bitrate;
215 cfg.max_bitrate = bitrate;
216 if let Some(DecoderSpecificInfo::Audio(a)) = cfg.decoder_specific_info.as_ref() {
217 sample_frequency = Some(a.sampling_frequency.as_hz());
218 }
219 }
220 audio.extensions.push(StsdExtension::Btrt(BtrtExtension {
221 buffer_size_db: 0,
222 avg_bitrate: bitrate,
223 max_bitrate: bitrate,
224 }));
225
226 if let Some(hz) = sample_frequency {
227 audio.sample_rate = hz as f32;
228 }
229 } else {
230 unreachable!("STSD constructed with invalid data")
232 }
233 }
234}
235
236#[cfg(feature = "experimental-trim")]
237impl<'a> TrakAtomRefMut<'a> {
238 pub(crate) fn trim_duration<R>(&mut self, movie_timescale: u64, trim_ranges: &[R]) -> Duration
240 where
241 R: RangeBounds<Duration> + Clone + Debug,
242 {
243 let mut mdia = self.media();
244 let media_timescale = u64::from(mdia.header().timescale);
245 let media_duration = mdia.header().duration;
246 let mut minf = mdia.media_information();
247 let mut stbl = minf.sample_table();
248
249 let scaled_ranges = trim_ranges
251 .iter()
252 .cloned()
253 .map(|range| {
254 convert_range(
255 media_duration,
256 scaled_duration_range(range, media_timescale),
257 )
258 })
259 .collect::<Vec<_>>();
260
261 let (remaining_duration, sample_indices_to_remove) =
263 stbl.time_to_sample().trim_duration(&scaled_ranges);
264
265 let remaining_duration = unscaled_duration(remaining_duration, media_timescale);
266
267 let removed_sample_sizes = stbl
269 .sample_size()
270 .remove_sample_indices(&sample_indices_to_remove);
271
272 let total_chunks = stbl.chunk_offset().chunk_count();
274 let chunk_offset_ops = stbl
275 .sample_to_chunk()
276 .remove_sample_indices(&sample_indices_to_remove, total_chunks);
277
278 let chunk_offsets = &stbl.chunk_offset().chunk_offsets;
280 let chunk_offset_ops = chunk_offset_ops
281 .into_iter()
282 .map(|op| op.resolve(chunk_offsets, &removed_sample_sizes))
283 .collect::<anyhow::Result<Vec<_>>>()
284 .expect("chunk offset ops should only involve removed sample indices and valid chunk indices");
285
286 stbl.chunk_offset().apply_operations(chunk_offset_ops);
288
289 mdia.header().update_duration(|_| remaining_duration);
291 self.header()
292 .update_duration(movie_timescale, |_| remaining_duration);
293
294 if self.as_ref().edit_list_container().edit_list().is_some() {
296 self.edit_list_container()
297 .edit_list()
298 .replace_entries(vec![EditEntry::builder()
299 .movie_timescale(movie_timescale)
300 .segment_duration(remaining_duration)
301 .build()]);
302 }
303
304 remaining_duration
305 }
306}
307
308#[cfg(feature = "experimental-trim")]
309fn convert_range(media_time: u64, range: impl RangeBounds<u64>) -> Range<u64> {
310 use std::ops::Bound;
311 let start = match range.start_bound() {
312 Bound::Included(start) => *start,
313 Bound::Excluded(start) => *start + 1,
314 Bound::Unbounded => 0,
315 };
316 let end = match range.end_bound() {
317 Bound::Included(end) => *end + 1,
318 Bound::Excluded(end) => *end,
319 Bound::Unbounded => media_time,
320 };
321 start..end
322}
323
324#[cfg(feature = "experimental-trim")]
325#[cfg(test)]
326pub(crate) mod trim_tests {
327 use std::{ops::Bound, time::Duration};
328
329 use bon::Builder;
330
331 use crate::atom::{
332 container::{DINF, MDIA, MINF, STBL, TRAK},
333 dref::{DataReferenceAtom, DataReferenceEntry, DREF},
334 gmin::GMIN,
335 hdlr::{HandlerReferenceAtom, HandlerType, HDLR},
336 mdhd::{MediaHeaderAtom, MDHD},
337 smhd::{SoundMediaHeaderAtom, SMHD},
338 stco_co64::{ChunkOffsetAtom, STCO},
339 stsc::{SampleToChunkAtom, SampleToChunkEntry, STSC},
340 stsd::{SampleDescriptionTableAtom, STSD},
341 stsz::{SampleSizeAtom, STSZ},
342 stts::{TimeToSampleAtom, TimeToSampleEntry, STTS},
343 text::TEXT,
344 tkhd::{TrackHeaderAtom, TKHD},
345 util::scaled_duration,
346 Atom, AtomHeader, BaseMediaInfoAtom, TextMediaInfoAtom, TrakAtomRef, TrakAtomRefMut, GMHD,
347 };
348
349 #[bon::builder(finish_fn(name = "build"), state_mod(vis = "pub(crate)"))]
350 pub fn create_test_track(
351 #[builder(getter)] movie_timescale: u32,
352 #[builder(getter)] media_timescale: u32,
353 #[builder(getter)] duration: Duration,
354 handler_reference: Option<HandlerReferenceAtom>,
355 minf_header: Option<Atom>,
356 stsc_entries: Option<Vec<SampleToChunkEntry>>,
357 sample_sizes: Option<Vec<u32>>,
358 ) -> Atom {
359 Atom::builder()
360 .header(AtomHeader::new(*TRAK))
361 .children(vec![
362 Atom::builder()
364 .header(AtomHeader::new(*TKHD))
365 .data(
366 TrackHeaderAtom::builder()
367 .track_id(1)
368 .duration(scaled_duration(duration, movie_timescale as u64))
369 .build(),
370 )
371 .build(),
372 create_test_media(media_timescale, duration)
374 .maybe_handler_reference(handler_reference)
375 .maybe_minf_header(minf_header)
376 .maybe_stsc_entries(stsc_entries)
377 .maybe_sample_sizes(sample_sizes)
378 .build(),
379 ])
380 .build()
381 }
382
383 #[bon::builder(finish_fn(name = "build"))]
384 fn create_test_media(
385 #[builder(start_fn)] media_timescale: u32,
386 #[builder(start_fn)] duration: Duration,
387 handler_reference: Option<HandlerReferenceAtom>,
388 minf_header: Option<Atom>,
389 stsc_entries: Option<Vec<SampleToChunkEntry>>,
390 sample_sizes: Option<Vec<u32>>,
391 ) -> Atom {
392 let handler_reference = handler_reference.unwrap_or_else(|| {
393 HandlerReferenceAtom::builder()
394 .handler_type(HandlerType::Audio)
395 .name("SoundHandler".to_string())
396 .build()
397 });
398
399 let minf_header = minf_header.unwrap_or_else(|| {
400 match &handler_reference.handler_type {
401 HandlerType::Audio => {
402 Atom::builder()
404 .header(AtomHeader::new(*SMHD))
405 .data(SoundMediaHeaderAtom::default())
406 .build()
407 }
408 HandlerType::Text => {
409 Atom::builder()
411 .header(AtomHeader::new(*GMHD))
412 .children(vec![
413 Atom::builder()
414 .header(AtomHeader::new(*GMIN))
415 .data(BaseMediaInfoAtom::default())
416 .build(),
417 Atom::builder()
418 .header(AtomHeader::new(*TEXT))
419 .data(TextMediaInfoAtom::default())
420 .build(),
421 ])
422 .build()
423 }
424 _ => {
425 todo!(
426 "no default minf header for {:?}",
427 &handler_reference.handler_type
428 )
429 }
430 }
431 });
432
433 Atom::builder()
434 .header(AtomHeader::new(*MDIA))
435 .children(vec![
436 Atom::builder()
438 .header(AtomHeader::new(*MDHD))
439 .data(
440 MediaHeaderAtom::builder()
441 .timescale(media_timescale)
442 .duration(scaled_duration(duration, media_timescale as u64))
443 .build(),
444 )
445 .build(),
446 Atom::builder()
448 .header(AtomHeader::new(*HDLR))
449 .data(handler_reference)
450 .build(),
451 create_test_media_info()
453 .media_timescale(media_timescale)
454 .duration(duration)
455 .header(minf_header)
456 .maybe_stsc_entries(stsc_entries)
457 .maybe_sample_sizes(sample_sizes)
458 .build(),
459 ])
460 .build()
461 }
462
463 #[bon::builder(finish_fn(name = "build"))]
464 fn create_test_media_info(
465 media_timescale: u32,
466 duration: Duration,
467 header: Atom,
468 stsc_entries: Option<Vec<SampleToChunkEntry>>,
469 sample_sizes: Option<Vec<u32>>,
470 ) -> Atom {
471 let stsc_entries = stsc_entries.unwrap_or_else(|| {
472 vec![SampleToChunkEntry::builder()
473 .first_chunk(1)
474 .samples_per_chunk(2)
475 .sample_description_index(1)
476 .build()]
477 });
478
479 let sample_sizes = sample_sizes.unwrap_or_else(|| {
480 let total_samples = duration.as_secs() as usize;
482 let sample_size = 256;
483 vec![sample_size; total_samples]
484 });
485
486 Atom::builder()
487 .header(AtomHeader::new(*MINF))
488 .children(vec![
489 header,
490 Atom::builder()
492 .header(AtomHeader::new(*DINF))
493 .children(vec![
494 Atom::builder()
496 .header(AtomHeader::new(*DREF))
497 .data(
498 DataReferenceAtom::builder()
499 .entry(DataReferenceEntry::builder().url("").build())
500 .build(),
501 )
502 .build(),
503 ])
504 .build(),
505 create_test_sample_table()
507 .media_timescale(media_timescale)
508 .stsc_entries(stsc_entries)
509 .sample_sizes(sample_sizes)
510 .build(),
511 ])
512 .build()
513 }
514
515 #[bon::builder(finish_fn(name = "build"))]
516 fn create_test_sample_table(
517 media_timescale: u32,
518 stsc_entries: Vec<SampleToChunkEntry>,
519 sample_sizes: Vec<u32>,
520 #[builder(default = 1000)] mdat_content_offset: u64,
521 ) -> Atom {
522 let total_samples = sample_sizes.len() as u32;
523
524 let chunk_offsets = {
526 let mut chunk_offsets = Vec::new();
527 let mut current_offset = mdat_content_offset;
528 let mut sample_size_index = 0;
529 let mut remaining_samples = total_samples;
530 let mut stsc_iter = stsc_entries.iter().peekable();
531 while let Some(entry) = stsc_iter.next() {
532 let n_chunks = match stsc_iter.peek() {
533 Some(next) => next.first_chunk - entry.first_chunk,
534 None => remaining_samples / entry.samples_per_chunk,
535 };
536
537 let n_samples = entry.samples_per_chunk * n_chunks;
538 remaining_samples = remaining_samples.saturating_sub(n_samples);
539
540 for _ in 0..n_chunks {
541 chunk_offsets.push(current_offset);
542 current_offset += sample_sizes
543 .iter()
544 .skip(sample_size_index)
545 .take(entry.samples_per_chunk as usize)
546 .map(|s| *s as u64)
547 .sum::<u64>();
548 sample_size_index += entry.samples_per_chunk as usize;
549 }
550 }
551
552 chunk_offsets
553 };
554
555 Atom::builder()
556 .header(AtomHeader::new(*STBL))
557 .children(vec![
558 Atom::builder()
560 .header(AtomHeader::new(*STSD))
561 .data(SampleDescriptionTableAtom::default())
562 .build(),
563 Atom::builder()
565 .header(AtomHeader::new(*STTS))
566 .data(
567 TimeToSampleAtom::builder()
568 .entry(
569 TimeToSampleEntry::builder()
570 .sample_count(total_samples)
571 .sample_duration(media_timescale)
572 .build(),
573 )
574 .build(),
575 )
576 .build(),
577 Atom::builder()
579 .header(AtomHeader::new(*STSC))
580 .data(SampleToChunkAtom::from(stsc_entries))
581 .build(),
582 Atom::builder()
584 .header(AtomHeader::new(*STSZ))
585 .data(SampleSizeAtom::builder().entry_sizes(sample_sizes).build())
586 .build(),
587 Atom::builder()
589 .header(AtomHeader::new(*STCO))
590 .data(
591 ChunkOffsetAtom::builder()
592 .chunk_offsets(chunk_offsets)
593 .build(),
594 )
595 .build(),
596 ])
597 .build()
598 }
599
600 #[derive(Debug, Builder)]
601 struct TrimDurationRange {
602 start_bound: Bound<Duration>,
603 end_bound: Bound<Duration>,
604 }
605
606 #[derive(Builder)]
607 struct TrimDurationTestCase<ECO> {
608 #[builder(field)]
609 ranges: Vec<TrimDurationRange>,
610 #[builder(default = 1_000)]
611 movie_timescale: u32,
612 #[builder(default = 10_000)]
613 media_timescale: u32,
614 expected_duration: Duration,
615 expected_chunk_offsets: ECO,
616 }
617
618 impl<ECO, S> TrimDurationTestCaseBuilder<ECO, S>
619 where
620 S: trim_duration_test_case_builder::State,
621 {
622 fn range(mut self, range: TrimDurationRange) -> Self {
623 self.ranges.push(range);
624 self
625 }
626 }
627
628 fn get_chunk_offsets(track: TrakAtomRef) -> Vec<u64> {
629 track
630 .media()
631 .media_information()
632 .sample_table()
633 .chunk_offset()
634 .unwrap()
635 .chunk_offsets
636 .clone()
637 .into_inner()
638 }
639
640 fn test_trim_duration<ECO>(mut track: Atom, test_case: TrimDurationTestCase<ECO>)
641 where
642 ECO: FnOnce(Vec<u64>) -> Vec<u64>,
643 {
644 let mut track = TrakAtomRefMut::new(&mut track);
645 let starting_chunk_offsets = get_chunk_offsets(track.as_ref());
646
647 let trim_ranges = test_case
648 .ranges
649 .into_iter()
650 .map(|r| (r.start_bound, r.end_bound))
651 .collect::<Vec<_>>();
652 let res = track.trim_duration(test_case.movie_timescale as u64, &trim_ranges);
653 assert_eq!(res, test_case.expected_duration);
654
655 let trimmed_chunk_offsets = get_chunk_offsets(track.as_ref());
656 let expected_chunk_offsets = (test_case.expected_chunk_offsets)(starting_chunk_offsets);
657 assert_eq!(
658 trimmed_chunk_offsets, expected_chunk_offsets,
659 "trimmed chunk offsets don't match what's expected"
660 );
661 }
662
663 macro_rules! test_trim_duration {
664 ($(
665 $name:ident {
666 @track(
667 $( $track_field:ident: $track_value:expr ),+,
668 ),
669 $( $field:ident: $value:expr ),+,
670 }
671 )*) => {
672 $(
673 #[test]
674 fn $name() {
675 test_trim_duration!(
676 @inner $($field: $value),+,
677 @track $($track_field: $track_value),+,
678 );
679 }
680 )*
681 };
682
683 (
684 @inner $( $field:ident: $value:expr ),+,
685 @track $( $track_field:ident: $track_value:expr ),+,
686 ) => {
687 let test_case = TrimDurationTestCase::builder()
688 .$( $field($value) ).+
689 .build();
690 let track = create_test_track()
691 .movie_timescale(test_case.movie_timescale)
692 .media_timescale(test_case.media_timescale)
693 .$( $track_field($track_value) ).+
694 .build();
695 test_trim_duration(track, test_case);
696 };
697 }
698
699 mod test_trim_duration {
700 use super::*;
701
702 test_trim_duration!(
703 trim_start_11_seconds {
708 @track(
709 duration: Duration::from_secs(100),
710 ),
711 range: TrimDurationRange::builder()
712 .start_bound(Bound::Included(Duration::from_secs(0)))
713 .end_bound(Bound::Excluded(Duration::from_secs(11))).build(),
714 expected_duration: Duration::from_secs(89),
715 expected_chunk_offsets: |mut orig_offsets: Vec<u64>| {
716 orig_offsets.drain(..5);
719 let removed_sample_size = 256; orig_offsets[0] += removed_sample_size; orig_offsets
722 },
723 }
724 );
725 }
726}