Skip to main content

firewheel_nodes/
fast_rms.rs

1use bevy_platform::sync::{
2    Arc,
3    atomic::{AtomicU32, Ordering},
4};
5use firewheel_core::{
6    StreamInfo,
7    atomic_float::AtomicF32,
8    channel_config::{ChannelConfig, ChannelCount},
9    diff::{Diff, Patch},
10    dsp::volume::amp_to_db,
11    event::ProcEvents,
12    node::{
13        AudioNode, AudioNodeInfo, AudioNodeProcessor, ConstructProcessorContext, EmptyConfig,
14        ProcBuffers, ProcExtra, ProcInfo, ProcStreamCtx, ProcessStatus,
15    },
16};
17
18use firewheel_core::node::NodeError;
19#[cfg(not(feature = "std"))]
20use num_traits::Float;
21
22/// A lightweight node that measures the loudness of a mono signal using a rough RMS
23/// (root mean square) estimate.
24///
25/// Note this node doesn't calculate the true RMS (That requires a much more expensive
26/// algorithm using a sliding window.) But it should be good enough for games that
27/// simply wish to react to player audio.
28#[derive(Debug, Diff, Patch, Clone, Copy, PartialEq)]
29#[cfg_attr(feature = "bevy", derive(bevy_ecs::prelude::Component))]
30#[cfg_attr(feature = "bevy_reflect", derive(bevy_reflect::Reflect))]
31#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
32pub struct FastRmsNode {
33    /// The size of the window used for measuring the RMS value.
34    ///
35    /// Smaller values are better at detecting short bursts of loudness (transients),
36    /// while larger values are better for measuring loudness on a broader time scale.
37    ///
38    /// By default this is set to `0.05` (50ms).
39    pub window_size_secs: f32,
40}
41
42impl Default for FastRmsNode {
43    fn default() -> Self {
44        Self {
45            window_size_secs: 50.0 / 1_000.0,
46        }
47    }
48}
49
50/// The state of a [`FastRmsNode`]. This contains the calculated RMS values.
51#[derive(Clone)]
52pub struct FastRmsState {
53    shared_state: Arc<SharedState>,
54}
55
56impl FastRmsState {
57    fn new() -> Self {
58        Self {
59            shared_state: Arc::new(SharedState {
60                rms_value: AtomicF32::new(0.0),
61                read_count: AtomicU32::new(1),
62            }),
63        }
64    }
65
66    /// Get the estimated RMS value in decibels.
67    ///
68    /// * `min_db` - If the RMS value is less than or equal to this value, then it
69    ///   will be clamped to `f32::NEG_INFINITY` (silence). (You can use
70    ///   [firewheel_core::dsp::volume::DEFAULT_MIN_DB].)
71    ///
72    /// If the node is currently disabled, then this will return a value
73    /// of `f32::NEG_INFINITY` (silence).
74    ///
75    /// Note this node doesn't calculate the true RMS (That requires a much more expensive
76    /// algorithm using a sliding window.) But it should be good enough for games that
77    /// simply wish to react to player audio.
78    pub fn rms_db(&self, min_db: f32) -> f32 {
79        let rms = amp_to_db(self.shared_state.rms_value.load(Ordering::Relaxed));
80        self.shared_state.read_count.fetch_add(1, Ordering::Release);
81
82        if rms <= min_db {
83            f32::NEG_INFINITY
84        } else {
85            rms
86        }
87    }
88}
89
90impl AudioNode for FastRmsNode {
91    type Configuration = EmptyConfig;
92
93    fn info(&self, _config: &Self::Configuration) -> Result<AudioNodeInfo, NodeError> {
94        Ok(AudioNodeInfo::new()
95            .debug_name("fast_rms")
96            .channel_config(ChannelConfig {
97                num_inputs: ChannelCount::MONO,
98                num_outputs: ChannelCount::ZERO,
99            })
100            .custom_state(FastRmsState::new()))
101    }
102
103    fn construct_processor(
104        &self,
105        _config: &Self::Configuration,
106        cx: ConstructProcessorContext,
107    ) -> Result<impl AudioNodeProcessor, NodeError> {
108        let window_frames =
109            (self.window_size_secs * cx.stream_info.sample_rate.get() as f32).round() as usize;
110
111        let custom_state = cx.custom_state::<FastRmsState>().unwrap();
112
113        Ok(Processor {
114            params: *self,
115            shared_state: Arc::clone(&custom_state.shared_state),
116            squares: 0.0,
117            num_squared_values: 0,
118            window_frames,
119            last_read_count: 0,
120        })
121    }
122}
123
124struct Processor {
125    params: FastRmsNode,
126    shared_state: Arc<SharedState>,
127    squares: f32,
128    num_squared_values: usize,
129    window_frames: usize,
130    last_read_count: u32,
131}
132
133impl AudioNodeProcessor for Processor {
134    fn events(&mut self, info: &ProcInfo, events: &mut ProcEvents, _extra: &mut ProcExtra) {
135        for patch in events.drain_patches::<FastRmsNode>() {
136            match patch {
137                FastRmsNodePatch::WindowSizeSecs(window_size_secs) => {
138                    let window_frames =
139                        (window_size_secs * info.sample_rate.get() as f32).round() as usize;
140
141                    if self.window_frames != window_frames {
142                        self.window_frames = window_frames;
143
144                        self.squares = 0.0;
145                        self.num_squared_values = 0;
146                    }
147                }
148            }
149
150            self.params.apply(patch);
151        }
152    }
153
154    fn bypassed(&mut self, _bypassed: bool) {
155        self.shared_state.rms_value.store(0.0, Ordering::Relaxed);
156
157        self.squares = 0.0;
158        self.num_squared_values = 0;
159    }
160
161    fn process(
162        &mut self,
163        info: &ProcInfo,
164        buffers: ProcBuffers,
165        _extra: &mut ProcExtra,
166    ) -> ProcessStatus {
167        let mut frames_processed = 0;
168        while frames_processed < info.frames {
169            let process_frames =
170                (info.frames - frames_processed).min(self.window_frames - self.num_squared_values);
171
172            if !info.in_silence_mask.is_channel_silent(0) {
173                for &s in
174                    buffers.inputs[0][frames_processed..frames_processed + process_frames].iter()
175                {
176                    self.squares += s * s;
177                }
178            }
179
180            self.num_squared_values += process_frames;
181            frames_processed += process_frames;
182
183            if self.num_squared_values == self.window_frames {
184                let mean = self.squares / self.window_frames as f32;
185                let rms = mean.sqrt();
186
187                let latest_read_count = self.shared_state.read_count.load(Ordering::Acquire);
188                let previous_rms = self.shared_state.rms_value.load(Ordering::Relaxed);
189
190                if latest_read_count != self.last_read_count || rms > previous_rms {
191                    self.shared_state.rms_value.store(rms, Ordering::Relaxed);
192                }
193
194                self.squares = 0.0;
195                self.num_squared_values = 0;
196                self.last_read_count = latest_read_count;
197            }
198        }
199
200        // There are no outputs in this node.
201        ProcessStatus::Bypass
202    }
203
204    fn new_stream(&mut self, stream_info: &StreamInfo, _context: &mut ProcStreamCtx) {
205        self.window_frames =
206            (self.params.window_size_secs * stream_info.sample_rate.get() as f32).round() as usize;
207
208        self.squares = 0.0;
209        self.num_squared_values = 0;
210    }
211}
212
213#[derive(Debug)]
214struct SharedState {
215    rms_value: AtomicF32,
216    // A simple counter used to keep track of when the processor should update
217    // the RMS value.
218    read_count: AtomicU32,
219}