use alloc::collections::VecDeque;
use alloc::format;
use alloc::string::{String, ToString};
use alloc::vec;
use alloc::vec::Vec;
use broadcast_common::{Demand, Stage, Timestamp};
use crate::dash::DashPackager;
use crate::error::{Error, Result};
use crate::media::Media;
use crate::movie_fragment::{
MovieFragmentBox, MovieFragmentHeaderBox, TFHD_DEFAULT_BASE_IS_MOOF, TRUN_DATA_OFFSET_PRESENT,
TRUN_SAMPLE_COMPOSITION_TIME_OFFSET_PRESENT, TRUN_SAMPLE_DURATION_PRESENT,
TRUN_SAMPLE_FLAGS_PRESENT, TRUN_SAMPLE_SIZE_PRESENT, TrackFragmentBaseMediaDecodeTimeBox,
TrackFragmentBox, TrackFragmentHeaderBox, TrackFragmentRunBox, TrunSample,
};
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 SAMPLE_FLAGS_SYNC: u32 = 0x0200_0000;
const SAMPLE_FLAGS_NON_SYNC: u32 = 0x0101_0000;
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 take_counts: Vec<usize> = self
.tracks
.iter()
.enumerate()
.map(|(i, t)| {
if final_chunk {
t.pending.len()
} else if i == anchor {
self.chunk_samples.min(t.pending.len())
} else {
0
}
})
.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(())
}
}
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 any_cts = ft.samples.iter().any(|s| s.composition_offset() != 0);
let samples: Vec<TrunSample> = ft
.samples
.iter()
.map(|s| TrunSample {
sample_duration: Some(s.duration.unwrap_or(0)),
sample_size: Some(s.data.len() as u32),
sample_flags: Some(if s.flags.is_sync {
SAMPLE_FLAGS_SYNC
} else {
SAMPLE_FLAGS_NON_SYNC
}),
sample_composition_time_offset: if any_cts {
Some(s.composition_offset())
} else {
None
},
})
.collect();
let mut tr_flags = TRUN_DATA_OFFSET_PRESENT
| TRUN_SAMPLE_DURATION_PRESENT
| TRUN_SAMPLE_SIZE_PRESENT
| TRUN_SAMPLE_FLAGS_PRESENT;
let version = if any_cts {
tr_flags |= TRUN_SAMPLE_COMPOSITION_TIME_OFFSET_PRESENT;
1u8
} else {
0u8
};
let trun = TrackFragmentRunBox {
version,
tr_flags,
data_offset: Some(0),
first_sample_flags: None,
samples,
};
let tfhd = TrackFragmentHeaderBox {
flags: TFHD_DEFAULT_BASE_IS_MOOF,
track_id: ft.track_id,
base_data_offset: None,
sample_description_index: None,
default_sample_duration: None,
default_sample_size: None,
default_sample_flags: None,
};
let tfdt = TrackFragmentBaseMediaDecodeTimeBox::new_v1(ft.base_media_decode_time);
traf_boxes.push(TrackFragmentBox {
tfhd,
tfdt: Some(tfdt),
trun: vec![trun],
});
}
let mut moof = MovieFragmentBox {
mfhd: MovieFragmentHeaderBox::new(sequence_number),
traf: 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)
}
#[derive(Debug, Clone)]
pub struct LlDashPackager {
pub base: DashPackager,
pub segment_duration_secs: f64,
pub chunk_duration_secs: f64,
pub latency_target_ms: u32,
pub playback_rate: Option<(f64, f64)>,
}
impl LlDashPackager {
pub fn new(
segment_duration_secs: f64,
chunk_duration_secs: f64,
latency_target_ms: u32,
availability_start_time: impl Into<String>,
) -> Result<Self> {
if !(segment_duration_secs.is_finite() && segment_duration_secs > 0.0) {
return Err(Error::InvalidInput(
"segment_duration_secs must be positive and finite",
));
}
if !(chunk_duration_secs.is_finite() && chunk_duration_secs > 0.0) {
return Err(Error::InvalidInput(
"chunk_duration_secs must be positive and finite",
));
}
if chunk_duration_secs > segment_duration_secs {
return Err(Error::InvalidInput(
"chunk_duration_secs must not exceed segment_duration_secs",
));
}
let base = DashPackager {
dynamic: true,
availability_start_time: Some(availability_start_time.into()),
..DashPackager::default()
};
Ok(Self {
base,
segment_duration_secs,
chunk_duration_secs,
latency_target_ms,
playback_rate: None,
})
}
pub fn with_playback_rate(mut self, min: f64, max: f64) -> Self {
self.playback_rate = Some((min, max));
self
}
pub fn availability_time_offset(&self) -> f64 {
(self.segment_duration_secs - self.chunk_duration_secs).max(0.0)
}
}
impl Package for LlDashPackager {
type Media = Media;
type Output = String;
type Error = Error;
fn package(&mut self, media: &Media) -> Result<String> {
let base_xml = self.base.package(media)?;
Ok(self.inject_ll(&base_xml))
}
}
impl LlDashPackager {
fn inject_ll(&self, xml: &str) -> String {
let ato = self.availability_time_offset();
let ato_str = format_secs(ato);
let mut out = String::with_capacity(xml.len() + 256);
for line in xml.lines() {
let trimmed = line.trim_start();
if trimmed.starts_with("<SegmentTemplate") && line.trim_end().ends_with("/>") {
let end = line.trim_end();
let head = &end[..end.len() - 2]; out.push_str(head);
out.push_str(" availabilityTimeOffset=\"");
out.push_str(&ato_str);
out.push_str("\" availabilityTimeComplete=\"false\"/>\n");
} else {
out.push_str(line);
out.push('\n');
}
}
let service = self.service_description();
if let Some(pos) = find_mpd_open_end(&out) {
let mut with_sd = String::with_capacity(out.len() + service.len());
with_sd.push_str(&out[..pos]);
with_sd.push_str(&service);
with_sd.push_str(&out[pos..]);
with_sd
} else {
out
}
}
fn service_description(&self) -> String {
let mut s = String::new();
s.push_str(" <ServiceDescription id=\"0\">\n");
s.push_str(" <Latency target=\"");
s.push_str(&self.latency_target_ms.to_string());
s.push_str("\"/>\n");
if let Some((min, max)) = self.playback_rate {
s.push_str(" <PlaybackRate min=\"");
s.push_str(&format_secs(min));
s.push_str("\" max=\"");
s.push_str(&format_secs(max));
s.push_str("\"/>\n");
}
s.push_str(" </ServiceDescription>\n");
s
}
}
fn find_mpd_open_end(xml: &str) -> Option<usize> {
let start = xml.find("<MPD")?;
let rel = xml[start..].find('>')?;
Some(start + rel + 1 + 1) }
fn format_secs(v: f64) -> String {
let millis = (v * 1000.0 + 0.5) as u64;
let whole = millis / 1000;
let frac = millis % 1000;
if frac == 0 {
return whole.to_string();
}
let mut f = format!("{frac:03}");
while f.ends_with('0') {
f.pop();
}
format!("{whole}.{f}")
}