Skip to main content

oximedia_codec/rate_control/
types.rs

1//! Rate control types and data structures.
2//!
3//! This module defines the core types used throughout the rate control system:
4//! - Configuration structures for different rate control modes
5//! - Statistics structures for frames and GOPs
6//! - Output structures for rate control decisions
7
8#![allow(clippy::cast_lossless)]
9#![allow(clippy::cast_precision_loss)]
10#![allow(clippy::cast_possible_truncation)]
11#![allow(clippy::cast_sign_loss)]
12#![allow(clippy::similar_names)]
13#![allow(clippy::missing_errors_doc)]
14#![allow(clippy::double_must_use)]
15#![allow(clippy::match_same_arms)]
16#![forbid(unsafe_code)]
17
18use crate::frame::FrameType;
19
20/// Rate control mode selection.
21#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)]
22pub enum RateControlMode {
23    /// Constant Quantization Parameter.
24    /// Uses a fixed QP for all frames.
25    Cqp,
26    /// Constant Bitrate.
27    /// Maintains a steady bitrate over time.
28    Cbr,
29    /// Variable Bitrate.
30    /// Allows bitrate to vary while staying within bounds.
31    Vbr,
32    /// Average Bitrate.
33    /// Targets an average bitrate over the entire encode.
34    Abr,
35    /// Constant Rate Factor (quality-based).
36    /// Maintains consistent perceptual quality.
37    #[default]
38    Crf,
39}
40
41impl RateControlMode {
42    /// Returns true if this mode targets a specific bitrate.
43    #[must_use]
44    pub fn is_bitrate_based(&self) -> bool {
45        matches!(self, Self::Cbr | Self::Vbr | Self::Abr)
46    }
47
48    /// Returns true if this mode targets consistent quality.
49    #[must_use]
50    pub fn is_quality_based(&self) -> bool {
51        matches!(self, Self::Cqp | Self::Crf)
52    }
53}
54
55/// Rate control configuration.
56#[derive(Clone, Debug)]
57pub struct RcConfig {
58    /// Rate control mode.
59    pub mode: RateControlMode,
60    /// Target bitrate in bits per second.
61    pub target_bitrate: u64,
62    /// Maximum bitrate in bits per second (for VBR).
63    pub max_bitrate: Option<u64>,
64    /// Minimum bitrate in bits per second (for VBR).
65    pub min_bitrate: Option<u64>,
66    /// Minimum QP value (higher quality limit).
67    pub min_qp: u8,
68    /// Maximum QP value (lower quality limit).
69    pub max_qp: u8,
70    /// Initial QP for the first frame.
71    pub initial_qp: u8,
72    /// CRF value for quality-based modes (0-63).
73    pub crf: f32,
74    /// VBV/HRD buffer size in bits.
75    pub buffer_size: u64,
76    /// Initial buffer fullness as a fraction (0.0-1.0).
77    pub initial_buffer_fullness: f32,
78    /// Frame rate numerator.
79    pub framerate_num: u32,
80    /// Frame rate denominator.
81    pub framerate_den: u32,
82    /// GOP (keyframe interval) length.
83    pub gop_length: u32,
84    /// Enable adaptive quantization.
85    pub enable_aq: bool,
86    /// AQ strength (0.0-2.0).
87    pub aq_strength: f32,
88    /// Lookahead depth (number of frames).
89    pub lookahead_depth: usize,
90    /// B-frame QP offset from P-frame.
91    pub b_qp_offset: i8,
92    /// I-frame QP offset from P-frame (typically negative).
93    pub i_qp_offset: i8,
94    /// Enable scene cut detection.
95    pub scene_cut_detection: bool,
96    /// Scene cut threshold (0.0-1.0).
97    pub scene_cut_threshold: f32,
98}
99
100impl Default for RcConfig {
101    fn default() -> Self {
102        Self {
103            mode: RateControlMode::Crf,
104            target_bitrate: 5_000_000, // 5 Mbps
105            max_bitrate: None,
106            min_bitrate: None,
107            min_qp: 1,
108            max_qp: 63,
109            initial_qp: 28,
110            crf: 23.0,
111            buffer_size: 10_000_000, // 10 Mb
112            initial_buffer_fullness: 0.75,
113            framerate_num: 30,
114            framerate_den: 1,
115            gop_length: 250,
116            enable_aq: true,
117            aq_strength: 1.0,
118            lookahead_depth: 40,
119            b_qp_offset: 2,
120            i_qp_offset: -2,
121            scene_cut_detection: true,
122            scene_cut_threshold: 0.4,
123        }
124    }
125}
126
127impl RcConfig {
128    /// Create a CQP configuration.
129    #[must_use]
130    pub fn cqp(qp: u8) -> Self {
131        Self {
132            mode: RateControlMode::Cqp,
133            initial_qp: qp,
134            ..Default::default()
135        }
136    }
137
138    /// Create a CBR configuration.
139    #[must_use]
140    pub fn cbr(bitrate: u64) -> Self {
141        Self {
142            mode: RateControlMode::Cbr,
143            target_bitrate: bitrate,
144            max_bitrate: Some(bitrate),
145            buffer_size: bitrate, // 1 second buffer
146            ..Default::default()
147        }
148    }
149
150    /// Create a VBR configuration.
151    #[must_use]
152    pub fn vbr(target: u64, max: u64) -> Self {
153        Self {
154            mode: RateControlMode::Vbr,
155            target_bitrate: target,
156            max_bitrate: Some(max),
157            buffer_size: max * 2, // 2 second buffer at max rate
158            ..Default::default()
159        }
160    }
161
162    /// Create a CRF configuration.
163    #[must_use]
164    pub fn crf(crf_value: f32) -> Self {
165        Self {
166            mode: RateControlMode::Crf,
167            crf: crf_value.clamp(0.0, 63.0),
168            ..Default::default()
169        }
170    }
171
172    /// Get the frame rate as a floating point value.
173    #[must_use]
174    pub fn framerate(&self) -> f64 {
175        if self.framerate_den == 0 {
176            30.0
177        } else {
178            f64::from(self.framerate_num) / f64::from(self.framerate_den)
179        }
180    }
181
182    /// Get the target bits per frame.
183    #[must_use]
184    pub fn target_bits_per_frame(&self) -> u64 {
185        let fps = self.framerate();
186        if fps <= 0.0 {
187            return 0;
188        }
189        (self.target_bitrate as f64 / fps) as u64
190    }
191
192    /// Validate configuration parameters.
193    #[must_use]
194    pub fn validate(&self) -> Result<(), RcConfigError> {
195        if self.min_qp > self.max_qp {
196            return Err(RcConfigError::InvalidQpRange);
197        }
198        if self.initial_qp < self.min_qp || self.initial_qp > self.max_qp {
199            return Err(RcConfigError::InitialQpOutOfRange);
200        }
201        if self.crf < 0.0 || self.crf > 63.0 {
202            return Err(RcConfigError::InvalidCrf);
203        }
204        if self.framerate_den == 0 {
205            return Err(RcConfigError::InvalidFramerate);
206        }
207        if self.mode.is_bitrate_based() && self.target_bitrate == 0 {
208            return Err(RcConfigError::ZeroBitrate);
209        }
210        Ok(())
211    }
212}
213
214/// Configuration validation errors.
215#[derive(Clone, Copy, Debug, PartialEq, Eq)]
216pub enum RcConfigError {
217    /// `min_qp` is greater than `max_qp`.
218    InvalidQpRange,
219    /// `initial_qp` is outside `min_qp/max_qp` range.
220    InitialQpOutOfRange,
221    /// CRF value is out of valid range.
222    InvalidCrf,
223    /// Frame rate denominator is zero.
224    InvalidFramerate,
225    /// Target bitrate is zero for bitrate-based mode.
226    ZeroBitrate,
227}
228
229/// Statistics for a single encoded frame.
230#[derive(Clone, Debug, Default)]
231pub struct FrameStats {
232    /// Frame number (display order).
233    pub frame_num: u64,
234    /// Frame type (I/P/B).
235    pub frame_type: FrameType,
236    /// Bits used for this frame.
237    pub bits: u64,
238    /// QP used for this frame.
239    pub qp: u8,
240    /// QP as floating point for averaging.
241    pub qp_f: f32,
242    /// Spatial complexity estimate.
243    pub spatial_complexity: f32,
244    /// Temporal complexity estimate (motion).
245    pub temporal_complexity: f32,
246    /// Combined complexity metric.
247    pub complexity: f32,
248    /// PSNR if calculated.
249    pub psnr: Option<f32>,
250    /// SSIM if calculated.
251    pub ssim: Option<f32>,
252    /// Was this frame a scene cut.
253    pub scene_cut: bool,
254    /// Target bits for this frame.
255    pub target_bits: u64,
256    /// Actual bits / target bits ratio.
257    pub bit_accuracy: f32,
258    /// Encoding time in microseconds.
259    pub encode_time_us: u64,
260}
261
262impl FrameStats {
263    /// Create new frame statistics.
264    #[must_use]
265    pub fn new(frame_num: u64, frame_type: FrameType) -> Self {
266        Self {
267            frame_num,
268            frame_type,
269            ..Default::default()
270        }
271    }
272
273    /// Calculate bits per pixel.
274    #[must_use]
275    pub fn bits_per_pixel(&self, width: u32, height: u32) -> f32 {
276        let pixels = u64::from(width) * u64::from(height);
277        if pixels == 0 {
278            return 0.0;
279        }
280        self.bits as f32 / pixels as f32
281    }
282
283    /// Check if frame exceeded target significantly.
284    #[must_use]
285    pub fn exceeded_target(&self, threshold: f32) -> bool {
286        self.target_bits > 0 && self.bit_accuracy > (1.0 + threshold)
287    }
288
289    /// Check if frame was under target significantly.
290    #[must_use]
291    pub fn under_target(&self, threshold: f32) -> bool {
292        self.target_bits > 0 && self.bit_accuracy < (1.0 - threshold)
293    }
294}
295
296/// Statistics for a Group of Pictures (GOP).
297#[derive(Clone, Debug, Default)]
298pub struct GopStats {
299    /// GOP index.
300    pub gop_index: u64,
301    /// Number of frames in this GOP.
302    pub frame_count: u32,
303    /// Number of I-frames.
304    pub i_frame_count: u32,
305    /// Number of P-frames.
306    pub p_frame_count: u32,
307    /// Number of B-frames.
308    pub b_frame_count: u32,
309    /// Total bits for this GOP.
310    pub total_bits: u64,
311    /// Target bits for this GOP.
312    pub target_bits: u64,
313    /// Average QP across all frames.
314    pub average_qp: f32,
315    /// Average complexity.
316    pub average_complexity: f32,
317    /// Total complexity (sum of frame complexities).
318    pub total_complexity: f32,
319    /// First frame number in this GOP.
320    pub first_frame: u64,
321    /// Last frame number in this GOP.
322    pub last_frame: u64,
323    /// Accumulated frame statistics.
324    frames: Vec<FrameStats>,
325}
326
327impl GopStats {
328    /// Create new GOP statistics.
329    #[must_use]
330    pub fn new(gop_index: u64, first_frame: u64) -> Self {
331        Self {
332            gop_index,
333            first_frame,
334            last_frame: first_frame,
335            ..Default::default()
336        }
337    }
338
339    /// Add frame statistics to this GOP.
340    pub fn add_frame(&mut self, stats: FrameStats) {
341        self.last_frame = stats.frame_num;
342        self.total_bits += stats.bits;
343        self.total_complexity += stats.complexity;
344
345        match stats.frame_type {
346            FrameType::Key => self.i_frame_count += 1,
347            FrameType::Inter => self.p_frame_count += 1,
348            FrameType::BiDir => self.b_frame_count += 1,
349            FrameType::Switch => self.p_frame_count += 1,
350        }
351
352        self.frames.push(stats);
353        self.frame_count = self.frames.len() as u32;
354
355        // Update averages
356        if self.frame_count > 0 {
357            let fc = self.frame_count as f32;
358            self.average_qp = self.frames.iter().map(|f| f.qp_f).sum::<f32>() / fc;
359            self.average_complexity = self.total_complexity / fc;
360        }
361    }
362
363    /// Get frame statistics.
364    #[must_use]
365    pub fn frames(&self) -> &[FrameStats] {
366        &self.frames
367    }
368
369    /// Get bit accuracy for this GOP.
370    #[must_use]
371    pub fn bit_accuracy(&self) -> f32 {
372        if self.target_bits == 0 {
373            return 1.0;
374        }
375        self.total_bits as f32 / self.target_bits as f32
376    }
377
378    /// Get bits per I-frame.
379    #[must_use]
380    pub fn bits_per_i_frame(&self) -> f64 {
381        if self.i_frame_count == 0 {
382            return 0.0;
383        }
384        let i_bits: u64 = self
385            .frames
386            .iter()
387            .filter(|f| f.frame_type == FrameType::Key)
388            .map(|f| f.bits)
389            .sum();
390        i_bits as f64 / f64::from(self.i_frame_count)
391    }
392
393    /// Get bits per P-frame.
394    #[must_use]
395    pub fn bits_per_p_frame(&self) -> f64 {
396        if self.p_frame_count == 0 {
397            return 0.0;
398        }
399        let p_bits: u64 = self
400            .frames
401            .iter()
402            .filter(|f| f.frame_type == FrameType::Inter)
403            .map(|f| f.bits)
404            .sum();
405        p_bits as f64 / f64::from(self.p_frame_count)
406    }
407
408    /// Get bits per B-frame.
409    #[must_use]
410    pub fn bits_per_b_frame(&self) -> f64 {
411        if self.b_frame_count == 0 {
412            return 0.0;
413        }
414        let b_bits: u64 = self
415            .frames
416            .iter()
417            .filter(|f| f.frame_type == FrameType::BiDir)
418            .map(|f| f.bits)
419            .sum();
420        b_bits as f64 / f64::from(self.b_frame_count)
421    }
422}
423
424/// Rate control output for a single frame.
425#[derive(Clone, Debug, Default)]
426pub struct RcOutput {
427    /// Quantization parameter to use.
428    pub qp: u8,
429    /// QP as floating point (before rounding).
430    pub qp_f: f32,
431    /// Target bits for this frame.
432    pub target_bits: u64,
433    /// Minimum bits (for underflow prevention).
434    pub min_bits: u64,
435    /// Maximum bits (for overflow prevention).
436    pub max_bits: u64,
437    /// Whether to drop this frame.
438    pub drop_frame: bool,
439    /// Whether to force a keyframe.
440    pub force_keyframe: bool,
441    /// Block-level QP offsets for AQ.
442    pub qp_offsets: Option<Vec<f32>>,
443    /// Lambda for RDO (rate-distortion optimization).
444    pub lambda: f64,
445    /// Lambda for motion estimation.
446    pub lambda_me: f64,
447}
448
449impl RcOutput {
450    /// Create a new rate control output with the given QP.
451    #[must_use]
452    pub fn with_qp(qp: u8) -> Self {
453        Self {
454            qp,
455            qp_f: qp as f32,
456            ..Default::default()
457        }
458    }
459
460    /// Create output that drops the frame.
461    #[must_use]
462    pub fn drop() -> Self {
463        Self {
464            drop_frame: true,
465            ..Default::default()
466        }
467    }
468
469    /// Calculate lambda from QP using standard formula.
470    /// Lambda = 0.85 * 2^((QP-12)/3)
471    #[must_use]
472    pub fn qp_to_lambda(qp: f32) -> f64 {
473        0.85 * 2.0_f64.powf((f64::from(qp) - 12.0) / 3.0)
474    }
475
476    /// Calculate motion estimation lambda from lambda.
477    #[must_use]
478    pub fn lambda_to_lambda_me(lambda: f64) -> f64 {
479        lambda.sqrt()
480    }
481
482    /// Set lambda values based on QP.
483    pub fn compute_lambda(&mut self) {
484        self.lambda = Self::qp_to_lambda(self.qp_f);
485        self.lambda_me = Self::lambda_to_lambda_me(self.lambda);
486    }
487}
488
489/// Rate control state summary.
490#[derive(Clone, Debug, Default)]
491pub struct RcState {
492    /// Total frames processed.
493    pub frames_encoded: u64,
494    /// Total bits produced.
495    pub total_bits: u64,
496    /// Current buffer level.
497    pub buffer_level: i64,
498    /// Average bitrate so far.
499    pub average_bitrate: f64,
500    /// Average QP so far.
501    pub average_qp: f32,
502    /// Frames dropped due to rate control.
503    pub frames_dropped: u64,
504    /// Keyframes forced by scene detection.
505    pub scene_cuts_detected: u64,
506    /// Current GOP index.
507    pub current_gop: u64,
508}
509
510impl RcState {
511    /// Update state with new frame statistics.
512    pub fn update(&mut self, stats: &FrameStats, elapsed_time: f64) {
513        self.frames_encoded += 1;
514        self.total_bits += stats.bits;
515
516        if elapsed_time > 0.0 {
517            self.average_bitrate = self.total_bits as f64 / elapsed_time;
518        }
519
520        // Running average of QP
521        let n = self.frames_encoded as f32;
522        self.average_qp = self.average_qp * (n - 1.0) / n + stats.qp_f / n;
523
524        if stats.scene_cut {
525            self.scene_cuts_detected += 1;
526        }
527    }
528}
529
530#[cfg(test)]
531mod tests {
532    use super::*;
533
534    #[test]
535    fn test_rate_control_mode() {
536        assert!(RateControlMode::Cbr.is_bitrate_based());
537        assert!(RateControlMode::Vbr.is_bitrate_based());
538        assert!(RateControlMode::Abr.is_bitrate_based());
539        assert!(!RateControlMode::Cqp.is_bitrate_based());
540        assert!(!RateControlMode::Crf.is_bitrate_based());
541
542        assert!(RateControlMode::Cqp.is_quality_based());
543        assert!(RateControlMode::Crf.is_quality_based());
544        assert!(!RateControlMode::Cbr.is_quality_based());
545    }
546
547    #[test]
548    fn test_rc_config_creation() {
549        let config = RcConfig::cqp(28);
550        assert_eq!(config.mode, RateControlMode::Cqp);
551        assert_eq!(config.initial_qp, 28);
552
553        let config = RcConfig::cbr(5_000_000);
554        assert_eq!(config.mode, RateControlMode::Cbr);
555        assert_eq!(config.target_bitrate, 5_000_000);
556
557        let config = RcConfig::vbr(5_000_000, 10_000_000);
558        assert_eq!(config.mode, RateControlMode::Vbr);
559        assert_eq!(config.max_bitrate, Some(10_000_000));
560
561        let config = RcConfig::crf(23.0);
562        assert_eq!(config.mode, RateControlMode::Crf);
563        assert!((config.crf - 23.0).abs() < f32::EPSILON);
564    }
565
566    #[test]
567    fn test_rc_config_validation() {
568        let mut config = RcConfig::default();
569        assert!(config.validate().is_ok());
570
571        config.min_qp = 50;
572        config.max_qp = 30;
573        assert_eq!(config.validate(), Err(RcConfigError::InvalidQpRange));
574
575        config.min_qp = 10;
576        config.max_qp = 40;
577        config.initial_qp = 5;
578        assert_eq!(config.validate(), Err(RcConfigError::InitialQpOutOfRange));
579
580        config.initial_qp = 25;
581        config.crf = 100.0;
582        assert_eq!(config.validate(), Err(RcConfigError::InvalidCrf));
583
584        config.crf = 23.0;
585        config.framerate_den = 0;
586        assert_eq!(config.validate(), Err(RcConfigError::InvalidFramerate));
587
588        config.framerate_den = 1;
589        config.mode = RateControlMode::Cbr;
590        config.target_bitrate = 0;
591        assert_eq!(config.validate(), Err(RcConfigError::ZeroBitrate));
592    }
593
594    #[test]
595    fn test_target_bits_per_frame() {
596        let config = RcConfig {
597            target_bitrate: 3_000_000,
598            framerate_num: 30,
599            framerate_den: 1,
600            ..Default::default()
601        };
602        assert_eq!(config.target_bits_per_frame(), 100_000);
603
604        let config = RcConfig {
605            target_bitrate: 6_000_000,
606            framerate_num: 60,
607            framerate_den: 1,
608            ..Default::default()
609        };
610        assert_eq!(config.target_bits_per_frame(), 100_000);
611    }
612
613    #[test]
614    fn test_frame_stats() {
615        let mut stats = FrameStats::new(0, FrameType::Key);
616        stats.bits = 100_000;
617        stats.target_bits = 80_000;
618        stats.bit_accuracy = stats.bits as f32 / stats.target_bits as f32;
619
620        assert!(stats.exceeded_target(0.1));
621        assert!(!stats.under_target(0.1));
622
623        let bpp = stats.bits_per_pixel(1920, 1080);
624        assert!(bpp > 0.0);
625    }
626
627    #[test]
628    fn test_gop_stats() {
629        let mut gop = GopStats::new(0, 0);
630
631        let mut i_frame = FrameStats::new(0, FrameType::Key);
632        i_frame.bits = 200_000;
633        i_frame.qp_f = 24.0;
634        i_frame.complexity = 1.5;
635        gop.add_frame(i_frame);
636
637        let mut p_frame = FrameStats::new(1, FrameType::Inter);
638        p_frame.bits = 50_000;
639        p_frame.qp_f = 26.0;
640        p_frame.complexity = 1.0;
641        gop.add_frame(p_frame);
642
643        assert_eq!(gop.frame_count, 2);
644        assert_eq!(gop.i_frame_count, 1);
645        assert_eq!(gop.p_frame_count, 1);
646        assert_eq!(gop.total_bits, 250_000);
647        assert!((gop.average_qp - 25.0).abs() < f32::EPSILON);
648    }
649
650    #[test]
651    fn test_rc_output_lambda() {
652        let lambda = RcOutput::qp_to_lambda(28.0);
653        assert!(lambda > 0.0);
654
655        let lambda_me = RcOutput::lambda_to_lambda_me(lambda);
656        assert!((lambda_me - lambda.sqrt()).abs() < f64::EPSILON);
657
658        let mut output = RcOutput::with_qp(28);
659        output.compute_lambda();
660        assert!(output.lambda > 0.0);
661        assert!(output.lambda_me > 0.0);
662    }
663
664    #[test]
665    fn test_rc_state_update() {
666        let mut state = RcState::default();
667        let mut stats = FrameStats::new(0, FrameType::Key);
668        stats.bits = 100_000;
669        stats.qp_f = 28.0;
670
671        state.update(&stats, 1.0 / 30.0);
672
673        assert_eq!(state.frames_encoded, 1);
674        assert_eq!(state.total_bits, 100_000);
675        assert!((state.average_qp - 28.0).abs() < f32::EPSILON);
676    }
677}