lib.rs

/*
 * Copyright (c) 2020 the Wavr Audio project.
 * This source file, as well as the binaries generated by it,
 * are licensed under MIT.
 */
//! # The Wavr Audio Buffer
//!
//! This crate implements an audio buffer that supports converting from and to interleaved data. It allows iterating over channels and single samples for each channels, enabling a declarative style of implementing audio processing.
//!
//! ## Usage
//!
//! ### Creating a zeroed buffer
//!
//! ```rust
//! const CHANNELS: usize = 2;
//! const BUFFER_SIZE: usize = 512;
//!
//! fn main() {
//!     let mut buffer = AudioBuffer::zeroed(CHANNELS, BUFFER_SIZE);
//!     generate_audio(&mut buffer);
//!     let data = buffer.interleave();
//!     assert_eq!(CHANNELS * BUFFER_SIZE, data.len());
//! }
//! ```
//!
//! ### Converting from interleaved data
//!
//! ```rust
//! const CHANNELS: usize = 2;
//!
//! fn process_interleaved(data: &mut [f64]) {
//!     let mut buffer = AudioBuffer::new(CHANNELS, data);
//!     process_audio(&mut buffer);
//!     buffer.copy_into_interleaved(data);
//! }
//! ```
//!
//! ### Getting per-channel RMS values
//!
//! ```rust
//! fn get_rms(buffer: &AudioBuffer) -> Vec<f64> {
//!     buffer
//!         .iter()
//!         .map(|ch|
//!             ch.fold(0.0, |acc, v|
//!                 acc + v
//!                     .powi(2))
//!                     .sqrt()
//!         )
//!         .collect()
//! }
//! ```

use std::borrow::{Borrow, BorrowMut};
use std::cmp::Ordering;
use std::ops::{Bound, Deref, Index, IndexMut, Range, RangeBounds};

use smallvec::{Array, SmallVec};

/// Structure holding per-channel buffers of audio data.
#[derive(Clone, Debug, PartialEq)]
pub struct AudioBuffer {
    audio_data: SmallVec<[Vec<f64>; 16]>,
    buffer_size: usize,
}

/// `AudioBuffer` iterator over channels data.
#[derive(Clone, Debug, PartialEq)]
pub struct Iter<'a> {
    buffer: &'a AudioBuffer,
    position: usize,
}

impl AudioBuffer {
    /// Create an `AudioBuffer` from interleaved data. Samples are copied from
    /// the buffer into itself.
    pub fn new(channels: usize, data: &[f64]) -> Self {
        let mut this = unsafe { Self::uninitialized(channels, data.len() / channels) };

        for (i, v) in data.iter().cloned().enumerate() {
            this.audio_data[i % channels][i / channels] = v;
        }

        this
    }

    /// Create an zeroed `AudioBuffer`.
    pub fn zeroed(channels: usize, buffer_size: usize) -> Self {
        Self {
            audio_data: SmallVec::from_vec(vec![vec![0.0; buffer_size]; channels]),
            buffer_size,
        }
    }

    /// Create an `AudioBuffer` containing unitialized data. This is faster than
    /// `AudioBuffer::zeroed` but reading from it is undefined behavior.
    ///
    /// ### Safety
    ///
    /// This function uses `Vec::set_len` to fill the inner buffer without
    /// explicitely zeroing it. The buffers will therefore contain uninitilized
    /// data and reading from it is undefined behavior (and playing it will hurt
    /// your ears!)
    pub unsafe fn uninitialized(channels: usize, buffer_size: usize) -> Self {
        Self {
            audio_data: SmallVec::from_vec(vec![
                {
                    let mut v = Vec::with_capacity(buffer_size);
                    v.set_len(buffer_size);
                    v
                };
                channels
            ]),
            buffer_size,
        }
    }

    /// Gets the number of channels in the buffer.
    #[inline]
    pub fn channels(&self) -> usize {
        self.audio_data.len()
    }

    /// Gets the sample size of the buffer.
    #[inline]
    pub fn buffer_size(&self) -> usize {
        self.buffer_size
    }

    /// Return a copy of the buffer with `channels` number of channels. It will
    /// copy the first N channels if the requested number is small than what the
    /// buffer holds, otherwise it will create zeroed channels to match the
    /// requested amount.
    pub fn with_channels(mut self, channels: usize) -> Self {
        match self.channels().cmp(&channels) {
            Ordering::Less => {
                for _ in 0..(channels - self.channels()) {
                    self.audio_data.push(vec![0.0; self.buffer_size]);
                }
                self
            }
            Ordering::Greater => {
                for _ in 0..(self.channels() - channels) {
                    self.audio_data.pop();
                }
                self
            }
            Ordering::Equal => self,
        }
    }

    /// Return a reference to the nth channel of the buffer, or `None` if the channel is not
    /// available.
    pub fn channel(&self, channel: usize) -> Option<&[f64]> {
        if channel < self.audio_data.len() {
            Some(self.audio_data[channel].borrow())
        } else {
            None
        }
    }

    /// Return a mutable reference to the nth channel of the buffer, of `None`
    /// if the channel is not available.
    pub fn channel_mut(&mut self, channel: usize) -> Option<&mut [f64]> {
        if channel < self.audio_data.len() {
            Some(self.audio_data[channel].borrow_mut())
        } else {
            None
        }
    }

    /// Copy a single sample out of the buffer, or None if it is not available.
    pub fn sample(&self, channel: usize, position: usize) -> Option<f64> {
        if channel < self.channels() && position < self.buffer_size {
            Some(self.audio_data[channel][position])
        } else {
            None
        }
    }

    /// Get a mutable reference to a single sample of the buffer, or None if it
    /// is not available.
    pub fn sample_mut(&mut self, channel: usize, position: usize) -> Option<&mut f64> {
        self.audio_data
            .get_mut(channel)
            .and_then(|v| v.get_mut(position))
    }

    /// Apply a constant gain factor across the whole buffer.
    pub fn apply_gain(&mut self, gain: f64) {
        for ch in &mut self.audio_data {
            ch.iter_mut().for_each(|v| *v *= gain);
        }
    }

    /// Copy an interleaved slice at the given sample position. The slice is
    /// assumed to contain as many channels as the buffer. The slice sample size
    /// and the position must fit so that the slice can be fully copied into the
    /// buffer without overflow.
    pub unsafe fn copy_interleaved(&mut self, slice: &[f64], position: usize) {
        let inner = Self::new(self.channels(), slice);
        for idx in 0..self.channels() {
            let src_ptr = inner[idx].as_ptr();
            let dst_ptr = self.audio_data[idx].as_mut_ptr().add(position);
            std::ptr::copy_nonoverlapping(src_ptr, dst_ptr, inner.buffer_size);
        }
    }

    /// Copies a slice of the buffer into a new one.
    pub fn copy_slice(&self, range: Range<usize>) -> Self {
        if range.start >= range.end {
            Self {
                audio_data: SmallVec::new(),
                buffer_size: 0,
            }
        } else {
            let range_len = range.end - range.start;
            let audio_data = (0..self.channels())
                .map(move |i| self.audio_data[i][range.clone()].to_vec())
                .collect();
            Self {
                audio_data,
                buffer_size: range_len,
            }
        }
    }

    /// Consume the buffer into an interleaved `Vec`.
    pub fn interleave(self) -> Vec<f64> {
        let channels = self.channels();
        let out_size = channels * self.buffer_size;
        (0..out_size)
            .map(|i| self.audio_data[i % channels][i / channels])
            .collect()
    }

    /// Consumes the buffer by moving the data into the given interleaved
    /// buffer.
    pub fn move_into_interleaved(self, data: &mut [f64]) {
        assert_eq!(self.channels() * self.buffer_size, data.len());
        let interleaved = self.interleave();
        unsafe {
            std::ptr::copy_nonoverlapping(interleaved.as_ptr(), data.as_mut_ptr(), data.len());
        }
    }

    /// Returns an iterator over the channels of this buffer.
    pub fn iter(&self) -> Iter {
        Iter {
            buffer: self,
            position: 0,
        }
    }
}

impl Index<usize> for AudioBuffer {
    type Output = [f64];

    fn index(&self, channel: usize) -> &Self::Output {
        self.channel(channel).unwrap()
    }
}

impl Index<(usize, usize)> for AudioBuffer {
    type Output = f64;

    fn index(&self, (channel, position): (usize, usize)) -> &Self::Output {
        &self.audio_data[channel][position]
    }
}

impl IndexMut<usize> for AudioBuffer {
    fn index_mut(&mut self, channel: usize) -> &mut Self::Output {
        self.channel_mut(channel).unwrap()
    }
}

impl IndexMut<(usize, usize)> for AudioBuffer {
    fn index_mut(&mut self, (channel, position): (usize, usize)) -> &mut Self::Output {
        &mut self.audio_data[channel][position]
    }
}

impl Into<Vec<f64>> for AudioBuffer {
    fn into(self) -> Vec<f64> {
        self.interleave()
    }
}

impl<'a> Iterator for Iter<'a> {
    type Item = &'a [f64];

    fn next(&mut self) -> Option<Self::Item> {
        self.buffer.channel({
            let pos = self.position;
            self.position += 1;
            pos
        })
    }
}