use std::collections::VecDeque;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ResolutionBitrate {
pub width: u32,
pub height: u32,
pub bitrate_kbps: u32,
pub codec: String,
}
impl ResolutionBitrate {
#[must_use]
pub fn new(width: u32, height: u32, bitrate_kbps: u32, codec: impl Into<String>) -> Self {
Self {
width,
height,
bitrate_kbps,
codec: codec.into(),
}
}
#[must_use]
pub fn pixel_count(&self) -> u64 {
u64::from(self.width) * u64::from(self.height)
}
#[must_use]
pub fn aspect_ratio(&self) -> f32 {
self.width as f32 / self.height as f32
}
}
#[derive(Debug, Clone)]
pub struct AbrLadder {
pub renditions: Vec<ResolutionBitrate>,
}
impl AbrLadder {
#[must_use]
pub fn new(mut renditions: Vec<ResolutionBitrate>) -> Self {
renditions.sort_by(|a, b| b.bitrate_kbps.cmp(&a.bitrate_kbps));
Self { renditions }
}
#[must_use]
pub fn optimal_bitrate_kbps(&self, bandwidth_kbps: u32) -> Option<&ResolutionBitrate> {
let threshold = (f64::from(bandwidth_kbps) * 0.8) as u32;
self.renditions.iter().find(|r| r.bitrate_kbps <= threshold)
}
#[must_use]
pub fn len(&self) -> usize {
self.renditions.len()
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.renditions.is_empty()
}
#[must_use]
pub fn highest_quality(&self) -> Option<&ResolutionBitrate> {
self.renditions.first()
}
#[must_use]
pub fn lowest_quality(&self) -> Option<&ResolutionBitrate> {
self.renditions.last()
}
}
#[derive(Debug, Clone, Default)]
pub struct AbrLadderGenerator;
impl AbrLadderGenerator {
#[must_use]
pub fn new() -> Self {
Self
}
#[must_use]
pub fn generate(&self, target_resolution: (u32, u32), codec: &str) -> AbrLadder {
let (_target_w, target_h) = target_resolution;
let standard_rungs: &[(u32, u32, u32)] = &[
(1920, 1080, 8000),
(1280, 720, 4500),
(960, 540, 2000),
(640, 360, 800),
(426, 240, 300),
];
let renditions = standard_rungs
.iter()
.filter(|(_, h, _)| *h <= target_h)
.map(|(w, h, kbps)| ResolutionBitrate::new(*w, *h, *kbps, codec))
.collect();
AbrLadder::new(renditions)
}
}
#[derive(Debug, Clone, Default)]
pub struct PerTitleEncoding;
impl PerTitleEncoding {
#[must_use]
pub fn new() -> Self {
Self
}
#[must_use]
pub fn analyze_complexity(frame_variance: &[f32]) -> f32 {
if frame_variance.is_empty() {
return 0.0;
}
let sum: f32 = frame_variance.iter().sum();
sum / frame_variance.len() as f32
}
#[must_use]
pub fn optimize_ladder(base_ladder: &AbrLadder, complexity: f32) -> AbrLadder {
let factor = complexity.clamp(0.5, 2.0);
let adjusted = base_ladder
.renditions
.iter()
.map(|r| {
let new_bitrate = ((r.bitrate_kbps as f32) * factor).round() as u32;
ResolutionBitrate::new(r.width, r.height, new_bitrate, r.codec.clone())
})
.collect();
AbrLadder::new(adjusted)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct BandwidthPoint {
pub timestamp_ms: u64,
pub kbps: u32,
}
#[derive(Debug, Clone)]
pub struct BandwidthEstimator {
samples: VecDeque<BandwidthPoint>,
capacity: usize,
}
impl BandwidthEstimator {
#[must_use]
pub fn new() -> Self {
Self {
samples: VecDeque::with_capacity(30),
capacity: 30,
}
}
pub fn add_sample(&mut self, point: BandwidthPoint) {
if self.samples.len() >= self.capacity {
self.samples.pop_front();
}
self.samples.push_back(point);
}
#[must_use]
pub fn sample_count(&self) -> usize {
self.samples.len()
}
#[must_use]
pub fn estimated_kbps(&self) -> Option<u32> {
if self.samples.is_empty() {
return None;
}
let sum: u64 = self.samples.iter().map(|s| u64::from(s.kbps)).sum();
Some((sum / self.samples.len() as u64) as u32)
}
#[must_use]
pub fn samples(&self) -> &VecDeque<BandwidthPoint> {
&self.samples
}
}
impl Default for BandwidthEstimator {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_resolution_bitrate_new() {
let r = ResolutionBitrate::new(1920, 1080, 8000, "h264");
assert_eq!(r.width, 1920);
assert_eq!(r.height, 1080);
assert_eq!(r.bitrate_kbps, 8000);
assert_eq!(r.codec, "h264");
}
#[test]
fn test_resolution_bitrate_pixel_count() {
let r = ResolutionBitrate::new(1920, 1080, 8000, "h264");
assert_eq!(r.pixel_count(), 1920 * 1080);
}
#[test]
fn test_resolution_bitrate_aspect_ratio() {
let r = ResolutionBitrate::new(1920, 1080, 8000, "h264");
let ar = r.aspect_ratio();
assert!((ar - 16.0 / 9.0).abs() < 0.01);
}
#[test]
fn test_abr_ladder_sorted_descending() {
let renditions = vec![
ResolutionBitrate::new(640, 360, 800, "h264"),
ResolutionBitrate::new(1920, 1080, 8000, "h264"),
ResolutionBitrate::new(1280, 720, 4500, "h264"),
];
let ladder = AbrLadder::new(renditions);
assert_eq!(ladder.renditions[0].bitrate_kbps, 8000);
assert_eq!(ladder.renditions[1].bitrate_kbps, 4500);
assert_eq!(ladder.renditions[2].bitrate_kbps, 800);
}
#[test]
fn test_abr_ladder_optimal_bitrate() {
let gen = AbrLadderGenerator::new();
let ladder = gen.generate((1920, 1080), "h264");
let opt = ladder.optimal_bitrate_kbps(10000);
assert!(opt.is_some());
assert_eq!(opt.expect("should succeed in test").bitrate_kbps, 8000);
let opt2 = ladder.optimal_bitrate_kbps(5000);
assert!(opt2.is_some());
assert_eq!(opt2.expect("should succeed in test").bitrate_kbps, 2000);
}
#[test]
fn test_abr_ladder_optimal_bitrate_too_low() {
let gen = AbrLadderGenerator::new();
let ladder = gen.generate((1920, 1080), "h264");
let opt = ladder.optimal_bitrate_kbps(100);
assert!(opt.is_none());
}
#[test]
fn test_abr_ladder_empty_check() {
let ladder = AbrLadder::new(vec![]);
assert!(ladder.is_empty());
assert_eq!(ladder.len(), 0);
assert!(ladder.highest_quality().is_none());
assert!(ladder.lowest_quality().is_none());
}
#[test]
fn test_abr_ladder_generator_full_1080p() {
let gen = AbrLadderGenerator::new();
let ladder = gen.generate((1920, 1080), "h264");
assert_eq!(ladder.len(), 5);
assert_eq!(
ladder
.highest_quality()
.expect("should succeed in test")
.bitrate_kbps,
8000
);
assert_eq!(
ladder
.lowest_quality()
.expect("should succeed in test")
.bitrate_kbps,
300
);
}
#[test]
fn test_abr_ladder_generator_720p_limit() {
let gen = AbrLadderGenerator::new();
let ladder = gen.generate((1280, 720), "vp9");
for r in &ladder.renditions {
assert!(r.height <= 720);
}
}
#[test]
fn test_per_title_analyze_complexity_empty() {
let c = PerTitleEncoding::analyze_complexity(&[]);
assert_eq!(c, 0.0);
}
#[test]
fn test_per_title_analyze_complexity() {
let variances = vec![1.0, 2.0, 3.0, 4.0];
let c = PerTitleEncoding::analyze_complexity(&variances);
assert!((c - 2.5).abs() < 1e-5);
}
#[test]
fn test_per_title_optimize_ladder_clamping() {
let gen = AbrLadderGenerator::new();
let base = gen.generate((1920, 1080), "h264");
let optimized = PerTitleEncoding::optimize_ladder(&base, 3.0);
let base_highest = base
.highest_quality()
.expect("should succeed in test")
.bitrate_kbps;
let opt_highest = optimized
.highest_quality()
.expect("should succeed in test")
.bitrate_kbps;
assert_eq!(opt_highest, (base_highest as f32 * 2.0).round() as u32);
}
#[test]
fn test_per_title_optimize_ladder_lower() {
let gen = AbrLadderGenerator::new();
let base = gen.generate((1920, 1080), "h264");
let optimized = PerTitleEncoding::optimize_ladder(&base, 0.5);
let base_highest = base
.highest_quality()
.expect("should succeed in test")
.bitrate_kbps;
let opt_highest = optimized
.highest_quality()
.expect("should succeed in test")
.bitrate_kbps;
assert_eq!(opt_highest, (base_highest as f32 * 0.5).round() as u32);
}
#[test]
fn test_bandwidth_estimator_add_sample() {
let mut est = BandwidthEstimator::new();
assert_eq!(est.sample_count(), 0);
assert!(est.estimated_kbps().is_none());
est.add_sample(BandwidthPoint {
timestamp_ms: 0,
kbps: 5000,
});
assert_eq!(est.sample_count(), 1);
assert_eq!(est.estimated_kbps(), Some(5000));
}
#[test]
fn test_bandwidth_estimator_ring_buffer() {
let mut est = BandwidthEstimator::new();
for i in 0..35u64 {
est.add_sample(BandwidthPoint {
timestamp_ms: i * 1000,
kbps: 1000,
});
}
assert_eq!(est.sample_count(), 30);
}
#[test]
fn test_bandwidth_estimator_average() {
let mut est = BandwidthEstimator::new();
est.add_sample(BandwidthPoint {
timestamp_ms: 0,
kbps: 4000,
});
est.add_sample(BandwidthPoint {
timestamp_ms: 1000,
kbps: 6000,
});
assert_eq!(est.estimated_kbps(), Some(5000));
}
}