peerspeak/src/dsp/wav.rs

//! A minimal canonical-PCM WAV reader (the analysis-side counterpart to
//! `audio::recorder::WavWriter`). Handles 16-bit integer PCM — the only format
//! PeerSpeak writes — in mono or interleaved multi-channel, downmixing to mono
//! `f32` in `[-1, 1]` for analysis. I/O lives here, at the edge; the rest of the
//! `dsp` module is pure.

use std::fs;
use std::path::Path;

/// A decoded WAV: mono `f32` samples plus the original sample rate.
pub struct WavData {
    pub samples: Vec<f32>,
    pub sample_rate: u32,
}

/// Reads a 16-bit PCM WAV file, downmixing any channels to mono `f32`.
///
/// Returns `Err` with a human-readable reason if the file is missing, truncated,
/// not RIFF/WAVE, or not 16-bit PCM. The parser walks the chunk list rather than
/// assuming a fixed 44-byte header, so files with extra chunks (`LIST`, `fact`,
/// …) still read.
pub fn read(path: &Path) -> Result<WavData, String> {
    let bytes = fs::read(path).map_err(|e| format!("cannot read {}: {e}", path.display()))?;
    if bytes.len() < 12 || &bytes[0..4] != b"RIFF" || &bytes[8..12] != b"WAVE" {
        return Err("not a RIFF/WAVE file".to_string());
    }

    let mut channels = 0u16;
    let mut sample_rate = 0u32;
    let mut bits = 0u16;
    let mut data: Option<&[u8]> = None;

    // Walk chunks starting after the 12-byte RIFF/WAVE header.
    let mut pos = 12usize;
    while pos + 8 <= bytes.len() {
        let id = &bytes[pos..pos + 4];
        let size = u32::from_le_bytes([
            bytes[pos + 4],
            bytes[pos + 5],
            bytes[pos + 6],
            bytes[pos + 7],
        ]) as usize;
        let body_start = pos + 8;
        let body_end = (body_start + size).min(bytes.len());
        match id {
            b"fmt " if size >= 16 => {
                let fmt = &bytes[body_start..body_end];
                let audio_format = u16::from_le_bytes([fmt[0], fmt[1]]);
                channels = u16::from_le_bytes([fmt[2], fmt[3]]);
                sample_rate = u32::from_le_bytes([fmt[4], fmt[5], fmt[6], fmt[7]]);
                bits = u16::from_le_bytes([fmt[14], fmt[15]]);
                if audio_format != 1 {
                    return Err(format!(
                        "unsupported WAV format tag {audio_format} (need PCM=1)"
                    ));
                }
            }
            b"data" => {
                data = Some(&bytes[body_start..body_end]);
            }
            _ => {}
        }
        // Chunks are word-aligned: an odd size is padded with one byte.
        pos = body_start + size + (size & 1);
    }

    if bits != 16 {
        return Err(format!("only 16-bit PCM supported, got {bits}-bit"));
    }
    let channels = channels.max(1);
    let data = data.ok_or("no data chunk")?;

    // Interleaved S16LE -> per-frame channel average -> mono f32.
    let frame_bytes = 2 * channels as usize;
    let mut samples = Vec::with_capacity(data.len() / frame_bytes.max(1));
    for frame in data.chunks_exact(frame_bytes) {
        let mut acc = 0i32;
        for ch in frame.chunks_exact(2) {
            acc += i16::from_le_bytes([ch[0], ch[1]]) as i32;
        }
        let avg = acc as f32 / channels as f32;
        samples.push(avg / 32768.0);
    }

    Ok(WavData {
        samples,
        sample_rate,
    })
}

/// Writes mono `f32` samples (clamped to `[-1, 1]`) as a 16-bit PCM WAV. Used by
/// the `specview gen` command to materialize conjured test signals.
pub fn write(path: &Path, samples: &[f32], sample_rate: u32) -> Result<(), String> {
    let pcm: Vec<i16> = samples
        .iter()
        .map(|&s| (s.clamp(-1.0, 1.0) * 32767.0).round() as i16)
        .collect();
    let data_bytes = (pcm.len() * 2) as u32;
    let byte_rate = sample_rate * 2; // mono, 2 bytes/sample
    let mut out = Vec::with_capacity(44 + pcm.len() * 2);
    out.extend_from_slice(b"RIFF");
    out.extend_from_slice(&(36 + data_bytes).to_le_bytes());
    out.extend_from_slice(b"WAVE");
    out.extend_from_slice(b"fmt ");
    out.extend_from_slice(&16u32.to_le_bytes());
    out.extend_from_slice(&1u16.to_le_bytes()); // PCM
    out.extend_from_slice(&1u16.to_le_bytes()); // mono
    out.extend_from_slice(&sample_rate.to_le_bytes());
    out.extend_from_slice(&byte_rate.to_le_bytes());
    out.extend_from_slice(&2u16.to_le_bytes()); // block align
    out.extend_from_slice(&16u16.to_le_bytes()); // bits
    out.extend_from_slice(b"data");
    out.extend_from_slice(&data_bytes.to_le_bytes());
    for s in pcm {
        out.extend_from_slice(&s.to_le_bytes());
    }
    fs::write(path, out).map_err(|e| format!("cannot write {}: {e}", path.display()))
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn write_then_read_round_trips() {
        let dir = std::env::temp_dir();
        let path = dir.join(format!("dsp-wav-rt-{}.wav", std::process::id()));
        let sig: Vec<f32> = (0..1000).map(|i| (i as f32 * 0.01).sin() * 0.5).collect();
        write(&path, &sig, 48_000).unwrap();

        let back = read(&path).unwrap();
        assert_eq!(back.sample_rate, 48_000);
        assert_eq!(back.samples.len(), sig.len());
        // 16-bit quantization tolerance.
        for (a, b) in sig.iter().zip(&back.samples) {
            assert!((a - b).abs() < 1e-3, "{a} vs {b}");
        }
        let _ = fs::remove_file(&path);
    }

    #[test]
    fn rejects_non_wave() {
        let dir = std::env::temp_dir();
        let path = dir.join(format!("dsp-wav-bad-{}.bin", std::process::id()));
        fs::write(&path, b"not a wav at all").unwrap();
        assert!(read(&path).is_err());
        let _ = fs::remove_file(&path);
    }
}