use std::{ ffi::c_void, io, pin::Pin, task::{Context, Poll}, }; use tokio::{ io::{AsyncRead, AsyncReadExt, AsyncWriteExt}, sync::{mpsc, oneshot}, }; use crate::{ error::FlacError, pcm::{PcmChunk, PcmFormat}, stream::ManagedAsyncReader, util::le_bytes_to_interleaved_i32, }; /// Channel capacity for async message passing between tasks. const CHANNEL_CAPACITY: usize = 8; /// Number of PCM frames to process per chunk (4096 frames = ~93ms at 44.1kHz). const PCM_FRAMES_PER_CHUNK: usize = 4096; /// An async stream that encodes PCM audio into FLAC format. /// /// This struct implements `AsyncRead`, allowing you to read encoded FLAC data /// as it becomes available. The encoding happens in a background task. /// /// # Example /// /// ```no_run /// use pmoflac::{encode_flac_stream, EncoderOptions, PcmFormat}; /// use tokio::io::AsyncReadExt; /// /// # #[tokio::main] /// # async fn main() -> Result<(), Box> { /// let pcm_data: &[u8] = &[/* 16-bit stereo PCM */]; /// let format = PcmFormat { /// sample_rate: 44_100, /// channels: 2, /// bits_per_sample: 16, /// }; /// /// let mut stream = encode_flac_stream(pcm_data, format, EncoderOptions::default()).await?; /// let mut flac_output = Vec::new(); /// stream.read_to_end(&mut flac_output).await?; /// stream.wait().await?; /// # Ok(()) /// # } /// ``` pub struct FlacEncodedStream { format: PcmFormat, reader: ManagedAsyncReader, } impl FlacEncodedStream { /// Returns the PCM format used for encoding. pub fn format(&self) -> PcmFormat { self.format } /// Consumes the stream and returns its components. pub fn into_parts(self) -> (PcmFormat, ManagedAsyncReader) { (self.format, self.reader) } /// Waits for the background encoding task to complete. /// /// This should be called after reading all data to ensure proper cleanup /// and to catch any errors that occurred during encoding. pub async fn wait(self) -> Result<(), FlacError> { self.reader.wait().await } } impl tokio::io::AsyncRead for FlacEncodedStream { fn poll_read( mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut tokio::io::ReadBuf<'_>, ) -> Poll> { Pin::new(&mut self.reader).poll_read(cx, buf) } } /// Options for configuring FLAC encoding. #[derive(Debug, Clone)] pub struct EncoderOptions { /// Compression level (0-12). Higher means better compression but slower. /// Default: 5 (balanced) pub compression_level: u32, /// Whether to verify the encoding by decoding in parallel. /// Default: false (disabled for performance) pub verify: bool, /// Total number of samples (optional). If known, improves seeking in output. pub total_samples: Option, /// Block size in samples (optional). If None, libFLAC chooses automatically. /// Typical values: 1152, 2304, 4096. pub block_size: Option, } impl Default for EncoderOptions { fn default() -> Self { Self { compression_level: 5, verify: false, total_samples: None, block_size: None, } } } /// Encodes PCM audio data into a FLAC stream. /// /// This function spawns background tasks to perform the encoding asynchronously. /// The returned `FlacEncodedStream` implements `AsyncRead` for streaming the FLAC output. /// /// # Threading Model /// /// - A Tokio task reads PCM chunks and converts them to i32 samples /// - A blocking task (via `spawn_blocking`) runs the libFLAC encoder /// - The encoder's write callback sends encoded data via a channel /// - Another Tokio task writes FLAC data to an internal duplex stream /// /// This architecture ensures true streaming: FLAC frames are produced as soon as /// enough PCM data is available, without waiting for the entire input. /// /// # Arguments /// /// * `reader` - Any async reader containing PCM audio in little-endian interleaved format /// * `format` - Describes the PCM format (sample rate, channels, bit depth) /// * `options` - Encoding options (compression level, verify, etc.) /// /// # PCM Input Format /// /// The PCM data must be: /// - **Little-endian** byte order /// - **Interleaved** channels (L, R, L, R for stereo) /// - **Signed integers** with bit depth matching `format.bits_per_sample` /// /// # Returns /// /// A `FlacEncodedStream` that can be read to obtain FLAC-encoded data. /// /// # Errors /// /// Returns an error if: /// - The PCM format is invalid (e.g., unsupported bit depth) /// - The input stream has incomplete sample data /// - libFLAC initialization fails /// - An I/O error occurs /// /// # Example /// /// ```no_run /// use pmoflac::{encode_flac_stream, EncoderOptions, PcmFormat}; /// use tokio::io::AsyncReadExt; /// /// # #[tokio::main] /// # async fn main() -> Result<(), Box> { /// // Generate 1 second of silence at 44.1kHz stereo 16-bit /// let sample_rate = 44_100u32; /// let channels = 2u8; /// let pcm_len = sample_rate as usize * channels as usize * 2; /// /// let format = PcmFormat { /// sample_rate, /// channels, /// bits_per_sample: 16, /// }; /// /// let options = EncoderOptions { /// compression_level: 8, /// total_samples: Some(sample_rate as u64), /// ..Default::default() /// }; /// /// let mut stream = encode_flac_stream( /// tokio::io::repeat(0).take(pcm_len as u64), /// format, /// options, /// ).await?; /// let mut flac_data = Vec::new(); /// stream.read_to_end(&mut flac_data).await?; /// stream.wait().await?; /// /// println!("Encoded {} bytes of PCM to {} bytes of FLAC", /// pcm_len, flac_data.len()); /// # Ok(()) /// # } /// ``` pub async fn encode_flac_stream( reader: R, format: PcmFormat, options: EncoderOptions, ) -> Result where R: AsyncRead + Unpin + Send + 'static, { format .validate() .map_err(|msg| FlacError::Unsupported(format!("invalid PCM format: {msg}")))?; if options.compression_level > 12 { return Err(FlacError::Unsupported( "compression level must be in 0..=12".into(), )); } let (pcm_tx, pcm_rx) = mpsc::channel::>(CHANNEL_CAPACITY); let format_for_reader = format; tokio::spawn(async move { let _ = feed_pcm_chunks(reader, format_for_reader, pcm_tx).await; }); let (flac_reader, mut flac_writer) = tokio::io::duplex(256 * 1024); let (flac_tx, mut flac_rx) = mpsc::channel::, FlacError>>(CHANNEL_CAPACITY); let (init_tx, init_rx) = oneshot::channel::>(); let format_for_encoder = format; let options_for_encoder = options; let blocking_handle = tokio::task::spawn_blocking(move || { run_encoder( format_for_encoder, options_for_encoder, pcm_rx, flac_tx, init_tx, ) }); let writer_handle = tokio::spawn(async move { while let Some(chunk) = flac_rx.recv().await { let bytes = chunk?; if bytes.is_empty() { continue; } flac_writer.write_all(&bytes).await?; } flac_writer.shutdown().await?; match blocking_handle.await { Ok(res) => res, Err(err) => Err(FlacError::TaskJoin { role: "flac-encode", details: err.to_string(), }), } }); init_rx.await.map_err(|_| FlacError::ChannelClosed)??; let reader = ManagedAsyncReader::new("flac-encode-writer", flac_reader, writer_handle); Ok(FlacEncodedStream { format, reader }) } async fn feed_pcm_chunks( reader: R, format: PcmFormat, tx: mpsc::Sender>, ) -> Result<(), FlacError> where R: AsyncRead + Unpin, { let bytes_per_frame = format.bytes_per_sample() * format.channels as usize; let chunk_bytes = PCM_FRAMES_PER_CHUNK * bytes_per_frame; let mut pending = Vec::with_capacity(chunk_bytes * 2); let mut reader = tokio::io::BufReader::new(reader); loop { while pending.len() >= chunk_bytes { let samples = le_bytes_to_interleaved_i32(&pending[..chunk_bytes], format.bits_per_sample) .map_err(|msg| FlacError::Encode(msg))?; pending.drain(..chunk_bytes); let frames = (samples.len() / format.channels as usize) as u32; let chunk = PcmChunk::new(samples, frames, format.channels); if tx.send(Ok(chunk)).await.is_err() { return Ok(()); } } let read = reader.read_buf(&mut pending).await?; if read == 0 { break; } } if !pending.is_empty() { if pending.len() % bytes_per_frame != 0 { let msg = "PCM stream ended with a partial frame (incomplete sample data)".to_string(); let _ = tx.send(Err(FlacError::Encode(msg.clone()))).await; return Err(FlacError::Encode(msg)); } let samples = le_bytes_to_interleaved_i32(&pending, format.bits_per_sample) .map_err(|msg| FlacError::Encode(msg))?; let frames = (samples.len() / format.channels as usize) as u32; let chunk = PcmChunk::new(samples, frames, format.channels); let _ = tx.send(Ok(chunk)).await; } Ok(()) } fn run_encoder( format: PcmFormat, options: EncoderOptions, mut rx: mpsc::Receiver>, tx: mpsc::Sender, FlacError>>, init_tx: oneshot::Sender>, ) -> Result<(), FlacError> { use libflac_sys::*; unsafe { let encoder = FLAC__stream_encoder_new(); if encoder.is_null() { let _ = init_tx.send(Err(FlacError::LibFlacInit( "FLAC__stream_encoder_new returned null".into(), ))); return Err(FlacError::LibFlacInit( "FLAC__stream_encoder_new returned null".into(), )); } let _encoder_guard = EncoderHandle { ptr: encoder }; let mut state = EncoderClientState::new(tx); let ensure = |ok: FLAC__bool, msg: &str| { if ok == 0 { Err(FlacError::LibFlacInit(msg.into())) } else { Ok(()) } }; ensure( FLAC__stream_encoder_set_channels(encoder, format.channels as u32), "set_channels failed", )?; ensure( FLAC__stream_encoder_set_bits_per_sample(encoder, format.bits_per_sample as u32), "set_bits_per_sample failed", )?; ensure( FLAC__stream_encoder_set_sample_rate(encoder, format.sample_rate), "set_sample_rate failed", )?; ensure( FLAC__stream_encoder_set_compression_level(encoder, options.compression_level), "set_compression_level failed", )?; ensure( FLAC__stream_encoder_set_streamable_subset(encoder, 1), "set_streamable_subset failed", )?; ensure( FLAC__stream_encoder_set_verify(encoder, if options.verify { 1 } else { 0 }), "set_verify failed", )?; if let Some(total) = options.total_samples { ensure( FLAC__stream_encoder_set_total_samples_estimate(encoder, total), "set_total_samples_estimate failed", )?; } if let Some(block_size) = options.block_size { ensure( FLAC__stream_encoder_set_blocksize(encoder, block_size), "set_blocksize failed", )?; } let init_status = FLAC__stream_encoder_init_stream( encoder, Some(write_callback), None, None, None, &mut state as *mut EncoderClientState as *mut c_void, ); if init_status != libflac_sys::FLAC__STREAM_ENCODER_INIT_STATUS_OK { let msg = format!("init_stream failed: status {init_status}"); let _ = init_tx.send(Err(FlacError::LibFlacInit(msg.clone()))); return Err(FlacError::LibFlacInit(msg)); } let _ = init_tx.send(Ok(())); while let Some(chunk_result) = rx.blocking_recv() { let chunk = match chunk_result { Ok(chunk) => chunk, Err(err) => { FLAC__stream_encoder_finish(encoder); return Err(err); } }; if chunk.frames == 0 { continue; } let success = FLAC__stream_encoder_process_interleaved( encoder, chunk.data.as_ptr(), chunk.frames, ); if success == 0 { if let Some(err) = state.error.take() { FLAC__stream_encoder_finish(encoder); return Err(err); } FLAC__stream_encoder_finish(encoder); return Err(FlacError::Encode("libFLAC reported encode failure".into())); } } let finish_ok = FLAC__stream_encoder_finish(encoder); if finish_ok == 0 { if let Some(err) = state.error.take() { return Err(err); } return Err(FlacError::Encode( "libFLAC failed to finalize stream".into(), )); } if let Some(err) = state.error.take() { return Err(err); } } Ok(()) } struct EncoderClientState { tx: mpsc::Sender, FlacError>>, error: Option, } impl EncoderClientState { fn new(tx: mpsc::Sender, FlacError>>) -> Self { Self { tx, error: None } } } struct EncoderHandle { ptr: *mut libflac_sys::FLAC__StreamEncoder, } impl Drop for EncoderHandle { fn drop(&mut self) { unsafe { if !self.ptr.is_null() { libflac_sys::FLAC__stream_encoder_delete(self.ptr); } } } } unsafe extern "C" fn write_callback( _encoder: *const libflac_sys::FLAC__StreamEncoder, buffer: *const libflac_sys::FLAC__byte, bytes: usize, _samples: u32, _current_frame: u32, client_data: *mut c_void, ) -> libflac_sys::FLAC__StreamEncoderWriteStatus { let state = &mut *(client_data as *mut EncoderClientState); let slice = std::slice::from_raw_parts(buffer, bytes); match state.tx.blocking_send(Ok(slice.to_vec())) { Ok(_) => libflac_sys::FLAC__STREAM_ENCODER_WRITE_STATUS_OK, Err(_) => { state.error = Some(FlacError::LibFlacWrite( "failed to send encoded data (receiver dropped)".into(), )); libflac_sys::FLAC__STREAM_ENCODER_WRITE_STATUS_FATAL_ERROR } } }