firewheel_nodes/
fast_rms.rs1use 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#[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 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#[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 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 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 read_count: AtomicU32,
219}