rodio/source/
agc.rs

1//
2//      Automatic Gain Control (AGC) Algorithm
3//      Designed by @UnknownSuperficialNight
4//
5//   Features:
6//   • Adaptive peak detection
7//   • RMS-based level estimation
8//   • Asymmetric attack/release
9//   • RMS-based general adjustments with peak limiting
10//
11//   Optimized for smooth and responsive gain control
12//
13//   Crafted with love. Enjoy! :)
14//
15
16use super::SeekError;
17use crate::{Sample, Source};
18#[cfg(feature = "experimental")]
19use atomic_float::AtomicF32;
20#[cfg(feature = "experimental")]
21use std::sync::atomic::{AtomicBool, Ordering};
22#[cfg(feature = "experimental")]
23use std::sync::Arc;
24use std::time::Duration;
25
26#[cfg(feature = "tracing")]
27use tracing;
28
29/// Ensures `RMS_WINDOW_SIZE` is a power of two
30const fn power_of_two(n: usize) -> usize {
31    assert!(
32        n.is_power_of_two(),
33        "RMS_WINDOW_SIZE must be a power of two"
34    );
35    n
36}
37
38/// Size of the circular buffer used for RMS calculation.
39/// A larger size provides more stable RMS values but increases latency.
40const RMS_WINDOW_SIZE: usize = power_of_two(8192);
41
42#[cfg(feature = "experimental")]
43/// Automatic Gain Control filter for maintaining consistent output levels.
44///
45/// This struct implements an AGC algorithm that dynamically adjusts audio levels
46/// based on both **peak** and **RMS** (Root Mean Square) measurements.
47#[derive(Clone, Debug)]
48pub struct AutomaticGainControl<I> {
49    input: I,
50    target_level: Arc<AtomicF32>,
51    absolute_max_gain: Arc<AtomicF32>,
52    current_gain: f32,
53    attack_coeff: Arc<AtomicF32>,
54    release_coeff: Arc<AtomicF32>,
55    min_attack_coeff: f32,
56    peak_level: f32,
57    rms_window: CircularBuffer,
58    is_enabled: Arc<AtomicBool>,
59}
60
61#[cfg(not(feature = "experimental"))]
62/// Automatic Gain Control filter for maintaining consistent output levels.
63///
64/// This struct implements an AGC algorithm that dynamically adjusts audio levels
65/// based on both **peak** and **RMS** (Root Mean Square) measurements.
66#[derive(Clone, Debug)]
67pub struct AutomaticGainControl<I> {
68    input: I,
69    target_level: f32,
70    absolute_max_gain: f32,
71    current_gain: f32,
72    attack_coeff: f32,
73    release_coeff: f32,
74    min_attack_coeff: f32,
75    peak_level: f32,
76    rms_window: CircularBuffer,
77    is_enabled: bool,
78}
79
80/// A circular buffer for efficient RMS calculation over a sliding window.
81///
82/// This structure allows for constant-time updates and mean calculations,
83/// which is crucial for real-time audio processing.
84#[derive(Clone, Debug)]
85struct CircularBuffer {
86    buffer: Box<[f32; RMS_WINDOW_SIZE]>,
87    sum: f32,
88    index: usize,
89}
90
91impl CircularBuffer {
92    /// Creates a new `CircularBuffer` with a fixed size determined at compile time.
93    #[inline]
94    fn new() -> Self {
95        CircularBuffer {
96            buffer: Box::new([0.0; RMS_WINDOW_SIZE]),
97            sum: 0.0,
98            index: 0,
99        }
100    }
101
102    /// Pushes a new value into the buffer and returns the old value.
103    ///
104    /// This method maintains a running sum for efficient mean calculation.
105    #[inline]
106    fn push(&mut self, value: f32) -> f32 {
107        let old_value = self.buffer[self.index];
108        // Update the sum by first subtracting the old value and then adding the new value; this is more accurate.
109        self.sum = self.sum - old_value + value;
110        self.buffer[self.index] = value;
111        // Use bitwise AND for efficient index wrapping since RMS_WINDOW_SIZE is a power of two.
112        self.index = (self.index + 1) & (RMS_WINDOW_SIZE - 1);
113        old_value
114    }
115
116    /// Calculates the mean of all values in the buffer.
117    ///
118    /// This operation is `O(1)` due to the maintained running sum.
119    #[inline]
120    fn mean(&self) -> f32 {
121        self.sum / RMS_WINDOW_SIZE as f32
122    }
123}
124
125/// Constructs an `AutomaticGainControl` object with specified parameters.
126///
127/// # Arguments
128///
129/// * `input` - The input audio source
130/// * `target_level` - The desired output level
131/// * `attack_time` - Time constant for gain increase
132/// * `release_time` - Time constant for gain decrease
133/// * `absolute_max_gain` - Maximum allowable gain
134#[inline]
135pub(crate) fn automatic_gain_control<I>(
136    input: I,
137    target_level: f32,
138    attack_time: f32,
139    release_time: f32,
140    absolute_max_gain: f32,
141) -> AutomaticGainControl<I>
142where
143    I: Source,
144    I::Item: Sample,
145{
146    let sample_rate = input.sample_rate();
147    let attack_coeff = (-1.0 / (attack_time * sample_rate as f32)).exp();
148    let release_coeff = (-1.0 / (release_time * sample_rate as f32)).exp();
149
150    #[cfg(feature = "experimental")]
151    {
152        AutomaticGainControl {
153            input,
154            target_level: Arc::new(AtomicF32::new(target_level)),
155            absolute_max_gain: Arc::new(AtomicF32::new(absolute_max_gain)),
156            current_gain: 1.0,
157            attack_coeff: Arc::new(AtomicF32::new(attack_coeff)),
158            release_coeff: Arc::new(AtomicF32::new(release_coeff)),
159            min_attack_coeff: release_time,
160            peak_level: 0.0,
161            rms_window: CircularBuffer::new(),
162            is_enabled: Arc::new(AtomicBool::new(true)),
163        }
164    }
165
166    #[cfg(not(feature = "experimental"))]
167    {
168        AutomaticGainControl {
169            input,
170            target_level,
171            absolute_max_gain,
172            current_gain: 1.0,
173            attack_coeff,
174            release_coeff,
175            min_attack_coeff: release_time,
176            peak_level: 0.0,
177            rms_window: CircularBuffer::new(),
178            is_enabled: true,
179        }
180    }
181}
182
183impl<I> AutomaticGainControl<I>
184where
185    I: Source,
186    I::Item: Sample,
187{
188    #[inline]
189    fn target_level(&self) -> f32 {
190        #[cfg(feature = "experimental")]
191        {
192            self.target_level.load(Ordering::Relaxed)
193        }
194        #[cfg(not(feature = "experimental"))]
195        {
196            self.target_level
197        }
198    }
199
200    #[inline]
201    fn absolute_max_gain(&self) -> f32 {
202        #[cfg(feature = "experimental")]
203        {
204            self.absolute_max_gain.load(Ordering::Relaxed)
205        }
206        #[cfg(not(feature = "experimental"))]
207        {
208            self.absolute_max_gain
209        }
210    }
211
212    #[inline]
213    fn attack_coeff(&self) -> f32 {
214        #[cfg(feature = "experimental")]
215        {
216            self.attack_coeff.load(Ordering::Relaxed)
217        }
218        #[cfg(not(feature = "experimental"))]
219        {
220            self.attack_coeff
221        }
222    }
223
224    #[inline]
225    fn release_coeff(&self) -> f32 {
226        #[cfg(feature = "experimental")]
227        {
228            self.release_coeff.load(Ordering::Relaxed)
229        }
230        #[cfg(not(feature = "experimental"))]
231        {
232            self.release_coeff
233        }
234    }
235
236    #[inline]
237    fn is_enabled(&self) -> bool {
238        #[cfg(feature = "experimental")]
239        {
240            self.is_enabled.load(Ordering::Relaxed)
241        }
242        #[cfg(not(feature = "experimental"))]
243        {
244            self.is_enabled
245        }
246    }
247
248    #[cfg(feature = "experimental")]
249    /// Access the target output level for real-time adjustment.
250    ///
251    /// Use this to dynamically modify the AGC's target level while audio is processing.
252    /// Adjust this value to control the overall output amplitude of the processed signal.
253    #[inline]
254    pub fn get_target_level(&self) -> Arc<AtomicF32> {
255        Arc::clone(&self.target_level)
256    }
257
258    #[cfg(feature = "experimental")]
259    /// Access the maximum gain limit for real-time adjustment.
260    ///
261    /// Use this to dynamically modify the AGC's maximum allowable gain during runtime.
262    /// Adjusting this value helps prevent excessive amplification in low-level signals.
263    #[inline]
264    pub fn get_absolute_max_gain(&self) -> Arc<AtomicF32> {
265        Arc::clone(&self.absolute_max_gain)
266    }
267
268    #[cfg(feature = "experimental")]
269    /// Access the attack coefficient for real-time adjustment.
270    ///
271    /// Use this to dynamically modify how quickly the AGC responds to level increases.
272    /// Smaller values result in faster response, larger values in slower response.
273    /// Adjust during runtime to fine-tune AGC behavior for different audio content.
274    #[inline]
275    pub fn get_attack_coeff(&self) -> Arc<AtomicF32> {
276        Arc::clone(&self.attack_coeff)
277    }
278
279    #[cfg(feature = "experimental")]
280    /// Access the release coefficient for real-time adjustment.
281    ///
282    /// Use this to dynamically modify how quickly the AGC responds to level decreases.
283    /// Smaller values result in faster response, larger values in slower response.
284    /// Adjust during runtime to optimize AGC behavior for varying audio dynamics.
285    #[inline]
286    pub fn get_release_coeff(&self) -> Arc<AtomicF32> {
287        Arc::clone(&self.release_coeff)
288    }
289
290    #[cfg(feature = "experimental")]
291    /// Access the AGC on/off control for real-time adjustment.
292    ///
293    /// Use this to dynamically enable or disable AGC processing during runtime.
294    /// Useful for comparing processed and unprocessed audio or for disabling/enabling AGC at runtime.
295    #[inline]
296    pub fn get_agc_control(&self) -> Arc<AtomicBool> {
297        Arc::clone(&self.is_enabled)
298    }
299
300    #[cfg(not(feature = "experimental"))]
301    /// Enable or disable AGC processing.
302    ///
303    /// Use this to enable or disable AGC processing.
304    /// Useful for comparing processed and unprocessed audio or for disabling/enabling AGC.
305    #[inline]
306    pub fn set_enabled(&mut self, enabled: bool) {
307        self.is_enabled = enabled;
308    }
309
310    /// Updates the peak level with an adaptive attack coefficient
311    ///
312    /// This method adjusts the peak level using a variable attack coefficient.
313    /// It responds faster to sudden increases in signal level by using a
314    /// minimum attack coefficient of `min_attack_coeff` when the sample value exceeds the
315    /// current peak level. This adaptive behavior helps capture transients
316    /// more accurately while maintaining smoother behavior for gradual changes.
317    #[inline]
318    fn update_peak_level(&mut self, sample_value: f32) {
319        let attack_coeff = if sample_value > self.peak_level {
320            self.attack_coeff().min(self.min_attack_coeff) // User-defined attack time limited via release_time
321        } else {
322            self.release_coeff()
323        };
324
325        self.peak_level = attack_coeff * self.peak_level + (1.0 - attack_coeff) * sample_value;
326    }
327
328    /// Updates the RMS (Root Mean Square) level using a circular buffer approach.
329    /// This method calculates a moving average of the squared input samples,
330    /// providing a measure of the signal's average power over time.
331    #[inline]
332    fn update_rms(&mut self, sample_value: f32) -> f32 {
333        let squared_sample = sample_value * sample_value;
334        self.rms_window.push(squared_sample);
335        self.rms_window.mean().sqrt()
336    }
337
338    /// Calculate gain adjustments based on peak levels
339    /// This method determines the appropriate gain level to apply to the audio
340    /// signal, considering the peak level.
341    /// The peak level helps prevent sudden spikes in the output signal.
342    #[inline]
343    fn calculate_peak_gain(&self) -> f32 {
344        if self.peak_level > 0.0 {
345            (self.target_level() / self.peak_level).min(self.absolute_max_gain())
346        } else {
347            self.absolute_max_gain()
348        }
349    }
350
351    #[inline]
352    fn process_sample(&mut self, sample: I::Item) -> I::Item {
353        // Convert the sample to its absolute float value for level calculations
354        let sample_value = sample.to_f32().abs();
355
356        // Dynamically adjust peak level using an adaptive attack coefficient
357        self.update_peak_level(sample_value);
358
359        // Calculate the current RMS (Root Mean Square) level using a sliding window approach
360        let rms = self.update_rms(sample_value);
361
362        // Compute the gain adjustment required to reach the target level based on RMS
363        let rms_gain = if rms > 0.0 {
364            self.target_level() / rms
365        } else {
366            self.absolute_max_gain() // Default to max gain if RMS is zero
367        };
368
369        // Calculate the peak limiting gain
370        let peak_gain = self.calculate_peak_gain();
371
372        // Use RMS for general adjustments, but limit by peak gain to prevent clipping
373        let desired_gain = rms_gain.min(peak_gain);
374
375        // Adaptive attack/release speed for AGC (Automatic Gain Control)
376        //
377        // This mechanism implements an asymmetric approach to gain adjustment:
378        // 1. **Slow increase**: Prevents abrupt amplification of noise during quiet periods.
379        // 2. **Fast decrease**: Rapidly attenuates sudden loud signals to avoid distortion.
380        //
381        // The asymmetry is crucial because:
382        // - Gradual gain increases sound more natural and less noticeable to listeners.
383        // - Quick gain reductions are necessary to prevent clipping and maintain audio quality.
384        //
385        // This approach addresses several challenges associated with high attack times:
386        // 1. **Slow response**: With a high attack time, the AGC responds very slowly to changes in input level.
387        //    This means it takes longer for the gain to adjust to new signal levels.
388        // 2. **Initial gain calculation**: When the audio starts or after a period of silence, the initial gain
389        //    calculation might result in a very high gain value, especially if the input signal starts quietly.
390        // 3. **Overshooting**: As the gain slowly increases (due to the high attack time), it might overshoot
391        //    the desired level, causing the signal to become too loud.
392        // 4. **Overcorrection**: The AGC then tries to correct this by reducing the gain, but due to the slow response,
393        //    it might reduce the gain too much, causing the sound to drop to near-zero levels.
394        // 5. **Slow recovery**: Again, due to the high attack time, it takes a while for the gain to increase
395        //    back to the appropriate level.
396        //
397        // By using a faster release time for decreasing gain, we can mitigate these issues and provide
398        // more responsive control over sudden level increases while maintaining smooth gain increases.
399        let attack_speed = if desired_gain > self.current_gain {
400            self.attack_coeff()
401        } else {
402            self.release_coeff()
403        };
404
405        // Gradually adjust the current gain towards the desired gain for smooth transitions
406        self.current_gain = self.current_gain * attack_speed + desired_gain * (1.0 - attack_speed);
407
408        // Ensure the calculated gain stays within the defined operational range
409        self.current_gain = self.current_gain.clamp(0.1, self.absolute_max_gain());
410
411        // Output current gain value for developers to fine tune their inputs to automatic_gain_control
412        #[cfg(feature = "tracing")]
413        tracing::debug!("AGC gain: {}", self.current_gain,);
414
415        // Apply the computed gain to the input sample and return the result
416        sample.amplify(self.current_gain)
417    }
418
419    /// Returns a mutable reference to the inner source.
420    pub fn inner(&self) -> &I {
421        &self.input
422    }
423
424    /// Returns the inner source.
425    pub fn inner_mut(&mut self) -> &mut I {
426        &mut self.input
427    }
428}
429
430impl<I> Iterator for AutomaticGainControl<I>
431where
432    I: Source,
433    I::Item: Sample,
434{
435    type Item = I::Item;
436
437    #[inline]
438    fn next(&mut self) -> Option<I::Item> {
439        self.input.next().map(|sample| {
440            if self.is_enabled() {
441                self.process_sample(sample)
442            } else {
443                sample
444            }
445        })
446    }
447
448    #[inline]
449    fn size_hint(&self) -> (usize, Option<usize>) {
450        self.input.size_hint()
451    }
452}
453
454impl<I> ExactSizeIterator for AutomaticGainControl<I>
455where
456    I: Source + ExactSizeIterator,
457    I::Item: Sample,
458{
459}
460
461impl<I> Source for AutomaticGainControl<I>
462where
463    I: Source,
464    I::Item: Sample,
465{
466    #[inline]
467    fn current_frame_len(&self) -> Option<usize> {
468        self.input.current_frame_len()
469    }
470
471    #[inline]
472    fn channels(&self) -> u16 {
473        self.input.channels()
474    }
475
476    #[inline]
477    fn sample_rate(&self) -> u32 {
478        self.input.sample_rate()
479    }
480
481    #[inline]
482    fn total_duration(&self) -> Option<Duration> {
483        self.input.total_duration()
484    }
485
486    #[inline]
487    fn try_seek(&mut self, pos: Duration) -> Result<(), SeekError> {
488        self.input.try_seek(pos)
489    }
490}