use std::sync::atomic::{AtomicU32, AtomicUsize, Ordering};
use std::sync::Arc;
use crate::alloc::AudioBuffer as FixedAudioBuffer;
use crate::media::MediaStream;
use crate::SampleRate;
#[derive(Clone, Debug)]
pub struct AudioBuffer {
channels: Vec<ChannelData>,
sample_rate: SampleRate,
}
use std::error::Error;
impl AudioBuffer {
pub fn new(channels: usize, len: usize, sample_rate: SampleRate) -> Self {
let silence = ChannelData::new(len);
Self {
channels: vec![silence; channels],
sample_rate,
}
}
pub fn from_channels(channels: Vec<ChannelData>, sample_rate: SampleRate) -> Self {
Self {
channels,
sample_rate,
}
}
pub fn number_of_channels(&self) -> usize {
self.channels.len()
}
pub fn sample_len(&self) -> usize {
self.channels.get(0).map(ChannelData::len).unwrap_or(0)
}
pub fn sample_rate(&self) -> SampleRate {
self.sample_rate
}
pub fn duration(&self) -> f64 {
self.sample_len() as f64 / self.sample_rate.0 as f64
}
pub fn channels(&self) -> &[ChannelData] {
&self.channels
}
pub fn channels_mut(&mut self) -> &mut [ChannelData] {
&mut self.channels
}
pub fn channel_data(&self, index: usize) -> &ChannelData {
&self.channels[index]
}
pub fn channel_data_mut(&mut self, index: usize) -> &mut ChannelData {
&mut self.channels[index]
}
pub fn modify_channels<F: Fn(&mut ChannelData)>(&mut self, fun: F) {
self.channels.iter_mut().for_each(fun)
}
pub fn extend(&mut self, other: &Self) {
assert_eq!(self.sample_rate, other.sample_rate);
assert_eq!(self.number_of_channels(), other.number_of_channels());
let data = self.channels_mut();
data.iter_mut()
.zip(other.channels.iter())
.for_each(|(channel, other_channel)| {
let cur_channel_data = Arc::make_mut(&mut channel.data);
cur_channel_data.extend(other_channel.as_slice());
})
}
pub fn extend_alloc(&mut self, other: &FixedAudioBuffer) {
self.channels_mut()
.iter_mut()
.zip(other.channels())
.for_each(|(channel, other_channel)| {
let cur_channel_data = Arc::make_mut(&mut channel.data);
cur_channel_data.extend_from_slice(&other_channel[..]);
})
}
pub fn split(mut self, sample_len: u32) -> Vec<AudioBuffer> {
let sample_len = sample_len as usize;
let total_len = self.sample_len();
let sample_rate = self.sample_rate();
let mut channels: Vec<_> = self
.channels_mut()
.iter()
.map(|channel_data| channel_data.as_slice().chunks(sample_len))
.collect();
(0..total_len)
.step_by(sample_len)
.map(|_| {
let cur: Vec<_> = channels
.iter_mut()
.map(|c| ChannelData::from(c.next().unwrap().to_vec()))
.collect();
AudioBuffer::from_channels(cur, sample_rate)
})
.collect()
}
pub fn split_off(&mut self, index: u32) -> AudioBuffer {
let index = index as usize;
let sample_rate = self.sample_rate();
let channels: Vec<_> = self
.channels_mut()
.iter_mut()
.map(|channel_data| Arc::make_mut(&mut channel_data.data).split_off(index))
.map(ChannelData::from)
.collect();
AudioBuffer::from_channels(channels, sample_rate)
}
pub fn resample(&mut self, sample_rate: SampleRate) {
if self.sample_rate() == sample_rate {
return;
}
let rate = sample_rate.0 as f32 / self.sample_rate.0 as f32;
self.modify_channels(|channel_data| {
let mut current = 0;
let resampled = channel_data
.data
.iter()
.enumerate()
.flat_map(|(i, v)| {
let target = ((i + 1) as f32 * rate) as usize;
let take = target - current.min(target);
current += take;
std::iter::repeat(*v).take(take)
})
.collect();
channel_data.data = Arc::new(resampled);
});
self.sample_rate = sample_rate;
}
}
#[derive(Clone, Debug, PartialEq)]
pub struct ChannelData {
data: Arc<Vec<f32>>,
}
impl ChannelData {
pub fn new(length: usize) -> Self {
let buffer = vec![0.; length];
let data = Arc::new(buffer);
Self { data }
}
pub fn from(data: Vec<f32>) -> Self {
Self {
data: Arc::new(data),
}
}
pub fn len(&self) -> usize {
self.data.len()
}
pub fn is_empty(&self) -> bool {
self.data.is_empty()
}
pub fn as_slice(&self) -> &[f32] {
&self.data[..]
}
pub fn as_mut_slice(&mut self) -> &mut [f32] {
&mut Arc::make_mut(&mut self.data)[..]
}
}
#[derive(Copy, Clone, PartialEq, Eq, Debug)]
pub enum ChannelCountMode {
Max,
ClampedMax,
Explicit,
}
impl From<u32> for ChannelCountMode {
fn from(i: u32) -> Self {
use ChannelCountMode::*;
match i {
0 => Max,
1 => ClampedMax,
2 => Explicit,
_ => unreachable!(),
}
}
}
#[derive(Copy, Clone, PartialEq, Eq, Debug)]
pub enum ChannelInterpretation {
Speakers,
Discrete,
}
impl From<u32> for ChannelInterpretation {
fn from(i: u32) -> Self {
use ChannelInterpretation::*;
match i {
0 => Speakers,
1 => Discrete,
_ => unreachable!(),
}
}
}
#[derive(Clone, Debug)]
pub struct ChannelConfigOptions {
pub count: usize,
pub mode: ChannelCountMode,
pub interpretation: ChannelInterpretation,
}
impl Default for ChannelConfigOptions {
fn default() -> Self {
Self {
count: 2,
mode: ChannelCountMode::Max,
interpretation: ChannelInterpretation::Speakers,
}
}
}
#[derive(Clone, Debug)]
pub struct ChannelConfig {
count: Arc<AtomicUsize>,
mode: Arc<AtomicU32>,
interpretation: Arc<AtomicU32>,
}
impl ChannelConfig {
pub fn count_mode(&self) -> ChannelCountMode {
self.mode.load(Ordering::SeqCst).into()
}
pub fn set_count_mode(&self, v: ChannelCountMode) {
self.mode.store(v as u32, Ordering::SeqCst)
}
pub fn interpretation(&self) -> ChannelInterpretation {
self.interpretation.load(Ordering::SeqCst).into()
}
pub fn set_interpretation(&self, v: ChannelInterpretation) {
self.interpretation.store(v as u32, Ordering::SeqCst)
}
pub fn count(&self) -> usize {
self.count.load(Ordering::SeqCst)
}
pub fn set_count(&self, v: usize) {
self.count.store(v, Ordering::SeqCst)
}
}
impl From<ChannelConfigOptions> for ChannelConfig {
fn from(opts: ChannelConfigOptions) -> Self {
ChannelConfig {
count: Arc::new(AtomicUsize::from(opts.count)),
mode: Arc::new(AtomicU32::from(opts.mode as u32)),
interpretation: Arc::new(AtomicU32::from(opts.interpretation as u32)),
}
}
}
impl std::iter::FromIterator<AudioBuffer> for AudioBuffer {
fn from_iter<I: IntoIterator<Item = AudioBuffer>>(iter: I) -> Self {
let mut iter = iter.into_iter();
let mut collect: AudioBuffer = match iter.next() {
None => return AudioBuffer::new(0, 0, SampleRate(0)),
Some(first) => first,
};
for elem in iter {
collect.extend(&elem);
}
collect
}
}
pub struct Resampler<I> {
sample_rate: SampleRate,
sample_len: u32,
input: I,
buffer: Option<AudioBuffer>,
}
impl<M: MediaStream> Resampler<M> {
pub fn new(sample_rate: SampleRate, sample_len: u32, input: M) -> Self {
Self {
sample_rate,
sample_len,
input,
buffer: None,
}
}
}
impl<M: MediaStream> Iterator for Resampler<M> {
type Item = Result<AudioBuffer, Box<dyn Error + Send>>;
fn next(&mut self) -> Option<Self::Item> {
let mut buffer = match self.buffer.take() {
None => match self.input.next() {
None => return None,
Some(Err(e)) => return Some(Err(e)),
Some(Ok(mut data)) => {
data.resample(self.sample_rate);
data
}
},
Some(data) => data,
};
while (buffer.sample_len() as u32) < self.sample_len {
match self.input.next() {
None => {
let padding = AudioBuffer::new(
buffer.number_of_channels(),
self.sample_len as usize - buffer.sample_len(),
self.sample_rate,
);
buffer.extend(&padding);
return Some(Ok(buffer));
}
Some(Err(e)) => return Some(Err(e)),
Some(Ok(mut data)) => {
data.resample(self.sample_rate);
buffer.extend(&data)
}
}
}
if buffer.sample_len() as u32 == self.sample_len {
return Some(Ok(buffer));
}
self.buffer = Some(buffer.split_off(self.sample_len));
Some(Ok(buffer))
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_silent() {
let b = AudioBuffer::new(2, 10, SampleRate(44_100));
assert_eq!(b.sample_len(), 10);
assert_eq!(b.number_of_channels(), 2);
assert_eq!(b.sample_rate().0, 44_100);
assert_eq!(b.channel_data(0).as_slice(), &[0.; 10]);
assert_eq!(b.channel_data(1).as_slice(), &[0.; 10]);
assert_eq!(b.channels().get(2), None);
}
#[test]
fn test_concat_split() {
let mut b1 = AudioBuffer::new(2, 5, SampleRate(44_100));
let b2 = AudioBuffer::new(2, 5, SampleRate(44_100));
b1.extend(&b2);
assert_eq!(b1.sample_len(), 10);
assert_eq!(b1.number_of_channels(), 2);
assert_eq!(b1.sample_rate().0, 44_100);
let channel_data = ChannelData::from(vec![1.; 5]);
let b3 = AudioBuffer::from_channels(vec![channel_data; 2], SampleRate(44_100));
b1.extend(&b3);
assert_eq!(b1.sample_len(), 15);
assert_eq!(b1.number_of_channels(), 2);
assert_eq!(b1.sample_rate().0, 44_100);
assert_eq!(
b1.channel_data(0).as_slice(),
&[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 1., 1., 1., 1.]
);
let split = b1.split(8);
assert_eq!(
split[0].channel_data(0).as_slice(),
&[0., 0., 0., 0., 0., 0., 0., 0.]
);
assert_eq!(
split[1].channel_data(0).as_slice(),
&[0., 0., 1., 1., 1., 1., 1.]
);
}
#[test]
fn test_resample_upmix() {
let channel = ChannelData::from(vec![1., 2., 3., 4., 5.]);
let mut buffer = AudioBuffer::from_channels(vec![channel], SampleRate(100));
buffer.resample(SampleRate(200));
assert_eq!(
buffer.channel_data(0).as_slice(),
&[1., 1., 2., 2., 3., 3., 4., 4., 5., 5.,]
);
assert_eq!(buffer.sample_rate().0, 200);
}
#[test]
fn test_resample_downmix() {
let channel = ChannelData::from(vec![1., 2., 3., 4., 5.]);
let mut buffer = AudioBuffer::from_channels(vec![channel], SampleRate(200));
buffer.resample(SampleRate(100));
assert_eq!(buffer.channel_data(0).as_slice(), &[2., 4.]);
assert_eq!(buffer.sample_rate().0, 100);
}
#[test]
fn test_resampler_concat() {
let channel = ChannelData::from(vec![1., 2., 3., 4., 5.]);
let input_buf = AudioBuffer::from_channels(vec![channel], SampleRate(44_100));
let input = vec![input_buf; 3].into_iter().map(|b| Ok(b));
let mut resampler = Resampler::new(SampleRate(44_100), 10, input);
let next = resampler.next().unwrap().unwrap();
assert_eq!(next.sample_len(), 10);
assert_eq!(
next.channel_data(0).as_slice(),
&[1., 2., 3., 4., 5., 1., 2., 3., 4., 5.,]
);
let next = resampler.next().unwrap().unwrap();
assert_eq!(next.sample_len(), 10);
assert_eq!(
next.channel_data(0).as_slice(),
&[1., 2., 3., 4., 5., 0., 0., 0., 0., 0.]
);
assert!(resampler.next().is_none());
}
#[test]
fn test_resampler_split() {
let channel = ChannelData::from(vec![1., 2., 3., 4., 5., 6., 7., 8., 9., 10.]);
let input_buf = Ok(AudioBuffer::from_channels(
vec![channel],
SampleRate(44_100),
));
let input = vec![input_buf].into_iter();
let mut resampler = Resampler::new(SampleRate(44_100), 5, input);
let next = resampler.next().unwrap().unwrap();
assert_eq!(next.sample_len(), 5);
assert_eq!(next.channel_data(0).as_slice(), &[1., 2., 3., 4., 5.,]);
let next = resampler.next().unwrap().unwrap();
assert_eq!(next.sample_len(), 5);
assert_eq!(next.channel_data(0).as_slice(), &[6., 7., 8., 9., 10.]);
assert!(resampler.next().is_none());
}
}