ringkernel_audio_fft/
messages.rs

1//! Message types for audio FFT processing and K2K communication.
2//!
3//! This module defines all message types used for communication between
4//! the host and GPU bin actors, as well as between neighboring bin actors.
5
6use bytemuck::{Pod, Zeroable};
7use ringkernel_core::prelude::*;
8use rkyv::{Archive, Deserialize, Serialize};
9
10/// Message type IDs for audio FFT messages.
11pub mod message_types {
12    /// Audio frame input message.
13    pub const AUDIO_FRAME: u64 = 0x4155_4449_4F00_0001; // "AUDIO" + 1
14    /// Frequency bin data message.
15    pub const FREQUENCY_BIN: u64 = 0x4155_4449_4F00_0002;
16    /// Neighbor data exchange message.
17    pub const NEIGHBOR_DATA: u64 = 0x4155_4449_4F00_0003;
18    /// Separated bin output message.
19    pub const SEPARATED_BIN: u64 = 0x4155_4449_4F00_0004;
20    /// Processing complete signal.
21    pub const FRAME_COMPLETE: u64 = 0x4155_4449_4F00_0005;
22    /// Bin actor control message.
23    pub const BIN_CONTROL: u64 = 0x4155_4449_4F00_0006;
24}
25
26/// A frame of audio samples for FFT processing.
27#[derive(Debug, Clone, Archive, Serialize, Deserialize)]
28pub struct AudioFrame {
29    /// Unique frame identifier.
30    pub frame_id: u64,
31    /// Sample rate in Hz.
32    pub sample_rate: u32,
33    /// Number of channels (1 = mono, 2 = stereo).
34    pub channels: u8,
35    /// Audio samples (interleaved if stereo).
36    pub samples: Vec<f32>,
37    /// Timestamp in samples from start.
38    pub timestamp_samples: u64,
39}
40
41impl AudioFrame {
42    /// Create a new audio frame.
43    pub fn new(
44        frame_id: u64,
45        sample_rate: u32,
46        channels: u8,
47        samples: Vec<f32>,
48        timestamp_samples: u64,
49    ) -> Self {
50        Self {
51            frame_id,
52            sample_rate,
53            channels,
54            samples,
55            timestamp_samples,
56        }
57    }
58
59    /// Get the duration of this frame in seconds.
60    pub fn duration_secs(&self) -> f64 {
61        let sample_count = self.samples.len() / self.channels as usize;
62        sample_count as f64 / self.sample_rate as f64
63    }
64
65    /// Get samples for a specific channel.
66    pub fn channel_samples(&self, channel: usize) -> Vec<f32> {
67        if channel >= self.channels as usize {
68            return Vec::new();
69        }
70        self.samples
71            .iter()
72            .skip(channel)
73            .step_by(self.channels as usize)
74            .copied()
75            .collect()
76    }
77}
78
79impl RingMessage for AudioFrame {
80    fn message_type() -> u64 {
81        message_types::AUDIO_FRAME
82    }
83
84    fn message_id(&self) -> MessageId {
85        MessageId::new(self.frame_id)
86    }
87
88    fn serialize(&self) -> Vec<u8> {
89        rkyv::to_bytes::<_, 256>(self)
90            .map(|v| v.to_vec())
91            .unwrap_or_default()
92    }
93
94    fn deserialize(bytes: &[u8]) -> ringkernel_core::error::Result<Self> {
95        // SAFETY: We trust the serialized data from our own FFT processing
96        let archived = unsafe { rkyv::archived_root::<Self>(bytes) };
97        Ok(archived.deserialize(&mut rkyv::Infallible).unwrap())
98    }
99}
100
101/// Complex number representation for FFT bins.
102#[derive(Debug, Clone, Copy, Default, Pod, Zeroable, Archive, Serialize, Deserialize)]
103#[repr(C)]
104pub struct Complex {
105    /// Real component.
106    pub re: f32,
107    /// Imaginary component.
108    pub im: f32,
109}
110
111impl Complex {
112    /// Create a new complex number.
113    pub const fn new(re: f32, im: f32) -> Self {
114        Self { re, im }
115    }
116
117    /// Create from polar coordinates.
118    pub fn from_polar(magnitude: f32, phase: f32) -> Self {
119        Self {
120            re: magnitude * phase.cos(),
121            im: magnitude * phase.sin(),
122        }
123    }
124
125    /// Get magnitude.
126    pub fn magnitude(&self) -> f32 {
127        (self.re * self.re + self.im * self.im).sqrt()
128    }
129
130    /// Get phase in radians.
131    pub fn phase(&self) -> f32 {
132        self.im.atan2(self.re)
133    }
134
135    /// Get magnitude squared (more efficient than magnitude).
136    pub fn magnitude_squared(&self) -> f32 {
137        self.re * self.re + self.im * self.im
138    }
139
140    /// Multiply by another complex number.
141    pub fn mul(&self, other: &Self) -> Self {
142        Self {
143            re: self.re * other.re - self.im * other.im,
144            im: self.re * other.im + self.im * other.re,
145        }
146    }
147
148    /// Complex conjugate.
149    pub fn conj(&self) -> Self {
150        Self {
151            re: self.re,
152            im: -self.im,
153        }
154    }
155
156    /// Add another complex number.
157    pub fn add(&self, other: &Self) -> Self {
158        Self {
159            re: self.re + other.re,
160            im: self.im + other.im,
161        }
162    }
163
164    /// Scale by a real number.
165    pub fn scale(&self, s: f32) -> Self {
166        Self {
167            re: self.re * s,
168            im: self.im * s,
169        }
170    }
171}
172
173/// Data for a single frequency bin.
174#[derive(Debug, Clone, Archive, Serialize, Deserialize)]
175pub struct FrequencyBin {
176    /// Frame this bin belongs to.
177    pub frame_id: u64,
178    /// Bin index (0 = DC, fft_size/2 = Nyquist).
179    pub bin_index: u32,
180    /// Total number of bins.
181    pub total_bins: u32,
182    /// Complex FFT value.
183    pub value: Complex,
184    /// Frequency in Hz for this bin.
185    pub frequency_hz: f32,
186    /// Previous frame's value (for temporal analysis).
187    pub prev_value: Option<Complex>,
188}
189
190impl FrequencyBin {
191    /// Create a new frequency bin.
192    pub fn new(
193        frame_id: u64,
194        bin_index: u32,
195        total_bins: u32,
196        value: Complex,
197        frequency_hz: f32,
198    ) -> Self {
199        Self {
200            frame_id,
201            bin_index,
202            total_bins,
203            value,
204            frequency_hz,
205            prev_value: None,
206        }
207    }
208
209    /// Set the previous frame's value.
210    pub fn with_prev_value(mut self, prev: Complex) -> Self {
211        self.prev_value = Some(prev);
212        self
213    }
214
215    /// Get magnitude in dB.
216    pub fn magnitude_db(&self) -> f32 {
217        let mag = self.value.magnitude();
218        if mag > 1e-10 {
219            20.0 * mag.log10()
220        } else {
221            -200.0 // Floor
222        }
223    }
224}
225
226impl RingMessage for FrequencyBin {
227    fn message_type() -> u64 {
228        message_types::FREQUENCY_BIN
229    }
230
231    fn message_id(&self) -> MessageId {
232        MessageId::new(self.frame_id * 10000 + self.bin_index as u64)
233    }
234
235    fn serialize(&self) -> Vec<u8> {
236        rkyv::to_bytes::<_, 128>(self)
237            .map(|v| v.to_vec())
238            .unwrap_or_default()
239    }
240
241    fn deserialize(bytes: &[u8]) -> ringkernel_core::error::Result<Self> {
242        // SAFETY: We trust the serialized data from our own FFT processing
243        let archived = unsafe { rkyv::archived_root::<Self>(bytes) };
244        Ok(archived.deserialize(&mut rkyv::Infallible).unwrap())
245    }
246}
247
248/// Neighbor data exchanged via K2K messaging.
249#[derive(Debug, Clone, Archive, Serialize, Deserialize)]
250pub struct NeighborData {
251    /// Source bin index.
252    pub source_bin: u32,
253    /// Frame ID.
254    pub frame_id: u64,
255    /// Complex value.
256    pub value: Complex,
257    /// Magnitude.
258    pub magnitude: f32,
259    /// Phase.
260    pub phase: f32,
261    /// Temporal derivative (phase change from previous frame).
262    pub phase_derivative: f32,
263    /// Spectral flux (magnitude change).
264    pub spectral_flux: f32,
265}
266
267impl NeighborData {
268    /// Create neighbor data from a frequency bin.
269    pub fn from_bin(bin: &FrequencyBin, phase_derivative: f32, spectral_flux: f32) -> Self {
270        Self {
271            source_bin: bin.bin_index,
272            frame_id: bin.frame_id,
273            value: bin.value,
274            magnitude: bin.value.magnitude(),
275            phase: bin.value.phase(),
276            phase_derivative,
277            spectral_flux,
278        }
279    }
280}
281
282impl RingMessage for NeighborData {
283    fn message_type() -> u64 {
284        message_types::NEIGHBOR_DATA
285    }
286
287    fn message_id(&self) -> MessageId {
288        MessageId::new(self.frame_id * 10000 + self.source_bin as u64)
289    }
290
291    fn serialize(&self) -> Vec<u8> {
292        rkyv::to_bytes::<_, 64>(self)
293            .map(|v| v.to_vec())
294            .unwrap_or_default()
295    }
296
297    fn deserialize(bytes: &[u8]) -> ringkernel_core::error::Result<Self> {
298        // SAFETY: We trust the serialized data from our own FFT processing
299        let archived = unsafe { rkyv::archived_root::<Self>(bytes) };
300        Ok(archived.deserialize(&mut rkyv::Infallible).unwrap())
301    }
302}
303
304/// Separated frequency bin with direct and ambience components.
305#[derive(Debug, Clone, Archive, Serialize, Deserialize)]
306pub struct SeparatedBin {
307    /// Frame ID.
308    pub frame_id: u64,
309    /// Bin index.
310    pub bin_index: u32,
311    /// Direct signal component.
312    pub direct: Complex,
313    /// Ambience component.
314    pub ambience: Complex,
315    /// Coherence score (0.0 = pure ambience, 1.0 = pure direct).
316    pub coherence: f32,
317    /// Transient score (0.0 = steady, 1.0 = transient).
318    pub transient: f32,
319}
320
321impl SeparatedBin {
322    /// Create a new separated bin.
323    pub fn new(
324        frame_id: u64,
325        bin_index: u32,
326        direct: Complex,
327        ambience: Complex,
328        coherence: f32,
329        transient: f32,
330    ) -> Self {
331        Self {
332            frame_id,
333            bin_index,
334            direct,
335            ambience,
336            coherence,
337            transient,
338        }
339    }
340
341    /// Get the original combined value.
342    pub fn combined(&self) -> Complex {
343        self.direct.add(&self.ambience)
344    }
345
346    /// Get direct signal magnitude ratio.
347    pub fn direct_ratio(&self) -> f32 {
348        let total = self.direct.magnitude() + self.ambience.magnitude();
349        if total > 1e-10 {
350            self.direct.magnitude() / total
351        } else {
352            0.5
353        }
354    }
355}
356
357impl RingMessage for SeparatedBin {
358    fn message_type() -> u64 {
359        message_types::SEPARATED_BIN
360    }
361
362    fn message_id(&self) -> MessageId {
363        MessageId::new(self.frame_id * 10000 + self.bin_index as u64)
364    }
365
366    fn serialize(&self) -> Vec<u8> {
367        rkyv::to_bytes::<_, 128>(self)
368            .map(|v| v.to_vec())
369            .unwrap_or_default()
370    }
371
372    fn deserialize(bytes: &[u8]) -> ringkernel_core::error::Result<Self> {
373        // SAFETY: We trust the serialized data from our own FFT processing
374        let archived = unsafe { rkyv::archived_root::<Self>(bytes) };
375        Ok(archived.deserialize(&mut rkyv::Infallible).unwrap())
376    }
377}
378
379/// Control message for bin actors.
380#[derive(Debug, Clone, Archive, Serialize, Deserialize)]
381pub enum BinControl {
382    /// Reset state for new audio stream.
383    Reset,
384    /// Update separation parameters.
385    UpdateParams {
386        /// Coherence threshold.
387        coherence_threshold: f32,
388        /// Transient sensitivity.
389        transient_sensitivity: f32,
390        /// Temporal smoothing factor.
391        temporal_smoothing: f32,
392    },
393    /// Request current state.
394    GetState,
395    /// Shutdown the actor.
396    Shutdown,
397}
398
399impl RingMessage for BinControl {
400    fn message_type() -> u64 {
401        message_types::BIN_CONTROL
402    }
403
404    fn message_id(&self) -> MessageId {
405        MessageId::generate()
406    }
407
408    fn serialize(&self) -> Vec<u8> {
409        rkyv::to_bytes::<_, 64>(self)
410            .map(|v| v.to_vec())
411            .unwrap_or_default()
412    }
413
414    fn deserialize(bytes: &[u8]) -> ringkernel_core::error::Result<Self> {
415        // SAFETY: We trust the serialized data from our own FFT processing
416        let archived = unsafe { rkyv::archived_root::<Self>(bytes) };
417        Ok(archived.deserialize(&mut rkyv::Infallible).unwrap())
418    }
419}
420
421/// Frame processing complete notification.
422#[derive(Debug, Clone, Archive, Serialize, Deserialize)]
423pub struct FrameComplete {
424    /// Frame ID.
425    pub frame_id: u64,
426    /// Processing time in microseconds.
427    pub processing_time_us: u64,
428    /// Number of bins processed.
429    pub bins_processed: u32,
430}
431
432impl RingMessage for FrameComplete {
433    fn message_type() -> u64 {
434        message_types::FRAME_COMPLETE
435    }
436
437    fn message_id(&self) -> MessageId {
438        MessageId::new(self.frame_id)
439    }
440
441    fn serialize(&self) -> Vec<u8> {
442        rkyv::to_bytes::<_, 32>(self)
443            .map(|v| v.to_vec())
444            .unwrap_or_default()
445    }
446
447    fn deserialize(bytes: &[u8]) -> ringkernel_core::error::Result<Self> {
448        // SAFETY: We trust the serialized data from our own FFT processing
449        let archived = unsafe { rkyv::archived_root::<Self>(bytes) };
450        Ok(archived.deserialize(&mut rkyv::Infallible).unwrap())
451    }
452}
453
454#[cfg(test)]
455mod tests {
456    use super::*;
457
458    #[test]
459    fn test_complex_operations() {
460        let c1 = Complex::new(3.0, 4.0);
461        assert!((c1.magnitude() - 5.0).abs() < 1e-6);
462
463        let c2 = Complex::from_polar(5.0, 0.927295); // approx atan2(4, 3)
464        assert!((c2.re - 3.0).abs() < 0.01);
465        assert!((c2.im - 4.0).abs() < 0.01);
466    }
467
468    #[test]
469    fn test_audio_frame_serialization() {
470        use ringkernel_core::RingMessage;
471        let frame = AudioFrame::new(1, 44100, 2, vec![0.0, 0.1, 0.2, 0.3], 0);
472        let bytes = RingMessage::serialize(&frame);
473        let restored = <AudioFrame as RingMessage>::deserialize(&bytes).unwrap();
474        assert_eq!(restored.frame_id, 1);
475        assert_eq!(restored.sample_rate, 44100);
476        assert_eq!(restored.samples.len(), 4);
477    }
478
479    #[test]
480    fn test_frequency_bin_db() {
481        let bin = FrequencyBin::new(0, 10, 1024, Complex::new(1.0, 0.0), 440.0);
482        assert!((bin.magnitude_db() - 0.0).abs() < 0.1); // 1.0 magnitude = 0 dB
483
484        let quiet_bin = FrequencyBin::new(0, 10, 1024, Complex::new(0.1, 0.0), 440.0);
485        assert!((quiet_bin.magnitude_db() - (-20.0)).abs() < 0.1); // 0.1 magnitude = -20 dB
486    }
487}