use alloc::vec::Vec;
use broadcast_common::Unpackage;
use crate::error::{Error, Result};
use crate::media::{Fmp4Demux, Media, TimelineOrigin, Track, relative_decode_times};
use crate::pipeline::TrackSpec;
use crate::segmenter::{Segmenter, choose_anchor};
fn anchor_index(media: &Media) -> Result<usize> {
choose_anchor(media.tracks.iter().map(|t| &t.spec.config))
}
fn presentation_times(track: &Track, origin: &TimelineOrigin) -> Vec<i64> {
relative_decode_times(track, origin)
.into_iter()
.zip(&track.samples)
.map(|(dts, s)| dts + s.composition_offset() as i64)
.collect()
}
fn rescale_floor(ticks: u64, from_timescale: u32, to_timescale: u32) -> i64 {
if from_timescale == 0 || from_timescale == to_timescale {
return ticks as i64;
}
((ticks as u128 * to_timescale as u128) / from_timescale as u128) as i64
}
fn rescale_signed_floor(ticks: i64, from_timescale: u32, to_timescale: u32) -> i64 {
if from_timescale == 0 || from_timescale == to_timescale {
return ticks;
}
(ticks as i128 * to_timescale as i128).div_euclid(from_timescale as i128) as i64
}
fn earliest_decode_time_track(media: &Media, dts: &[u128], cursors: &[usize]) -> Option<usize> {
let mut best: Option<usize> = None;
for (ti, track) in media.tracks.iter().enumerate() {
if cursors[ti] >= track.samples.len() {
continue;
}
let Some(current) = best else {
best = Some(ti);
continue;
};
let lhs = dts[ti] * media.tracks[current].spec.timescale.max(1) as u128;
let rhs = dts[current] * track.spec.timescale.max(1) as u128;
if lhs < rhs {
best = Some(ti);
}
}
best
}
impl Media {
pub fn select_tracks(&self, indices: &[usize]) -> Result<Media> {
if indices.is_empty() {
return Err(Error::InvalidInput("select_tracks: empty track selection"));
}
let mut tracks = Vec::with_capacity(indices.len());
for &i in indices {
let t = self.tracks.get(i).ok_or(Error::InvalidInput(
"select_tracks: track index out of range",
))?;
tracks.push(t.clone());
}
Ok(Media::new(tracks, self.movie_timescale))
}
pub fn select_tracks_by<F>(&self, mut keep: F) -> Result<Media>
where
F: FnMut(&Track) -> bool,
{
let tracks: Vec<Track> = self.tracks.iter().filter(|t| keep(t)).cloned().collect();
if tracks.is_empty() {
return Err(Error::InvalidInput(
"select_tracks_by: predicate kept no track",
));
}
Ok(Media::new(tracks, self.movie_timescale))
}
pub fn trim(&self, start: u64, end: u64) -> Result<Media> {
if start >= end {
return Err(Error::InvalidInput("trim: start must be < end"));
}
if self.tracks.is_empty() {
return Err(Error::InvalidInput("trim: media has no tracks"));
}
let anchor = anchor_index(self)?;
let origin = TimelineOrigin::of(&self.tracks);
let anchor_track = &self.tracks[anchor];
let anchor_ts = anchor_track.spec.timescale;
let anchor_pts = presentation_times(anchor_track, &origin);
let anchor_lo = rescale_floor(start, self.movie_timescale, anchor_ts);
let anchor_hi = rescale_floor(end, self.movie_timescale, anchor_ts);
let snapped = anchor_pts
.iter()
.position(|&p| p >= anchor_lo && p < anchor_hi)
.map(|mut idx| {
while idx > 0 && !anchor_track.samples[idx].flags.is_sync {
idx -= 1;
}
idx
});
let mut out_tracks = Vec::with_capacity(self.tracks.len());
let mut kept_any = false;
for (ti, track) in self.tracks.iter().enumerate() {
let ts = track.spec.timescale;
let pts = presentation_times(track, &origin);
let hi = rescale_floor(end, self.movie_timescale, ts);
let start_idx = if ti == anchor {
snapped
} else {
let lo = match snapped {
Some(idx) => rescale_signed_floor(anchor_pts[idx], anchor_ts, ts),
None => rescale_floor(start, self.movie_timescale, ts),
};
pts.iter().position(|&p| p >= lo && p < hi)
};
let mut kept = Vec::new();
if let Some(start_idx) = start_idx {
for (s, &p) in track.samples[start_idx..].iter().zip(&pts[start_idx..]) {
if p >= hi {
break;
}
kept.push(s.clone());
}
}
if !kept.is_empty() {
kept_any = true;
}
let first_dts = kept.iter().find_map(|s| s.dts).unwrap_or(0).max(0) as u64;
out_tracks.push(Track::new_at(track.spec.clone(), kept, first_dts));
}
if !kept_any {
return Err(Error::InvalidInput(
"trim: window selected no samples on any track",
));
}
Ok(Media::new(out_tracks, self.movie_timescale))
}
pub fn anchor_duration(&self) -> Option<(u64, u32)> {
let anchor = anchor_index(self).ok()?;
let t = &self.tracks[anchor];
let ticks: u64 = t
.samples
.iter()
.map(|s| s.duration.unwrap_or(0) as u64)
.sum();
Some((ticks, t.spec.timescale))
}
}
#[derive(Debug, Clone)]
pub struct Repackage {
target_duration_secs: f64,
select: Option<Vec<usize>>,
trim: Option<(u64, u64)>,
}
impl Repackage {
pub fn new(target_duration_secs: f64) -> Self {
Self {
target_duration_secs,
select: None,
trim: None,
}
}
pub fn select_tracks(mut self, indices: &[usize]) -> Self {
self.select = Some(indices.to_vec());
self
}
pub fn trim(mut self, start: u64, end: u64) -> Self {
self.trim = Some((start, end));
self
}
pub fn run_media(&self, media: &Media) -> Result<RepackageOutput> {
let mut work = media.clone();
if let Some(indices) = &self.select {
work = work.select_tracks(indices)?;
}
if let Some((start, end)) = self.trim {
work = work.trim(start, end)?;
}
self.segment(&work)
}
pub fn run(&self, fmp4: &[u8]) -> Result<RepackageOutput> {
let media = Fmp4Demux::new().unpackage(fmp4)?;
self.run_media(&media)
}
fn segment(&self, media: &Media) -> Result<RepackageOutput> {
if media.tracks.is_empty() {
return Err(Error::InvalidInput("repackage: media has no tracks"));
}
let specs: Vec<TrackSpec> = media.tracks.iter().map(|t| t.spec.clone()).collect();
let movie_timescale = media.movie_timescale;
let mut seg = Segmenter::new(specs, movie_timescale, self.target_duration_secs)?;
let origin = TimelineOrigin::of(&media.tracks);
let mut dts = alloc::vec![0u128; media.tracks.len()];
for (ti, track) in media.tracks.iter().enumerate() {
let first = relative_decode_times(track, &origin)
.first()
.map_or(0, |&t| t.max(0) as u64);
seg.set_start_decode_time(track.spec.track_id, first)?;
dts[ti] = u128::from(first);
}
let init_segment = seg.init_segment()?;
let mut seg_ready: Vec<Vec<u8>> = Vec::new();
let mut cursors = alloc::vec![0usize; media.tracks.len()];
while let Some(ti) = earliest_decode_time_track(media, &dts, &cursors) {
let track = &media.tracks[ti];
let sample = &track.samples[cursors[ti]];
seg.push(track.spec.track_id, sample.clone())?;
dts[ti] += sample.duration.unwrap_or(0) as u128;
cursors[ti] += 1;
for s in seg.take_ready() {
seg_ready.push(s);
}
}
seg.flush()?;
seg_ready.extend(seg.take_ready());
Ok(RepackageOutput {
init_segment,
media_segments: seg_ready,
})
}
}
#[derive(Debug, Clone)]
pub struct RepackageOutput {
pub init_segment: Vec<u8>,
pub media_segments: Vec<Vec<u8>>,
}
impl RepackageOutput {
pub fn to_contiguous(&self) -> Vec<u8> {
let total =
self.init_segment.len() + self.media_segments.iter().map(Vec::len).sum::<usize>();
let mut out = Vec::with_capacity(total);
out.extend_from_slice(&self.init_segment);
for seg in &self.media_segments {
out.extend_from_slice(seg);
}
out
}
pub fn segment_count(&self) -> usize {
self.media_segments.len()
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::ir::CodecConfig;
use crate::pipeline::Sample;
use alloc::vec;
fn track(timescale: u32, count: usize) -> Track {
let samples = (0..count)
.map(|i| {
let dts = i as i64;
Sample::new(vec![0u8; 4], Some(dts), Some(dts), Some(1), true)
})
.collect();
Track::new(
TrackSpec::new(
1,
timescale,
CodecConfig::Aac {
esds: crate::mp4esds::EsdsBox::new(crate::mp4esds::ESDescriptor::new(
1,
0,
Some(crate::mp4esds::DecoderConfigDescriptor::new(
0x40,
0x05,
false,
0,
0,
0,
Some(crate::mp4esds::DecoderSpecificInfo::new(vec![0x12, 0x10])),
)),
Some(crate::mp4esds::SLConfigDescriptor::predefined_two()),
)),
channel_count: 2,
sample_rate: timescale,
sample_size: 16,
},
),
samples,
)
}
fn media(timescales: &[u32]) -> Media {
Media::new(timescales.iter().map(|&t| track(t, 4)).collect(), 1000)
}
fn pick(media: &Media, dts: &[u128], cursors: &[usize]) -> Option<usize> {
earliest_decode_time_track(media, dts, cursors)
}
#[test]
fn earliest_is_compared_as_seconds_not_ticks() {
let m = media(&[1000, 100]);
assert_eq!(pick(&m, &[1, 2], &[0, 0]), Some(0));
assert_eq!(pick(&m, &[500, 3], &[0, 0]), Some(1));
}
#[test]
fn a_tie_picks_the_lowest_track_index() {
let m = media(&[5000, 1000]);
assert_eq!(pick(&m, &[5, 1], &[0, 0]), Some(0));
let m = media(&[1000, 5000]);
assert_eq!(pick(&m, &[1, 5], &[0, 0]), Some(0));
let m = media(&[1000, 2000, 3000]);
assert_eq!(pick(&m, &[1, 2, 3], &[0, 0, 0]), Some(0));
}
#[test]
fn the_whole_interleave_order_is_stable() {
let m = media(&[1000, 2000]);
let mut dts = vec![0u128; 2];
let mut cursors = vec![0usize; 2];
let mut order = Vec::new();
while let Some(ti) = pick(&m, &dts, &cursors) {
order.push(ti);
dts[ti] += 1;
cursors[ti] += 1;
}
assert_eq!(order, vec![0, 1, 1, 0, 1, 1, 0, 0]);
}
#[test]
fn zero_timescale_is_treated_as_one() {
let m = media(&[0, 1]);
assert_eq!(pick(&m, &[0, 0], &[0, 0]), Some(0));
assert_eq!(pick(&m, &[1, 1], &[0, 0]), Some(0));
let m = media(&[1, 1_000_000]);
assert_eq!(pick(&m, &[1, 0], &[0, 0]), Some(1));
}
#[test]
fn exhausted_tracks_are_skipped() {
let m = media(&[1000, 1000]);
assert_eq!(pick(&m, &[0, 0], &[0, 4]), Some(0), "track 1 is exhausted");
assert_eq!(pick(&m, &[0, 0], &[4, 0]), Some(1), "track 0 is exhausted");
assert_eq!(pick(&m, &[0, 0], &[4, 4]), None, "both exhausted");
}
#[test]
fn huge_coprime_timescales_compare_correctly() {
const A: u32 = 4_294_967_291;
const B: u32 = 4_294_967_279;
const C: u32 = 4_294_967_231;
let m = media(&[A, B, C]);
assert_eq!(pick(&m, &[0, 0, 0], &[0, 0, 0]), Some(0));
assert_eq!(pick(&m, &[1, 0, 0], &[0, 0, 0]), Some(1));
assert_eq!(pick(&m, &[0, 0, 1], &[0, 0, 0]), Some(0));
assert_eq!(pick(&m, &[0, 1, 1], &[0, 0, 0]), Some(0));
assert_eq!(pick(&m, &[0, 1, 1], &[4, 0, 0]), Some(1));
assert_eq!(pick(&m, &[1, 1, 1], &[0, 0, 0]), Some(0));
let lcm_overflows = (A as u128) * (B as u128) * (C as u128);
assert!(lcm_overflows > u64::MAX as u128);
}
}