Lightningbeam/daw-backend/src/io/wav_writer.rs

126 lines
4.2 KiB
Rust

/// Incremental WAV file writer for streaming audio to disk
use std::fs::File;
use std::io::{self, Seek, SeekFrom, Write};
use std::path::Path;
/// WAV file writer that supports incremental writing
pub struct WavWriter {
file: File,
sample_rate: u32,
channels: u32,
frames_written: usize,
}
impl WavWriter {
/// Create a new WAV file and write initial header
/// The header is written with placeholder sizes that will be updated on finalization
pub fn create(path: impl AsRef<Path>, sample_rate: u32, channels: u32) -> io::Result<Self> {
let mut file = File::create(path)?;
// Write initial WAV header with placeholder sizes
write_wav_header(&mut file, sample_rate, channels, 0)?;
Ok(Self {
file,
sample_rate,
channels,
frames_written: 0,
})
}
/// Append audio samples to the file
/// Expects interleaved f32 samples in range [-1.0, 1.0]
pub fn write_samples(&mut self, samples: &[f32]) -> io::Result<()> {
// Convert f32 samples to 16-bit PCM
let pcm_data: Vec<u8> = samples
.iter()
.flat_map(|&sample| {
let clamped = sample.clamp(-1.0, 1.0);
let pcm_value = (clamped * 32767.0) as i16;
pcm_value.to_le_bytes()
})
.collect();
self.file.write_all(&pcm_data)?;
self.frames_written += samples.len() / self.channels as usize;
Ok(())
}
/// Get the current number of frames written
pub fn frames_written(&self) -> usize {
self.frames_written
}
/// Get the current duration in seconds
pub fn duration(&self) -> f64 {
self.frames_written as f64 / self.sample_rate as f64
}
/// Finalize the WAV file by updating the header with correct sizes
pub fn finalize(mut self) -> io::Result<()> {
// Flush any remaining data
self.file.flush()?;
// Calculate total data size
let data_size = self.frames_written * self.channels as usize * 2; // 2 bytes per sample (16-bit)
// WAV file structure:
// RIFF header (12 bytes): "RIFF" + size + "WAVE"
// fmt chunk (24 bytes): "fmt " + size + format data
// data chunk header (8 bytes): "data" + size
// Total header = 44 bytes
// RIFF chunk size = everything after offset 8 = 4 (WAVE) + 24 (fmt) + 8 (data header) + data_size
let riff_chunk_size = 36 + data_size; // 36 = size from "WAVE" to end of data chunk header
// Seek to RIFF chunk size (offset 4)
self.file.seek(SeekFrom::Start(4))?;
self.file.write_all(&(riff_chunk_size as u32).to_le_bytes())?;
// Seek to data chunk size (offset 40)
self.file.seek(SeekFrom::Start(40))?;
self.file.write_all(&(data_size as u32).to_le_bytes())?;
// Flush and sync to ensure all data is written to disk before file is closed
self.file.flush()?;
self.file.sync_all()?;
Ok(())
}
}
/// Write WAV header with specified parameters
fn write_wav_header(file: &mut File, sample_rate: u32, channels: u32, frames: usize) -> io::Result<()> {
let bytes_per_sample = 2u16; // 16-bit PCM
let data_size = (frames * channels as usize * bytes_per_sample as usize) as u32;
// RIFF chunk size = everything after offset 8
// = 4 (WAVE) + 24 (fmt chunk) + 8 (data chunk header) + data_size
let riff_chunk_size = 36 + data_size;
// RIFF header
file.write_all(b"RIFF")?;
file.write_all(&riff_chunk_size.to_le_bytes())?;
file.write_all(b"WAVE")?;
// fmt chunk
file.write_all(b"fmt ")?;
file.write_all(&16u32.to_le_bytes())?; // fmt chunk size
file.write_all(&1u16.to_le_bytes())?; // PCM format
file.write_all(&(channels as u16).to_le_bytes())?;
file.write_all(&sample_rate.to_le_bytes())?;
let byte_rate = sample_rate * channels * bytes_per_sample as u32;
file.write_all(&byte_rate.to_le_bytes())?;
let block_align = channels as u16 * bytes_per_sample;
file.write_all(&block_align.to_le_bytes())?;
file.write_all(&(bytes_per_sample * 8).to_le_bytes())?; // bits per sample
// data chunk header
file.write_all(b"data")?;
file.write_all(&data_size.to_le_bytes())?;
Ok(())
}