#![allow(clippy::cast_lossless)]
#![allow(clippy::cast_precision_loss)]
#![allow(clippy::cast_possible_truncation)]
#![allow(clippy::cast_sign_loss)]
use std::collections::VecDeque;
#[derive(Debug, Clone)]
pub enum SimpleRateControlMode {
ConstantQuality {
crf: u8,
},
AverageBitrate {
target_kbps: u32,
},
ConstantBitrate {
target_kbps: u32,
vbv_size_kb: u32,
},
VariableBitrate {
min_kbps: u32,
max_kbps: u32,
target_kbps: u32,
},
TwoPass {
target_kbps: u32,
first_pass_stats: Option<String>,
},
}
#[derive(Debug, Clone)]
pub struct SimpleRateControlConfig {
pub mode: SimpleRateControlMode,
pub fps: f32,
pub resolution: (u32, u32),
pub keyframe_interval: u32,
pub b_frames: u8,
pub scene_change_sensitivity: u8,
}
impl SimpleRateControlConfig {
pub fn crf(crf: u8, fps: f32, resolution: (u32, u32)) -> Self {
Self {
mode: SimpleRateControlMode::ConstantQuality { crf },
fps,
resolution,
keyframe_interval: 120,
b_frames: 2,
scene_change_sensitivity: 40,
}
}
pub fn abr(target_kbps: u32, fps: f32, resolution: (u32, u32)) -> Self {
Self {
mode: SimpleRateControlMode::AverageBitrate { target_kbps },
fps,
resolution,
keyframe_interval: 120,
b_frames: 2,
scene_change_sensitivity: 40,
}
}
pub fn cbr(target_kbps: u32, vbv_size_kb: u32, fps: f32, resolution: (u32, u32)) -> Self {
Self {
mode: SimpleRateControlMode::ConstantBitrate {
target_kbps,
vbv_size_kb,
},
fps,
resolution,
keyframe_interval: 120,
b_frames: 0,
scene_change_sensitivity: 20,
}
}
}
#[derive(Debug, Clone)]
pub struct SimpleRateControlStats {
pub frames_encoded: u64,
pub total_bits: u64,
pub avg_bits_per_frame: f64,
pub avg_complexity: f32,
pub vbv_fullness: f64,
pub target_bitrate_kbps: u32,
pub actual_bitrate_kbps: f64,
}
const COMPLEXITY_HISTORY_LEN: usize = 64;
pub struct SimpleRateController {
config: SimpleRateControlConfig,
frame_count: u64,
bits_spent: u64,
target_bits: u64,
vbv_fullness: f64,
complexity_history: VecDeque<f32>,
first_pass_complexities: Vec<f32>,
first_pass_mean: f32,
}
impl SimpleRateController {
pub fn new(config: SimpleRateControlConfig) -> Self {
let target_bits = compute_target_bits_per_frame(&config);
let vbv_fullness = match &config.mode {
SimpleRateControlMode::ConstantBitrate { .. } => 0.5,
_ => 0.0,
};
let (first_pass_complexities, first_pass_mean) = parse_first_pass_stats(&config.mode);
Self {
config,
frame_count: 0,
bits_spent: 0,
target_bits,
vbv_fullness,
complexity_history: VecDeque::with_capacity(COMPLEXITY_HISTORY_LEN),
first_pass_complexities,
first_pass_mean,
}
}
fn avg_complexity(&self) -> f32 {
if self.complexity_history.is_empty() {
return 1.0;
}
let sum: f32 = self.complexity_history.iter().sum();
sum / self.complexity_history.len() as f32
}
pub fn target_bits_per_frame(&self) -> u64 {
self.target_bits
}
pub fn allocate_frame_bits(&mut self, complexity: f32) -> u32 {
let complexity = complexity.max(0.0001_f32);
match &self.config.mode.clone() {
SimpleRateControlMode::ConstantQuality { crf } => {
let (w, h) = self.config.resolution;
let qp = (*crf as f32).max(1.0);
let bits = (w as f64 * h as f64 / (qp * qp) as f64) as u64;
bits.min(u32::MAX as u64) as u32
}
SimpleRateControlMode::AverageBitrate { .. } => {
let avg = self.avg_complexity();
let ratio = (complexity / avg) as f64;
let allocated = (self.target_bits as f64 * ratio).round() as u64;
allocated.clamp(1, self.target_bits.saturating_mul(4)) as u32
}
SimpleRateControlMode::VariableBitrate {
min_kbps, max_kbps, ..
} => {
let min_bits = (*min_kbps as f64 * 1000.0 / self.config.fps as f64) as u64;
let max_bits = (*max_kbps as f64 * 1000.0 / self.config.fps as f64) as u64;
let avg = self.avg_complexity();
let ratio = (complexity / avg) as f64;
let allocated = (self.target_bits as f64 * ratio).round() as u64;
allocated.clamp(min_bits.max(1), max_bits.max(1)) as u32
}
SimpleRateControlMode::ConstantBitrate { vbv_size_kb, .. } => {
let avg = self.avg_complexity();
let ratio = (complexity / avg) as f64;
let mut allocated = (self.target_bits as f64 * ratio).round() as u64;
let vbv_bits = (*vbv_size_kb as u64).saturating_mul(1000);
if self.vbv_fullness > 0.9 {
let cap = (self.target_bits as f64 * 0.5).round() as u64;
allocated = allocated.min(cap);
} else if self.vbv_fullness < 0.3 {
let floor = (self.target_bits as f64 * 1.5).round() as u64;
allocated = allocated.max(floor);
}
allocated.clamp(1, vbv_bits.max(1)) as u32
}
SimpleRateControlMode::TwoPass { .. } => {
let idx = self.frame_count as usize;
if !self.first_pass_complexities.is_empty() && self.first_pass_mean > 0.0 {
let frame_cplx = if idx < self.first_pass_complexities.len() {
self.first_pass_complexities[idx]
} else {
self.first_pass_mean
};
let ratio = (frame_cplx / self.first_pass_mean) as f64;
let allocated = (self.target_bits as f64 * ratio).round() as u64;
allocated.clamp(1, self.target_bits.saturating_mul(4)) as u32
} else {
self.target_bits.clamp(1, u32::MAX as u64) as u32
}
}
}
}
pub fn record_frame(&mut self, actual_bits: u32, complexity: f32) {
self.bits_spent = self.bits_spent.saturating_add(actual_bits as u64);
self.frame_count = self.frame_count.saturating_add(1);
if self.complexity_history.len() >= COMPLEXITY_HISTORY_LEN {
self.complexity_history.pop_front();
}
self.complexity_history.push_back(complexity.max(0.0));
if let SimpleRateControlMode::ConstantBitrate {
target_kbps,
vbv_size_kb,
} = &self.config.mode
{
let vbv_capacity = (*vbv_size_kb as f64) * 1000.0;
if vbv_capacity > 0.0 {
let drained = actual_bits as f64;
let refilled = (*target_kbps as f64 * 1000.0) / self.config.fps as f64;
let delta = (refilled - drained) / vbv_capacity;
self.vbv_fullness = (self.vbv_fullness + delta).clamp(0.0, 1.0);
}
}
}
pub fn vbv_status(&self) -> f64 {
self.vbv_fullness
}
pub fn is_keyframe(&self) -> bool {
let interval = self.config.keyframe_interval.max(1) as u64;
self.frame_count % interval == 0
}
pub fn stats(&self) -> SimpleRateControlStats {
let target_bitrate_kbps = extract_target_kbps(&self.config.mode);
let avg_bits_per_frame = if self.frame_count == 0 {
0.0
} else {
self.bits_spent as f64 / self.frame_count as f64
};
let actual_bitrate_kbps = avg_bits_per_frame * self.config.fps as f64 / 1000.0;
SimpleRateControlStats {
frames_encoded: self.frame_count,
total_bits: self.bits_spent,
avg_bits_per_frame,
avg_complexity: self.avg_complexity(),
vbv_fullness: self.vbv_fullness,
target_bitrate_kbps,
actual_bitrate_kbps,
}
}
}
fn compute_target_bits_per_frame(config: &SimpleRateControlConfig) -> u64 {
let fps = config.fps.max(0.001_f32) as f64;
match &config.mode {
SimpleRateControlMode::ConstantQuality { .. } => 0,
SimpleRateControlMode::AverageBitrate { target_kbps }
| SimpleRateControlMode::ConstantBitrate { target_kbps, .. }
| SimpleRateControlMode::VariableBitrate { target_kbps, .. }
| SimpleRateControlMode::TwoPass { target_kbps, .. } => {
(*target_kbps as f64 * 1000.0 / fps).round() as u64
}
}
}
fn extract_target_kbps(mode: &SimpleRateControlMode) -> u32 {
match mode {
SimpleRateControlMode::ConstantQuality { .. } => 0,
SimpleRateControlMode::AverageBitrate { target_kbps }
| SimpleRateControlMode::ConstantBitrate { target_kbps, .. }
| SimpleRateControlMode::VariableBitrate { target_kbps, .. }
| SimpleRateControlMode::TwoPass { target_kbps, .. } => *target_kbps,
}
}
fn parse_first_pass_stats(mode: &SimpleRateControlMode) -> (Vec<f32>, f32) {
if let SimpleRateControlMode::TwoPass {
first_pass_stats: Some(stats),
..
} = mode
{
let values: Vec<f32> = stats
.lines()
.filter_map(|l| l.trim().parse::<f32>().ok())
.collect();
let mean = if values.is_empty() {
0.0
} else {
values.iter().sum::<f32>() / values.len() as f32
};
return (values, mean);
}
(Vec::new(), 0.0)
}
#[cfg(test)]
mod tests {
use super::*;
fn make_crf(crf: u8) -> SimpleRateController {
SimpleRateController::new(SimpleRateControlConfig::crf(crf, 30.0, (1920, 1080)))
}
fn make_abr(kbps: u32) -> SimpleRateController {
SimpleRateController::new(SimpleRateControlConfig::abr(kbps, 30.0, (1920, 1080)))
}
fn make_cbr(kbps: u32, vbv_kb: u32) -> SimpleRateController {
SimpleRateController::new(SimpleRateControlConfig::cbr(
kbps,
vbv_kb,
30.0,
(1920, 1080),
))
}
#[test]
fn crf_allocation_nonzero() {
let mut rc = make_crf(23);
let bits = rc.allocate_frame_bits(1.0);
assert!(bits > 0, "CRF allocation must be positive");
}
#[test]
fn crf_lower_crf_more_bits() {
let mut hi_q = make_crf(10);
let mut lo_q = make_crf(40);
let hi = hi_q.allocate_frame_bits(1.0);
let lo = lo_q.allocate_frame_bits(1.0);
assert!(hi > lo, "crf=10 should yield more bits than crf=40");
}
#[test]
fn abr_target_bits_per_frame() {
let rc = make_abr(4000);
let expected = (4_000_000_u64 / 30) as f64;
let got = rc.target_bits_per_frame() as f64;
assert!(
(got - expected).abs() / expected < 0.01,
"ABR target bits/frame off: got {got}, expected ~{expected}"
);
}
#[test]
fn crf_target_bits_zero() {
let rc = make_crf(28);
assert_eq!(rc.target_bits_per_frame(), 0);
}
#[test]
fn abr_high_complexity_more_bits() {
let mut rc = make_abr(4000);
for _ in 0..16 {
rc.record_frame(100_000, 1.0);
}
let low = rc.allocate_frame_bits(0.5);
let high = rc.allocate_frame_bits(2.0);
assert!(
high > low,
"higher complexity should yield more bits in ABR"
);
}
#[test]
fn record_frame_increments_count() {
let mut rc = make_abr(2000);
for _ in 0..10 {
rc.record_frame(50_000, 1.0);
}
assert_eq!(rc.stats().frames_encoded, 10);
}
#[test]
fn record_frame_accumulates_bits() {
let mut rc = make_abr(2000);
for _ in 0..5 {
rc.record_frame(10_000, 1.0);
}
assert_eq!(rc.stats().total_bits, 50_000);
}
#[test]
fn cbr_vbv_fullness_in_range() {
let mut rc = make_cbr(4000, 8000);
for i in 0..60 {
let bits = if i % 5 == 0 { 1_000_000 } else { 50_000 };
rc.record_frame(bits, 1.0);
let f = rc.vbv_status();
assert!((0.0..=1.0).contains(&f), "VBV fullness out of range: {f}");
}
}
#[test]
fn cbr_vbv_reduces_when_full() {
let mut rc = make_cbr(4000, 1000); for _ in 0..100 {
rc.record_frame(0, 1.0); }
let bits_full = rc.allocate_frame_bits(1.0);
let target = rc.target_bits_per_frame() as u32;
assert!(
bits_full <= target.saturating_mul(5) / 10 + 1,
"CBR should reduce bits when VBV is full; got {bits_full}, target {target}"
);
}
#[test]
fn is_keyframe_at_correct_intervals() {
let cfg = SimpleRateControlConfig {
mode: SimpleRateControlMode::AverageBitrate { target_kbps: 2000 },
fps: 30.0,
resolution: (1280, 720),
keyframe_interval: 10,
b_frames: 0,
scene_change_sensitivity: 0,
};
let mut rc = SimpleRateController::new(cfg);
assert!(rc.is_keyframe(), "frame_count=0 should be keyframe");
for i in 1..10u64 {
rc.record_frame(10_000, 1.0); assert_eq!(rc.frame_count, i);
assert!(!rc.is_keyframe(), "frame_count={i} should NOT be keyframe");
}
rc.record_frame(10_000, 1.0);
assert_eq!(rc.frame_count, 10);
assert!(rc.is_keyframe(), "frame_count=10 should be keyframe");
}
#[test]
fn stats_avg_bits_per_frame() {
let mut rc = make_abr(2000);
rc.record_frame(200_000, 1.0);
rc.record_frame(100_000, 1.0);
let s = rc.stats();
assert!((s.avg_bits_per_frame - 150_000.0).abs() < 1.0);
}
#[test]
fn vbr_allocation_clamped() {
let cfg = SimpleRateControlConfig {
mode: SimpleRateControlMode::VariableBitrate {
min_kbps: 1000,
max_kbps: 8000,
target_kbps: 4000,
},
fps: 30.0,
resolution: (1920, 1080),
keyframe_interval: 120,
b_frames: 2,
scene_change_sensitivity: 40,
};
let mut rc = SimpleRateController::new(cfg);
for _ in 0..16 {
rc.record_frame(130_000, 1.0);
}
let bits_low = rc.allocate_frame_bits(0.01); let bits_high = rc.allocate_frame_bits(100.0); let min_bits = (1000_u64 * 1000 / 30) as u32;
let max_bits = (8000_u64 * 1000 / 30) as u32;
assert!(
bits_low >= min_bits.saturating_sub(1),
"VBR low allocation {bits_low} below min {min_bits}"
);
assert!(
bits_high <= max_bits + 1,
"VBR high allocation {bits_high} above max {max_bits}"
);
}
#[test]
fn two_pass_uses_first_pass_stats() {
let stats = "0.5\n1.0\n2.0\n0.5\n1.0";
let cfg = SimpleRateControlConfig {
mode: SimpleRateControlMode::TwoPass {
target_kbps: 4000,
first_pass_stats: Some(stats.to_owned()),
},
fps: 30.0,
resolution: (1920, 1080),
keyframe_interval: 120,
b_frames: 2,
scene_change_sensitivity: 40,
};
let mut rc = SimpleRateController::new(cfg);
let bits_easy = rc.allocate_frame_bits(0.5);
rc.record_frame(bits_easy, 0.5);
let bits_avg = rc.allocate_frame_bits(1.0);
rc.record_frame(bits_avg, 1.0);
let bits_hard = rc.allocate_frame_bits(2.0);
assert!(
bits_hard > bits_easy,
"Two-pass: complex frame should get more bits than easy frame"
);
}
#[test]
fn stats_actual_bitrate_nonzero() {
let mut rc = make_abr(4000);
for _ in 0..30 {
rc.record_frame(133_333, 1.0);
}
let s = rc.stats();
assert!(
s.actual_bitrate_kbps > 0.0,
"actual bitrate should be positive"
);
}
#[test]
fn two_pass_no_stats_flat_fallback() {
let cfg = SimpleRateControlConfig {
mode: SimpleRateControlMode::TwoPass {
target_kbps: 2000,
first_pass_stats: None,
},
fps: 25.0,
resolution: (1280, 720),
keyframe_interval: 50,
b_frames: 2,
scene_change_sensitivity: 40,
};
let mut rc = SimpleRateController::new(cfg);
let bits = rc.allocate_frame_bits(1.0);
assert!(bits > 0, "fallback two-pass allocation must be positive");
let target = rc.target_bits_per_frame() as u32;
assert_eq!(
bits, target,
"no-stats two-pass should produce flat allocation equal to target"
);
}
}