use alloc::collections::VecDeque;
use alloc::vec::Vec;
use broadcast_common::{Demand, Stage, Timestamp};
use crate::error::{Error, Result};
use crate::ll_dash::build_chunk;
use crate::pipeline::{FragmentTrackData, Sample, TrackSpec, build_init_segment};
use crate::segmenter::{
MAX_PENDING_SAMPLES_PER_TRACK, MediaClock, choose_anchor, no_sync_sample_error,
};
#[derive(Debug, Clone)]
pub struct PartInfo {
pub bytes: Vec<u8>,
pub duration: f64,
pub independent: bool,
pub segment_seq: u32,
pub part_index: u32,
}
#[derive(Debug, Clone)]
pub struct SegmentInfo {
pub bytes: Vec<u8>,
pub duration: f64,
pub segment_seq: u32,
pub part_count: u32,
}
struct TrackState {
spec: TrackSpec,
pending: Vec<Sample>,
part_start: usize,
part_base_decode: u64,
seg_base_decode: u64,
part_clock: MediaClock,
seg_clock: MediaClock,
}
pub struct LlHlsSegmenter {
tracks: Vec<TrackState>,
movie_timescale: u32,
anchor: usize,
target_ticks: u64,
part_target_ticks: u64,
anchor_seg_dur: u64,
anchor_part_dur: u64,
anchor_clock: MediaClock,
next_seq: u32,
current_segment: u32,
next_part_index: u32,
ready: VecDeque<LlHlsStageOutput>,
}
impl LlHlsSegmenter {
pub fn with_part_target(
tracks: Vec<TrackSpec>,
movie_timescale: u32,
target_duration_secs: f64,
part_target_ms: u32,
) -> Result<Self> {
if tracks.is_empty() {
return Err(Error::InvalidInput(
"ll-hls 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 part_target_ms == 0 {
return Err(Error::InvalidInput("part_target_ms 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 part_target_ticks =
((part_target_ms as u64 * tracks[anchor].timescale as u64) / 1000).max(1);
let tracks = tracks
.into_iter()
.map(|spec| TrackState {
spec,
pending: Vec::new(),
part_start: 0,
part_base_decode: 0,
seg_base_decode: 0,
part_clock: MediaClock::new(),
seg_clock: MediaClock::new(),
})
.collect();
Ok(Self {
tracks,
movie_timescale,
anchor,
target_ticks,
part_target_ticks,
anchor_seg_dur: 0,
anchor_part_dur: 0,
anchor_clock: MediaClock::new(),
next_seq: 1,
current_segment: 1,
next_part_index: 0,
ready: VecDeque::new(),
})
}
pub fn part_target_secs(&self) -> f64 {
self.part_target_ticks as f64 / self.tracks[self.anchor].spec.timescale as f64
}
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_seg_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(no_sync_sample_error());
}
if idx == self.anchor {
let elapsed = self.anchor_clock.tick(&sample);
self.anchor_seg_dur += elapsed;
self.anchor_part_dur += elapsed;
}
self.tracks[idx].pending.push(sample);
if idx == self.anchor
&& self.anchor_part_dur >= self.part_target_ticks
&& self.anchor_seg_dur < self.target_ticks
{
self.emit_part(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_parts(&mut self) -> Vec<PartInfo> {
let mut parts = Vec::new();
let mut remaining = VecDeque::with_capacity(self.ready.len());
for item in self.ready.drain(..) {
match item {
LlHlsStageOutput::Part(p) => parts.push(p),
other => remaining.push_back(other),
}
}
self.ready = remaining;
parts
}
pub fn take_ready_segments(&mut self) -> Vec<SegmentInfo> {
let mut segments = Vec::new();
let mut remaining = VecDeque::with_capacity(self.ready.len());
for item in self.ready.drain(..) {
match item {
LlHlsStageOutput::Segment(s) => segments.push(s),
other => remaining.push_back(other),
}
}
self.ready = remaining;
segments
}
fn finish_segment(&mut self) -> Result<()> {
let has_tail = self.tracks.iter().any(|t| t.part_start < t.pending.len());
if has_tail {
self.emit_part(true)?;
}
let seg_seq = self.next_seq;
let seg_bytes = {
let frags: Vec<FragmentTrackData<'_>> = self
.tracks
.iter()
.filter(|t| !t.pending.is_empty())
.map(|t| FragmentTrackData {
track_id: t.spec.track_id,
base_media_decode_time: t.seg_base_decode,
samples: &t.pending,
})
.collect();
crate::pipeline::build_media_segment(seg_seq, &frags)?
};
self.next_seq += 1;
let seg_duration =
self.anchor_seg_dur as f64 / self.tracks[self.anchor].spec.timescale as f64;
self.ready.push_back(LlHlsStageOutput::Segment(SegmentInfo {
bytes: seg_bytes,
duration: seg_duration,
segment_seq: self.current_segment,
part_count: self.next_part_index,
}));
for t in &mut self.tracks {
let clock = &mut t.seg_clock;
let dur: u64 = t.pending.iter().map(|s| clock.tick(s)).sum();
t.seg_base_decode += dur;
t.part_base_decode = t.seg_base_decode;
t.pending.clear();
t.part_start = 0;
}
self.anchor_seg_dur = 0;
self.anchor_part_dur = 0;
self.current_segment += 1;
self.next_part_index = 0;
Ok(())
}
fn emit_part(&mut self, final_part: bool) -> Result<()> {
let anchor = self.anchor;
let take_ends: Vec<usize> = self
.tracks
.iter()
.enumerate()
.map(|(i, t)| {
if final_part || i == anchor {
t.pending.len()
} else {
t.part_start
}
})
.collect();
if take_ends
.iter()
.zip(&self.tracks)
.all(|(&end, t)| end == t.part_start)
{
return Ok(());
}
let anchor_state = &self.tracks[anchor];
let independent = anchor_state
.pending
.get(anchor_state.part_start)
.map(|s| s.flags.is_sync)
.unwrap_or(false);
let part_secs = self.anchor_part_dur as f64 / self.tracks[anchor].spec.timescale as f64;
let seq = self.next_seq;
let part_bytes = {
let frags: Vec<FragmentTrackData<'_>> = self
.tracks
.iter()
.zip(&take_ends)
.filter(|&(t, &end)| end > t.part_start)
.map(|(t, &end)| FragmentTrackData {
track_id: t.spec.track_id,
base_media_decode_time: t.part_base_decode,
samples: &t.pending[t.part_start..end],
})
.collect();
build_chunk(seq, &frags, false)?
};
self.next_seq += 1;
for (t, &end) in self.tracks.iter_mut().zip(&take_ends) {
let clock = &mut t.part_clock;
let dur: u64 = t.pending[t.part_start..end]
.iter()
.map(|s| clock.tick(s))
.sum();
t.part_base_decode += dur;
t.part_start = end;
}
let part_index = self.next_part_index;
self.next_part_index += 1;
self.anchor_part_dur = 0;
self.ready.push_back(LlHlsStageOutput::Part(PartInfo {
bytes: part_bytes,
duration: part_secs,
independent,
segment_seq: self.current_segment,
part_index,
}));
Ok(())
}
}
#[derive(Debug, Clone)]
#[non_exhaustive]
pub enum LlHlsStageOutput {
Part(PartInfo),
Segment(SegmentInfo),
}
impl Stage for LlHlsSegmenter {
type In<'a> = (u32, Sample);
type Out = LlHlsStageOutput;
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()
}
}
}