use alloc::collections::VecDeque;
use alloc::vec;
use alloc::vec::Vec;
use broadcast_common::{Demand, Stage, Timestamp};
use crate::error::{Error, Result};
use crate::media::Media;
use crate::movie_fragment::{
MediaRunSample, MovieFragmentBox, MovieFragmentHeaderBox, TrackFragmentBaseMediaDecodeTimeBox,
TrackFragmentBox, media_fragment_run,
};
use crate::pipeline::{FragmentTrackData, Sample, TrackSpec, build_init_segment};
use crate::segmenter::{MAX_PENDING_SAMPLES_PER_TRACK, MediaClock, choose_anchor};
use crate::segments::{MediaDataBox, SegmentTypeBox};
use broadcast_common::{Package, Serialize};
const STYP_MAJOR_BRAND: [u8; 4] = *b"msdh";
const STYP_COMPATIBLE_BRANDS: [[u8; 4]; 2] = [*b"msdh", *b"msix"];
#[derive(Debug, Clone)]
pub struct Chunk {
pub data: Vec<u8>,
pub segment_number: u64,
pub is_segment_start: bool,
pub sequence_number: u32,
}
struct TrackState {
spec: TrackSpec,
pending: Vec<Sample>,
base_decode: u64,
flush_clock: MediaClock,
}
pub struct LlSegmenter {
tracks: Vec<TrackState>,
movie_timescale: u32,
anchor: usize,
target_ticks: u64,
chunk_samples: usize,
anchor_pending_dur: u64,
anchor_clock: MediaClock,
next_seq: u32,
current_segment: u64,
segment_open: bool,
ready: VecDeque<Chunk>,
}
impl LlSegmenter {
pub fn new(
tracks: Vec<TrackSpec>,
movie_timescale: u32,
target_duration_secs: f64,
chunk_samples: usize,
) -> Result<Self> {
if tracks.is_empty() {
return Err(Error::InvalidInput("ll segmenter needs at least one track"));
}
if !(target_duration_secs.is_finite() && target_duration_secs > 0.0) {
return Err(Error::InvalidInput(
"target_duration_secs must be positive and finite",
));
}
if chunk_samples == 0 {
return Err(Error::InvalidInput("chunk_samples must be >= 1"));
}
for (i, a) in tracks.iter().enumerate() {
if tracks[i + 1..].iter().any(|b| b.track_id == a.track_id) {
return Err(Error::InvalidInput("duplicate track_id"));
}
}
let anchor = choose_anchor(tracks.iter().map(|t| &t.config))?;
let anchor_timescale = tracks[anchor].timescale as f64;
let target_ticks = ((target_duration_secs * anchor_timescale) as u64).max(1);
let tracks = tracks
.into_iter()
.map(|spec| TrackState {
spec,
pending: Vec::new(),
base_decode: 0,
flush_clock: MediaClock::new(),
})
.collect();
Ok(Self {
tracks,
movie_timescale,
anchor,
target_ticks,
chunk_samples,
anchor_pending_dur: 0,
anchor_clock: MediaClock::new(),
next_seq: 1,
current_segment: 1,
segment_open: false,
ready: VecDeque::new(),
})
}
pub fn init_segment(&self) -> Result<Vec<u8>> {
let specs: Vec<TrackSpec> = self.tracks.iter().map(|t| t.spec.clone()).collect();
build_init_segment(&specs, self.movie_timescale)
}
pub fn push(&mut self, track_id: u32, sample: Sample) -> Result<()> {
let idx = self
.tracks
.iter()
.position(|t| t.spec.track_id == track_id)
.ok_or(Error::InvalidInput("push: unknown track_id"))?;
if idx == self.anchor
&& sample.flags.is_sync
&& self.anchor_pending_dur >= self.target_ticks
&& !self.tracks[self.anchor].pending.is_empty()
{
self.finish_segment()?;
}
if self.tracks[idx].pending.len() >= MAX_PENDING_SAMPLES_PER_TRACK {
return Err(crate::segmenter::no_sync_sample_error());
}
if idx == self.anchor {
self.anchor_pending_dur += self.anchor_clock.tick(&sample);
}
self.tracks[idx].pending.push(sample);
if idx == self.anchor
&& self.tracks[self.anchor].pending.len() >= self.chunk_samples
&& self.anchor_pending_dur < self.target_ticks
{
self.emit_chunk(false)?;
}
Ok(())
}
pub fn flush(&mut self) -> Result<()> {
if self.tracks.iter().any(|t| !t.pending.is_empty()) {
self.finish_segment()?;
}
Ok(())
}
pub fn take_ready(&mut self) -> Vec<Chunk> {
self.ready.drain(..).collect()
}
fn finish_segment(&mut self) -> Result<()> {
while self.tracks[self.anchor].pending.len() > self.chunk_samples {
self.emit_chunk(false)?;
}
self.emit_chunk(true)?;
self.current_segment += 1;
self.segment_open = false;
self.anchor_pending_dur = 0;
Ok(())
}
fn emit_chunk(&mut self, final_chunk: bool) -> Result<()> {
let is_start = !self.segment_open;
let anchor = self.anchor;
let anchor_scale = self.tracks[anchor].spec.timescale.max(1) as u64;
let anchor_take = if final_chunk {
self.tracks[anchor].pending.len()
} else {
self.chunk_samples.min(self.tracks[anchor].pending.len())
};
let chunk_span = if final_chunk {
u64::MAX
} else {
let t = &self.tracks[anchor];
let mut clock =
MediaClock::resumed_at(Some(i64::try_from(t.base_decode).unwrap_or(i64::MAX)));
let mut span: u64 = 0;
for s in &t.pending[..anchor_take] {
span += clock.tick(s);
}
span
};
let take_counts: Vec<usize> = self
.tracks
.iter()
.enumerate()
.map(|(i, t)| {
if final_chunk {
return t.pending.len();
}
if i == anchor {
return anchor_take;
}
let track_scale = t.spec.timescale.max(1) as u64;
let track_span = chunk_span.saturating_mul(track_scale) / anchor_scale;
take_span(t, track_span)
})
.collect();
if take_counts.iter().all(|&n| n == 0) {
return Ok(());
}
let seq = self.next_seq;
let chunk_bytes = {
let frags: Vec<FragmentTrackData<'_>> = self
.tracks
.iter()
.zip(&take_counts)
.filter(|&(_, &n)| n > 0)
.map(|(t, &n)| FragmentTrackData {
track_id: t.spec.track_id,
base_media_decode_time: t.base_decode,
samples: &t.pending[..n],
})
.collect();
build_chunk(seq, &frags, is_start)?
};
for (t, &n) in self.tracks.iter_mut().zip(&take_counts) {
let clock = &mut t.flush_clock;
let dur: u64 = t.pending[..n].iter().map(|s| clock.tick(s)).sum();
t.base_decode += dur;
t.pending.drain(..n);
}
self.next_seq += 1;
self.segment_open = true;
self.ready.push_back(Chunk {
data: chunk_bytes,
segment_number: self.current_segment,
is_segment_start: is_start,
sequence_number: seq,
});
Ok(())
}
}
fn take_span(t: &TrackState, target_ticks: u64) -> usize {
if t.pending.is_empty() {
return 0;
}
let mut clock = MediaClock::resumed_at(Some(i64::try_from(t.base_decode).unwrap_or(i64::MAX)));
let mut span: u64 = 0;
let mut end = 0usize;
for s in &t.pending {
let next = clock.tick(s);
if end > 0 && span + next > target_ticks {
break;
}
span += next;
end += 1;
}
end.max(1).min(t.pending.len())
}
impl Stage for LlSegmenter {
type In<'a> = (u32, Sample);
type Out = Chunk;
type Error = Error;
fn feed(&mut self, (track_id, sample): Self::In<'_>, _now: Timestamp) -> Result<()> {
self.push(track_id, sample)
}
fn poll(&mut self) -> Option<Self::Out> {
self.ready.pop_front()
}
fn finish(&mut self) -> Result<()> {
self.flush()
}
fn next_deadline(&self) -> Option<Timestamp> {
None
}
fn on_deadline(&mut self, _now: Timestamp) {}
fn demand(&self) -> Demand {
if self
.tracks
.iter()
.any(|t| t.pending.len() >= MAX_PENDING_SAMPLES_PER_TRACK)
{
Demand::saturated()
} else {
Demand::default()
}
}
}
impl Package for LlSegmenter {
type Media = Media;
type Output = Vec<Chunk>;
type Error = Error;
fn package(&mut self, media: &Media) -> Result<Vec<Chunk>> {
if media.tracks.is_empty() {
return Err(Error::InvalidInput("cannot package a Media with no tracks"));
}
let anchor_id = self.tracks[self.anchor].spec.track_id;
struct Cursor<'a> {
track_id: u32,
timescale: u64,
samples: &'a [Sample],
idx: usize,
dts_ticks: u64,
clock: MediaClock,
is_anchor: bool,
}
let mut cursors: Vec<Cursor<'_>> = media
.tracks
.iter()
.map(|t| Cursor {
track_id: t.spec.track_id,
timescale: (t.spec.timescale as u64).max(1),
samples: &t.samples,
idx: 0,
dts_ticks: 0,
clock: MediaClock::new(),
is_anchor: t.spec.track_id == anchor_id,
})
.collect();
loop {
let mut best: Option<usize> = None;
for (i, c) in cursors.iter().enumerate() {
if c.idx >= c.samples.len() {
continue;
}
best = Some(match best {
None => i,
Some(b) => {
let lhs = c.dts_ticks as u128 * cursors[b].timescale as u128;
let rhs = cursors[b].dts_ticks as u128 * c.timescale as u128;
if lhs < rhs || (lhs == rhs && !c.is_anchor && cursors[b].is_anchor) {
i
} else {
b
}
}
});
}
let Some(i) = best else { break };
let (track_id, sample) = {
let c = &mut cursors[i];
let s = c.samples[c.idx].clone();
c.dts_ticks += c.clock.tick(&s);
c.idx += 1;
(c.track_id, s)
};
self.push(track_id, sample)?;
}
self.flush()?;
Ok(self.take_ready())
}
}
pub(crate) fn build_chunk(
sequence_number: u32,
tracks: &[FragmentTrackData<'_>],
with_styp: bool,
) -> Result<Vec<u8>> {
let styp = with_styp.then(|| SegmentTypeBox {
major_brand: STYP_MAJOR_BRAND,
minor_version: 0,
compatible_brands: STYP_COMPATIBLE_BRANDS.to_vec(),
});
let mut traf_boxes = Vec::with_capacity(tracks.len());
for ft in tracks {
let (tfhd, trun) = media_fragment_run(
ft.track_id,
ft.samples.iter().map(|s| MediaRunSample {
duration: s.duration.unwrap_or(0),
size: s.data.len(),
is_sync: s.flags.is_sync,
composition_offset: s.composition_offset(),
}),
)?;
let tfdt = TrackFragmentBaseMediaDecodeTimeBox::new_v1(ft.base_media_decode_time);
traf_boxes.push(TrackFragmentBox::new(tfhd, Some(tfdt), vec![trun]));
}
let mut moof = MovieFragmentBox::new(MovieFragmentHeaderBox::new(sequence_number), traf_boxes);
let moof_size = moof.serialized_len();
let mut cursor = moof_size + 8;
let mut mdat_data = Vec::new();
for (i, ft) in tracks.iter().enumerate() {
moof.traf[i].trun[0].data_offset = Some(cursor as i32);
for s in ft.samples {
mdat_data.extend_from_slice(&s.data);
cursor += s.data.len();
}
}
let mdat = MediaDataBox { data: mdat_data };
let styp_len = styp.as_ref().map_or(0, |s| s.serialized_len());
let total = styp_len + moof.serialized_len() + mdat.serialized_len();
let mut out = vec![0u8; total];
let mut c = 0usize;
if let Some(styp) = &styp {
c += styp.serialize_into(&mut out[c..])?;
}
c += moof.serialize_into(&mut out[c..])?;
c += mdat.serialize_into(&mut out[c..])?;
out.truncate(c);
Ok(out)
}
#[cfg(feature = "std")]
mod mpd;
#[cfg(feature = "std")]
pub use mpd::*;