use alloc::collections::BTreeMap;
use alloc::sync::Arc;
use alloc::vec;
use alloc::vec::Vec;
use crate::descriptors::{AnyDescriptor, DescriptorLoop, DescriptorRegistry};
use broadcast_common::Parse;
use mpeg_ts::section::Section;
mod bat;
mod eit;
mod nit;
mod sdt;
pub use bat::*;
pub use eit::*;
pub use nit::*;
pub use sdt::*;
pub const DEFAULT_MAX_PARTIAL_KEYS: usize = 256;
pub type CollectResult<T> = core::result::Result<T, CollectError>;
#[derive(Debug, thiserror::Error)]
#[non_exhaustive]
pub enum CollectError {
#[error("section parse failed: {0}")]
Section(#[from] crate::Error),
#[error(
"table_id {table_id:#04x} is a short-form section and cannot be multi-section collected"
)]
ShortFormSection {
table_id: u8,
},
#[error(
"section_number {section_number} exceeds last_section_number {last_section_number} for table_id {table_id:#04x}"
)]
SectionNumberOutOfRange {
table_id: u8,
section_number: u8,
last_section_number: u8,
},
#[error("conflicting bytes for table_id {table_id:#04x} section {section_number}")]
ConflictingSection {
table_id: u8,
section_number: u8,
},
#[error(
"EIT schedule table_id {table_id:#04x} is outside advertised range {first_table_id:#04x}..={last_table_id:#04x}"
)]
EitTableIdOutOfRange {
table_id: u8,
first_table_id: u8,
last_table_id: u8,
},
}
impl From<mpeg_ts::Error> for CollectError {
fn from(e: mpeg_ts::Error) -> Self {
CollectError::Section(crate::Error::from(e))
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
#[non_exhaustive]
pub struct SectionSetKey {
pub pid: Option<u16>,
pub table_id: u8,
pub extension_id: u16,
pub current_next_indicator: bool,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub struct SectionSetMeta {
pub key: SectionSetKey,
pub version_number: u8,
pub last_section_number: u8,
}
#[derive(Debug)]
struct PartialSectionSet {
meta: SectionSetMeta,
slots: Vec<Option<Arc<[u8]>>>,
filled: usize,
emitted: bool,
last_touch: u64,
}
impl PartialSectionSet {
fn new(meta: SectionSetMeta) -> Self {
let len = meta.last_section_number as usize + 1;
Self {
meta,
slots: vec![None; len],
filled: 0,
emitted: false,
last_touch: 0,
}
}
fn reset(&mut self, meta: SectionSetMeta) {
let last_touch = self.last_touch;
*self = Self::new(meta);
self.last_touch = last_touch;
}
fn insert(&mut self, section_number: u8, bytes: Arc<[u8]>) -> CollectResult<bool> {
let index = section_number as usize;
if let Some(existing) = &self.slots[index] {
if existing.as_ref() == bytes.as_ref() {
return Ok(false);
}
return Err(CollectError::ConflictingSection {
table_id: self.meta.key.table_id,
section_number,
});
}
self.slots[index] = Some(bytes);
self.filled += 1;
self.emitted = false;
Ok(true)
}
fn complete(&self) -> bool {
self.filled == self.slots.len()
}
fn to_complete(&self) -> Option<CompleteSectionSet> {
if !self.complete() || self.emitted {
return None;
}
let sections = self
.slots
.iter()
.map(|slot| slot.as_ref().expect("complete set has no holes").clone())
.collect();
Some(CompleteSectionSet {
meta: self.meta,
sections,
})
}
}
#[derive(Debug)]
pub struct SectionSetCollector {
partial: BTreeMap<SectionSetKey, PartialSectionSet>,
touch_order: BTreeMap<u64, SectionSetKey>,
tick: u64,
max_partial_keys: usize,
evicted_for_capacity: u64,
}
impl Default for SectionSetCollector {
fn default() -> Self {
Self {
partial: BTreeMap::new(),
touch_order: BTreeMap::new(),
tick: 0,
max_partial_keys: DEFAULT_MAX_PARTIAL_KEYS,
evicted_for_capacity: 0,
}
}
}
impl SectionSetCollector {
#[must_use]
pub fn new() -> Self {
Self::default()
}
#[must_use]
pub fn with_max_partial_keys(mut self, max_partial_keys: usize) -> Self {
self.max_partial_keys = max_partial_keys;
self
}
#[must_use]
pub fn evicted_for_capacity(&self) -> u64 {
self.evicted_for_capacity
}
fn touch_or_evict(&mut self, key: SectionSetKey) -> bool {
self.tick = self.tick.wrapping_add(1).max(1);
if !self.partial.contains_key(&key) && self.partial.len() >= self.max_partial_keys {
let oldest = self.touch_order.keys().next().copied();
match oldest {
Some(oldest_tick) => {
if let Some(oldest_key) = self.touch_order.remove(&oldest_tick) {
self.partial.remove(&oldest_key);
self.evicted_for_capacity += 1;
}
}
None => return false,
}
}
true
}
pub fn push_section(
&mut self,
bytes: impl AsRef<[u8]>,
) -> CollectResult<Option<CompleteSectionSet>> {
self.push_section_with_pid(None, bytes)
}
pub fn push_section_with_pid(
&mut self,
pid: Option<u16>,
bytes: impl AsRef<[u8]>,
) -> CollectResult<Option<CompleteSectionSet>> {
let raw = bytes.as_ref();
let section = Section::parse(raw)?;
if !section.section_syntax_indicator {
return Err(CollectError::ShortFormSection {
table_id: section.table_id,
});
}
if section.section_number > section.last_section_number {
return Err(CollectError::SectionNumberOutOfRange {
table_id: section.table_id,
section_number: section.section_number,
last_section_number: section.last_section_number,
});
}
section.validate_crc(raw)?;
let key = SectionSetKey {
pid,
table_id: section.table_id,
extension_id: section.extension_id,
current_next_indicator: section.current_next_indicator,
};
let meta = SectionSetMeta {
key,
version_number: section.version_number,
last_section_number: section.last_section_number,
};
let bytes: Arc<[u8]> = Arc::from(raw);
if !self.touch_or_evict(key) {
return Ok(None);
}
let now = self.tick;
let partial = self
.partial
.entry(key)
.or_insert_with(|| PartialSectionSet::new(meta));
if partial.meta.version_number != meta.version_number
|| partial.meta.last_section_number != meta.last_section_number
{
partial.reset(meta);
}
self.touch_order.remove(&partial.last_touch);
partial.last_touch = now;
self.touch_order.insert(now, key);
partial.insert(section.section_number, bytes)?;
let complete = partial.to_complete();
if complete.is_some() {
partial.emitted = true;
}
Ok(complete)
}
pub fn clear(&mut self) {
self.partial.clear();
self.touch_order.clear();
}
#[must_use]
pub fn len(&self) -> usize {
self.partial.len()
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.partial.is_empty()
}
}
#[derive(Debug, Clone)]
pub struct CompleteSectionSet {
meta: SectionSetMeta,
sections: Vec<Arc<[u8]>>,
}
#[derive(Debug)]
pub struct CompleteTable<T> {
meta: SectionSetMeta,
sections: Vec<T>,
}
impl<T> CompleteTable<T> {
#[must_use]
pub const fn meta(&self) -> SectionSetMeta {
self.meta
}
#[must_use]
pub fn sections(&self) -> &[T] {
&self.sections
}
#[must_use]
pub fn into_sections(self) -> Vec<T> {
self.sections
}
}
impl CompleteSectionSet {
#[must_use]
pub const fn meta(&self) -> SectionSetMeta {
self.meta
}
#[must_use]
pub fn section_bytes(&self) -> impl ExactSizeIterator<Item = &[u8]> {
self.sections.iter().map(AsRef::as_ref)
}
pub fn parse_sections<'a, T>(&'a self) -> crate::Result<Vec<T>>
where
T: Parse<'a, Error = crate::Error>,
{
self.section_bytes().map(T::parse).collect()
}
pub fn table<'a, T>(&'a self) -> crate::Result<CompleteTable<T>>
where
T: Parse<'a, Error = crate::Error>,
{
Ok(CompleteTable {
meta: self.meta,
sections: self.parse_sections()?,
})
}
pub fn nit(&self) -> crate::Result<CompleteNit<'_>> {
CompleteNit::parse(self, None)
}
pub fn nit_with_registry<'a>(
&'a self,
registry: &'a DescriptorRegistry,
) -> crate::Result<CompleteNit<'a>> {
CompleteNit::parse(self, Some(registry))
}
pub fn bat(&self) -> crate::Result<CompleteBat<'_>> {
CompleteBat::parse(self, None)
}
pub fn bat_with_registry<'a>(
&'a self,
registry: &'a DescriptorRegistry,
) -> crate::Result<CompleteBat<'a>> {
CompleteBat::parse(self, Some(registry))
}
pub fn sdt(&self) -> crate::Result<CompleteSdt<'_>> {
CompleteSdt::parse(self, None)
}
pub fn sdt_with_registry<'a>(
&'a self,
registry: &'a DescriptorRegistry,
) -> crate::Result<CompleteSdt<'a>> {
CompleteSdt::parse(self, Some(registry))
}
pub fn eit(&self) -> crate::Result<CompleteEit<'_>> {
CompleteEit::parse(self, None)
}
pub fn eit_with_registry<'a>(
&'a self,
registry: &'a DescriptorRegistry,
) -> crate::Result<CompleteEit<'a>> {
CompleteEit::parse(self, Some(registry))
}
}
#[derive(Debug)]
pub struct ParsedDescriptorLoop<'a> {
raw: Vec<DescriptorLoop<'a>>,
descriptors: Vec<crate::Result<AnyDescriptor<'a>>>,
}
impl<'a> ParsedDescriptorLoop<'a> {
pub(crate) fn parse(raw: DescriptorLoop<'a>, registry: Option<&'a DescriptorRegistry>) -> Self {
Self::parse_all(core::iter::once(raw), registry)
}
pub(crate) fn parse_all(
raws: impl IntoIterator<Item = DescriptorLoop<'a>>,
registry: Option<&'a DescriptorRegistry>,
) -> Self {
let raw: Vec<DescriptorLoop<'a>> = raws.into_iter().collect();
let mut descriptors = Vec::new();
for loop_ in &raw {
match registry {
Some(registry) => descriptors.extend(registry.parse_loop(loop_.raw())),
None => descriptors.extend(loop_.iter()),
}
}
Self { raw, descriptors }
}
#[must_use]
pub fn raw(&self) -> &[DescriptorLoop<'a>] {
&self.raw
}
pub fn descriptors(&self) -> &[crate::Result<AnyDescriptor<'a>>] {
&self.descriptors
}
}
#[cfg(test)]
mod tests {
use super::*;
const TEST_TABLE_ID: u8 = 0x42;
fn min_section(extension_id: u16) -> Vec<u8> {
let section_length: u16 = 9; let mut buf = vec![0u8; 12];
buf[0] = TEST_TABLE_ID;
buf[1] = 0xB0 | ((section_length >> 8) as u8 & 0x0F);
buf[2] = (section_length & 0xFF) as u8;
buf[3..5].copy_from_slice(&extension_id.to_be_bytes());
buf[5] = 0xC1;
buf[6] = 0;
buf[7] = 0;
let crc = broadcast_common::crc32_mpeg2::compute(&buf[..8]);
buf[8..12].copy_from_slice(&crc.to_be_bytes());
buf
}
#[test]
fn collect_single_section_is_complete() {
let mut c = SectionSetCollector::new();
let sec = min_section(0);
let result = c.push_section(&sec).unwrap();
assert!(result.is_some());
assert_eq!(c.len(), 1);
}
#[test]
fn new_key_beyond_cap_evicts_least_recently_touched_completed_entry() {
let mut c = SectionSetCollector::new().with_max_partial_keys(3);
for eid in 0..3u16 {
let sec = min_section(eid);
let result = c.push_section(&sec).unwrap();
assert!(
result.is_some(),
"single-section set for eid {eid} completes"
);
}
assert_eq!(c.len(), 3);
assert_eq!(c.evicted_for_capacity(), 0);
let sec4 = min_section(3);
let result = c.push_section(&sec4).unwrap();
assert!(
result.is_some(),
"new key beyond cap must evict, not be skipped"
);
assert_eq!(c.len(), 3, "cap stays honored, not just grown");
assert_eq!(c.evicted_for_capacity(), 1);
let sec0_again = min_section(0);
let result = c.push_section(&sec0_again).unwrap();
assert!(result.is_some(), "evicted key re-admits like a fresh one");
c.clear();
assert!(c.is_empty());
let result = c.push_section(&sec4).unwrap();
assert!(result.is_some());
assert_eq!(c.len(), 1);
}
#[test]
fn partial_keys_cap_does_not_skip_existing_key() {
let mut c = SectionSetCollector::new().with_max_partial_keys(1);
let sec0 = {
let mut buf = min_section(0xAB);
buf[7] = 1;
let crc = broadcast_common::crc32_mpeg2::compute(&buf[..8]);
buf[8..12].copy_from_slice(&crc.to_be_bytes());
buf
};
let result = c.push_section(&sec0).unwrap();
assert!(result.is_none(), "incomplete section set yields None");
let mut sec1 = min_section(0xAB);
sec1[6] = 1; sec1[7] = 1; let crc = broadcast_common::crc32_mpeg2::compute(&sec1[..8]);
sec1[8..12].copy_from_slice(&crc.to_be_bytes());
let result = c.push_section(&sec1).unwrap();
assert!(
result.is_some(),
"existing key must NOT be skipped when cap full"
);
assert_eq!(c.len(), 1);
}
}