Skip to main content

adk_audio/
mixer.rs

1//! Multi-track audio mixer with per-track volume control.
2
3use std::collections::HashMap;
4
5use async_trait::async_trait;
6use bytes::Bytes;
7
8use crate::error::{AudioError, AudioResult};
9use crate::frame::AudioFrame;
10use crate::traits::AudioProcessor;
11
12/// A track in the mixer with volume and buffered audio.
13struct MixerTrack {
14    volume: f32,
15    buffer: Option<AudioFrame>,
16}
17
18/// Multi-track audio mixer.
19///
20/// Combines multiple named audio tracks into a single output with
21/// per-track volume control. Missing tracks are treated as silence.
22///
23/// # Example
24///
25/// ```ignore
26/// let mut mixer = Mixer::new(24000);
27/// mixer.add_track("narration", 1.0);
28/// mixer.add_track("music", 0.3);
29/// mixer.push_frame("narration", narration_frame);
30/// mixer.push_frame("music", music_frame);
31/// let mixed = mixer.mix()?;
32/// ```
33pub struct Mixer {
34    tracks: HashMap<String, MixerTrack>,
35    output_sample_rate: u32,
36}
37
38impl Mixer {
39    /// Create a new mixer with the given output sample rate.
40    pub fn new(output_sample_rate: u32) -> Self {
41        Self { tracks: HashMap::new(), output_sample_rate }
42    }
43
44    /// Add a named track with the given volume (0.0–1.0).
45    pub fn add_track(&mut self, name: impl Into<String>, volume: f32) {
46        self.tracks
47            .insert(name.into(), MixerTrack { volume: volume.clamp(0.0, 1.0), buffer: None });
48    }
49
50    /// Set the volume for a named track.
51    pub fn set_volume(&mut self, name: &str, volume: f32) {
52        if let Some(track) = self.tracks.get_mut(name) {
53            track.volume = volume.clamp(0.0, 1.0);
54        }
55    }
56
57    /// Push an audio frame to a named track.
58    pub fn push_frame(&mut self, track: &str, frame: AudioFrame) {
59        if let Some(t) = self.tracks.get_mut(track) {
60            t.buffer = Some(frame);
61        }
62    }
63
64    /// Mix all tracks into a single output frame.
65    ///
66    /// Tracks without buffered audio are treated as silence.
67    /// All tracks are mixed at the output sample rate.
68    pub fn mix(&mut self) -> AudioResult<AudioFrame> {
69        if self.tracks.is_empty() {
70            return Err(AudioError::Fx("mixer has no tracks".into()));
71        }
72
73        // Find the maximum sample count across all buffered tracks
74        let max_samples = self
75            .tracks
76            .values()
77            .filter_map(|t| t.buffer.as_ref())
78            .map(|f| f.data.len() / 2)
79            .max()
80            .unwrap_or(0);
81
82        if max_samples == 0 {
83            return Ok(AudioFrame::silence(self.output_sample_rate, 1, 0));
84        }
85
86        let mut mixed = vec![0i32; max_samples];
87
88        for track in self.tracks.values() {
89            let volume = track.volume;
90            if let Some(ref frame) = track.buffer {
91                let samples = frame.samples();
92                for (i, &s) in samples.iter().enumerate() {
93                    if i < max_samples {
94                        mixed[i] += (s as f32 * volume) as i32;
95                    }
96                }
97            }
98        }
99
100        // Clamp to i16 range
101        let pcm: Vec<u8> = mixed
102            .iter()
103            .flat_map(|&s| {
104                let clamped = s.clamp(-32768, 32767) as i16;
105                clamped.to_le_bytes()
106            })
107            .collect();
108
109        // Clear buffers
110        for track in self.tracks.values_mut() {
111            track.buffer = None;
112        }
113
114        Ok(AudioFrame::new(Bytes::from(pcm), self.output_sample_rate, 1))
115    }
116}
117
118#[async_trait]
119impl AudioProcessor for Mixer {
120    async fn process(&self, frame: &AudioFrame) -> AudioResult<AudioFrame> {
121        // Single-track passthrough: apply first track's volume
122        let volume = self.tracks.values().next().map(|t| t.volume).unwrap_or(1.0);
123
124        let samples = frame.samples();
125        let pcm: Vec<u8> = samples
126            .iter()
127            .flat_map(|&s| {
128                let scaled = (s as f32 * volume) as i32;
129                let clamped = scaled.clamp(-32768, 32767) as i16;
130                clamped.to_le_bytes()
131            })
132            .collect();
133
134        Ok(AudioFrame::new(Bytes::from(pcm), frame.sample_rate, frame.channels))
135    }
136}