MP3 Support (#14)

src/encoder/flac.rs

@@ -0,0 +1,52 @@
+use flacenc::bitsink::ByteSink;
+use flacenc::component::BitRepr;
+use flacenc::error::Verify;
+
+use super::execute_with_result;
+use super::EncodedStream;
+use super::Encoder;
+use super::Samples;
+
+#[derive(Debug)]
+pub struct FlacEncoder;
+
+#[async_trait::async_trait]
+impl Encoder for FlacEncoder {
+    async fn encode(&self, samples: Samples) -> anyhow::Result<EncodedStream> {
+        let source = flacenc::source::MemSource::from_samples(
+            &samples.samples,
+            samples.channels as usize,
+            samples.bits_per_sample as usize,
+            samples.sample_rate as usize,
+        );
+
+        let config = flacenc::config::Encoder::default()
+            .into_verified()
+            .map_err(|e| anyhow::anyhow!("Failed to verify encoder config: {:?}", e))?;
+
+        let (tx, rx) = tokio::sync::oneshot::channel();
+
+        rayon::spawn(execute_with_result(
+            move || {
+                let flac_stream = flacenc::encode_with_fixed_block_size(
+                    &config,
+                    source,
+                    config.block_size,
+                )
+                .map_err(|e| anyhow::anyhow!("Failed to encode flac: {:?}", e))?;
+
+                let mut byte_sink = ByteSink::new();
+                flac_stream
+                    .write(&mut byte_sink)
+                    .map_err(|e| anyhow::anyhow!("Failed to write flac stream: {:?}", e))?;
+
+                Ok(byte_sink.into_inner())
+            },
+            tx,
+        ));
+
+        let byte_sink: Vec<u8> = rx.await??;
+
+        Ok(EncodedStream::new(byte_sink))
+    }
+}

src/encoder/mod.rs

@@ -0,0 +1,111 @@
+mod flac;
+#[cfg(feature = "mp3")]
+mod mp3;
+
+use std::{path::Path, str::FromStr};
+
+use anyhow::Result;
+use tokio::sync::oneshot::Sender;
+
+use self::{flac::FlacEncoder, mp3::Mp3Encoder};
+
+#[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
+pub enum Format {
+    Flac,
+    #[cfg(feature = "mp3")]
+    Mp3,
+}
+
+impl FromStr for Format {
+    type Err = anyhow::Error;
+
+    fn from_str(s: &str) -> Result<Self> {
+        match s {
+            "flac" => Ok(Format::Flac),
+            #[cfg(feature = "mp3")]
+            "mp3" => Ok(Format::Mp3),
+            _ => Err(anyhow::anyhow!("Unsupported format")),
+        }
+    }
+}
+
+impl Format {
+
+    pub fn extension(&self) -> &'static str {
+        match self {
+            Format::Flac => "flac",
+            #[cfg(feature = "mp3")]
+            Format::Mp3 => "mp3",
+        }
+    }
+
+}
+
+const FLAC_ENCODER: &FlacEncoder = &FlacEncoder;
+#[cfg(feature = "mp3")]
+const MP3_ENCODER: &Mp3Encoder = &Mp3Encoder;
+
+pub fn get_encoder(format: Format) -> &'static dyn Encoder {
+    match format {
+        Format::Flac => FLAC_ENCODER,
+        #[cfg(feature = "mp3")]
+        Format::Mp3 => MP3_ENCODER,
+    }
+}
+
+#[async_trait::async_trait]
+pub trait Encoder {
+    async fn encode(&self, samples: Samples) -> Result<EncodedStream>;
+}
+
+pub struct Samples {
+    pub samples: Vec<i32>,
+    pub sample_rate: u32,
+    pub channels: u32,
+    pub bits_per_sample: u32,
+}
+
+impl Samples {
+    pub fn new(samples: Vec<i32>, sample_rate: u32, channels: u32, bits_per_sample: u32) -> Self {
+        Samples {
+            samples,
+            sample_rate,
+            channels,
+            bits_per_sample,
+        }
+    }
+}
+
+pub struct EncodedStream {
+    pub stream: Vec<u8>,
+}
+
+impl EncodedStream {
+    pub fn new(stream: Vec<u8>) -> Self {
+        EncodedStream { stream }
+    }
+
+    pub async fn write_to_file<P: AsRef<Path>>(&self, path: P) -> Result<()> {
+        if !path.as_ref().exists() {
+            tokio::fs::create_dir_all(
+                path.as_ref()
+                    .parent()
+                    .ok_or(anyhow::anyhow!("Could not create path"))?,
+            ).await?;
+        }
+        tokio::fs::write(path, &self.stream).await?;
+        Ok(())
+    }
+}
+
+pub fn execute_with_result<F, T>(func: F, tx: Sender<anyhow::Result<T>>) -> impl FnOnce()
+where
+    F: FnOnce() -> anyhow::Result<T> + Send + 'static,
+    T: Send + 'static,
+{
+    move || {
+        let result = func();
+        // Ignore the error if the receiver has been dropped
+        let _ = tx.send(result);
+    }
+}

src/encoder/mp3.rs

@@ -0,0 +1,72 @@
+use anyhow::anyhow;
+use anyhow::Ok;
+use mp3lame_encoder::Builder;
+use mp3lame_encoder::FlushNoGap;
+use mp3lame_encoder::InterleavedPcm;
+
+use super::execute_with_result;
+use super::EncodedStream;
+use super::Encoder;
+use super::Samples;
+
+pub struct Mp3Encoder;
+
+impl Mp3Encoder {
+    fn build_encoder(
+        &self,
+        sample_rate: u32,
+        channels: u32,
+    ) -> anyhow::Result<mp3lame_encoder::Encoder> {
+        let mut builder = Builder::new().ok_or(anyhow::anyhow!("Failed to create mp3 encoder"))?;
+
+        builder
+            .set_sample_rate(sample_rate)
+            .map_err(|e| anyhow::anyhow!("Failed to set sample rate for mp3 encoder: {}", e))?;
+        builder.set_num_channels(channels as u8).map_err(|e| {
+            anyhow::anyhow!("Failed to set number of channels for mp3 encoder: {}", e)
+        })?;
+        builder
+            .set_brate(mp3lame_encoder::Birtate::Kbps160)
+            .map_err(|e| anyhow::anyhow!("Failed to set bitrate for mp3 encoder: {}", e))?;
+
+        builder
+            .build()
+            .map_err(|e| anyhow::anyhow!("Failed to build mp3 encoder: {}", e))
+    }
+}
+
+#[async_trait::async_trait]
+impl Encoder for Mp3Encoder {
+    async fn encode(&self, samples: Samples) -> anyhow::Result<EncodedStream> {
+        let mut mp3_encoder = self.build_encoder(samples.sample_rate, samples.channels)?;
+
+        let (tx, rx) = tokio::sync::oneshot::channel();
+
+        rayon::spawn(execute_with_result(
+            move || {
+                let samples: Vec<i16> = samples.samples.iter().map(|&x| x as i16).collect();
+                let input = InterleavedPcm(samples.as_slice());
+                let mut mp3_out_buffer = Vec::with_capacity(mp3lame_encoder::max_required_buffer_size(samples.len()));
+                let encoded_size = mp3_encoder
+                    .encode(input, mp3_out_buffer.spare_capacity_mut())
+                    .map_err(|e| anyhow!("Failed to encode mp3: {}", e))?;
+                unsafe {
+                    mp3_out_buffer.set_len(mp3_out_buffer.len().wrapping_add(encoded_size));
+                }
+
+                let encoded_size = mp3_encoder
+                    .flush::<FlushNoGap>(mp3_out_buffer.spare_capacity_mut())
+                    .map_err(|e| anyhow!("Failed to flush mp3 encoder: {}", e))?;
+                unsafe {
+                    mp3_out_buffer.set_len(mp3_out_buffer.len().wrapping_add(encoded_size));
+                }
+                Ok(mp3_out_buffer)
+            },
+            tx,
+        ));
+
+        let mp3_out_buffer = rx.await??;
+
+        Ok(EncodedStream::new(mp3_out_buffer))
+    }
+}