wav-codec 0.1.0

use std::fmt;
use std::fs::File;
use std::io::{self, BufWriter, Seek, SeekFrom, Write};
use std::path::Path;

use crate::sample::Codec;
use crate::{
    Error, Result, Sample, SampleFormat, WavSpec, WavSpecEx, GUID_IEEE_FLOAT, GUID_PCM,
    TAG_EXTENSIBLE, TAG_IEEE_FLOAT, TAG_PCM,
};

/// Size of the encode block flushed to the sink in one write.
const BLOCK: usize = 64 * 1024;

/// Block-buffered WAV encoder over any `io::Write + io::Seek`.
///
/// The header is written at construction with placeholder sizes. Samples
/// are encoded into a 64 KiB block that is flushed whole. [`finalize`]
/// writes the pad byte for an odd `data` size and patches the RIFF size,
/// the `data` size, and (for float formats) the `fact` frame count.
/// An incomplete last frame is truncated to whole frames in the header
/// and returns [`Error::UnfinishedFrame`]. Dropping the writer runs the
/// same steps and discards the result, so call [`finalize`] to observe
/// that error.
///
/// Header layout by spec:
///
/// * integer PCM, up to 2 channels, valid bits filling the container:
///   `fmt ` size 16.
/// * IEEE float, up to 2 channels: `fmt ` size 18 plus a `fact` chunk.
/// * a container wider than the valid bits, or more than 2 channels:
///   `fmt ` size 40 (`WAVE_FORMAT_EXTENSIBLE`, channel mask 0), plus
///   `fact` for float.
///
/// [`finalize`]: Self::finalize
pub struct WavWriter<W: Write + Seek> {
    inner: W,
    spec: WavSpecEx,
    codec: Codec,
    block: Vec<u8>,
    /// Bytes of sample data written or buffered.
    data_len: u64,
    /// Largest `data_len` the u32 size fields can describe.
    max_data: u64,
    /// Stream position of the `RIFF` tag.
    base: u64,
    /// Offset from `base` of the `fact` frame-count field.
    fact_pos: Option<u64>,
    /// Offset from `base` of the `data` size field.
    data_size_pos: u64,
    /// Offset from `base` of the first sample.
    data_start: u64,
    finalized: bool,
}

impl WavWriter<BufWriter<File>> {
    /// Create (or truncate) a file and write the header for `spec`.
    pub fn create<P: AsRef<Path>>(path: P, spec: WavSpec) -> Result<Self> {
        let file = File::create(path)?;
        WavWriter::new(BufWriter::new(file), spec)
    }
}

impl<W: Write + Seek> WavWriter<W> {
    /// Write the header for `spec` to `writer`, using the narrowest
    /// container that holds `bits_per_sample`.
    pub fn new(writer: W, spec: WavSpec) -> Result<Self> {
        if spec.bits_per_sample == 0 {
            return Err(Error::Format("zero bits per sample"));
        }
        let bits = u32::from(spec.bits_per_sample);
        match spec.sample_format {
            SampleFormat::Int => {
                if bits > 32 {
                    return Err(Error::TooWide);
                }
            }
            SampleFormat::Float => {
                if bits != 32 && bits != 64 {
                    // A well-formed spec that is simply outside this crate's scope is
                    // Unsupported, not a malformed-container Format error. This matches
                    // new_with_spec_ex and the reader for the same condition.
                    return Err(Error::Unsupported(
                        "float samples other than 32-bit and 64-bit",
                    ));
                }
            }
        }
        let bytes_per_sample = match spec.sample_format {
            SampleFormat::Int => ((bits + 7) / 8) as u16,
            SampleFormat::Float => (bits / 8) as u16,
        };
        WavWriter::new_with_spec_ex(
            writer,
            WavSpecEx {
                spec,
                bytes_per_sample,
            },
        )
    }

    /// Write the header for `spec` to `writer`, with an explicit container
    /// width (for example 24 valid bits in 4-byte containers).
    pub fn new_with_spec_ex(mut writer: W, spec: WavSpecEx) -> Result<Self> {
        let codec = validate(&spec)?;
        let channels = spec.spec.channels;
        let container = spec.bytes_per_sample;
        let bits = spec.spec.bits_per_sample;
        let container_bits = container * 8;
        let block_align = u32::from(channels) * u32::from(container);
        if block_align > u32::from(u16::MAX) {
            return Err(Error::Unsupported("block align does not fit in 16 bits"));
        }
        let byte_rate = match spec.spec.sample_rate.checked_mul(block_align) {
            Some(rate) => rate,
            None => return Err(Error::Unsupported("byte rate does not fit in 32 bits")),
        };
        let is_float = spec.spec.sample_format == SampleFormat::Float;
        let extensible = container_bits != bits || channels > 2;

        let mut header = Vec::with_capacity(80);
        header.extend_from_slice(b"RIFF");
        header.extend_from_slice(&0u32.to_le_bytes());
        header.extend_from_slice(b"WAVE");
        header.extend_from_slice(b"fmt ");
        let (tag, fmt_len, fmt_bits) = if extensible {
            (TAG_EXTENSIBLE, 40u32, container_bits)
        } else if is_float {
            (TAG_IEEE_FLOAT, 18u32, bits)
        } else {
            (TAG_PCM, 16u32, bits)
        };
        header.extend_from_slice(&fmt_len.to_le_bytes());
        header.extend_from_slice(&tag.to_le_bytes());
        header.extend_from_slice(&channels.to_le_bytes());
        header.extend_from_slice(&spec.spec.sample_rate.to_le_bytes());
        header.extend_from_slice(&byte_rate.to_le_bytes());
        header.extend_from_slice(&(block_align as u16).to_le_bytes());
        header.extend_from_slice(&fmt_bits.to_le_bytes());
        if extensible {
            header.extend_from_slice(&22u16.to_le_bytes());
            header.extend_from_slice(&bits.to_le_bytes());
            header.extend_from_slice(&0u32.to_le_bytes());
            header.extend_from_slice(if is_float {
                &GUID_IEEE_FLOAT
            } else {
                &GUID_PCM
            });
        } else if is_float {
            header.extend_from_slice(&0u16.to_le_bytes());
        }
        let fact_pos = if is_float {
            header.extend_from_slice(b"fact");
            header.extend_from_slice(&4u32.to_le_bytes());
            let pos = header.len() as u64;
            header.extend_from_slice(&0u32.to_le_bytes());
            Some(pos)
        } else {
            None
        };
        header.extend_from_slice(b"data");
        let data_size_pos = header.len() as u64;
        header.extend_from_slice(&0u32.to_le_bytes());
        let data_start = header.len() as u64;

        // file_len = data_start + data_len + pad, and riff_size = file_len - 8
        // must fit u32. The pad byte is reserved up front.
        let max_data = (u64::from(u32::MAX) + 7 - data_start).min(u64::from(u32::MAX));

        let base = writer.stream_position()?;
        writer.write_all(&header)?;

        Ok(WavWriter {
            inner: writer,
            spec,
            codec,
            block: Vec::with_capacity(BLOCK),
            data_len: 0,
            max_data,
            base,
            fact_pos,
            data_size_pos,
            data_start,
            finalized: false,
        })
    }

    /// Encode one sample.
    ///
    /// Integer types write to integer files and float types to float
    /// files. A mismatch, an `f64` into a 32-bit float file, or an integer
    /// value outside the file's range is [`Error::TooWide`].
    #[inline]
    pub fn write_sample<S: Sample>(&mut self, sample: S) -> Result<()> {
        self.check_write::<S>()?;
        let width = self.codec.width() as u64;
        if self.data_len + width > self.max_data {
            return Err(Error::TooLarge);
        }
        if self.block.len() + self.codec.width() > BLOCK {
            self.flush_block()?;
        }
        <S as crate::sample::SampleOps>::encode(sample, &self.codec, &mut self.block)?;
        self.data_len += width;
        Ok(())
    }

    /// Encode a slice of samples.
    ///
    /// Type and size checks happen before any sample is written. A
    /// per-value range failure leaves the preceding samples written.
    pub fn write_samples<S: Sample>(&mut self, samples: &[S]) -> Result<()> {
        self.check_write::<S>()?;
        let width = self.codec.width();
        let total = samples.len() as u64 * width as u64;
        if self.data_len + total > self.max_data {
            return Err(Error::TooLarge);
        }
        for &sample in samples {
            if self.block.len() + width > BLOCK {
                self.flush_block()?;
            }
            <S as crate::sample::SampleOps>::encode(sample, &self.codec, &mut self.block)?;
            self.data_len += width as u64;
        }
        Ok(())
    }

    /// Push buffered samples to the sink and flush it. Sizes in the header
    /// stay as they are. [`finalize`](Self::finalize) updates them.
    pub fn flush(&mut self) -> Result<()> {
        self.flush_block()?;
        self.inner.flush()?;
        Ok(())
    }

    /// Write the pad byte if needed, patch the header sizes, and flush.
    ///
    /// Returns [`Error::UnfinishedFrame`] when the sample count is not a
    /// multiple of the channel count. The header is patched to a whole
    /// number of frames first so the container stays readable. Drop runs
    /// these same steps and discards the result, so Drop cannot report
    /// [`Error::UnfinishedFrame`]. Call this method to observe that error.
    /// Sample bytes written before the incomplete final frame may remain in
    /// the sink past the declared RIFF boundary. The declared sizes still
    /// bound a reader, and a physical truncation is not possible over a
    /// generic `Write + Seek` sink.
    pub fn finalize(mut self) -> Result<()> {
        self.finalize_inner()
    }

    /// Frames (samples per channel) written so far.
    pub fn duration(&self) -> u32 {
        self.len() / u32::from(self.spec.spec.channels)
    }

    /// Samples written so far, across all channels.
    pub fn len(&self) -> u32 {
        (self.data_len / self.codec.width() as u64) as u32
    }

    /// Whether no sample has been written yet.
    pub fn is_empty(&self) -> bool {
        self.data_len == 0
    }

    /// The stream parameters.
    pub fn spec(&self) -> WavSpec {
        self.spec.spec
    }

    /// The stream parameters plus the container width.
    pub fn spec_ex(&self) -> WavSpecEx {
        self.spec
    }

    fn check_write<S: Sample>(&self) -> Result<()> {
        if <S as crate::sample::SampleOps>::accepts_write(&self.codec) {
            Ok(())
        } else {
            Err(Error::TooWide)
        }
    }

    fn flush_block(&mut self) -> Result<()> {
        if !self.block.is_empty() {
            self.inner.write_all(&self.block)?;
            self.block.clear();
        }
        Ok(())
    }

    fn finalize_inner(&mut self) -> Result<()> {
        self.finalized = true;
        self.flush_block()?;
        let block_align = u64::from(self.block_align());
        let unfinished = self.data_len % block_align != 0;
        if unfinished {
            self.data_len -= self.data_len % block_align;
            self.inner
                .seek(SeekFrom::Start(self.base + self.data_start + self.data_len))?;
        }
        let pad = self.data_len & 1;
        if pad == 1 {
            self.inner.write_all(&[0])?;
        }
        let file_len = self.data_start + self.data_len + pad;
        let riff_size = (file_len - 8) as u32;
        let data_size = self.data_len as u32;
        self.patch(4, riff_size)?;
        self.patch(self.data_size_pos, data_size)?;
        if let Some(pos) = self.fact_pos {
            let frames = (self.data_len / block_align) as u32;
            self.patch(pos, frames)?;
        }
        self.inner.seek(SeekFrom::Start(self.base + file_len))?;
        self.inner.flush()?;
        if unfinished {
            return Err(Error::UnfinishedFrame);
        }
        Ok(())
    }

    fn patch(&mut self, offset: u64, value: u32) -> io::Result<()> {
        self.inner.seek(SeekFrom::Start(self.base + offset))?;
        self.inner.write_all(&value.to_le_bytes())
    }

    fn block_align(&self) -> u32 {
        u32::from(self.spec.spec.channels) * u32::from(self.spec.bytes_per_sample)
    }
}

impl<W: Write + Seek> Drop for WavWriter<W> {
    fn drop(&mut self) {
        if !self.finalized {
            let _ = self.finalize_inner();
        }
    }
}

impl<W: Write + Seek> fmt::Debug for WavWriter<W> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("WavWriter")
            .field("spec", &self.spec)
            .field("data_len", &self.data_len)
            .field("finalized", &self.finalized)
            .finish_non_exhaustive()
    }
}

fn validate(spec: &WavSpecEx) -> Result<Codec> {
    let s = spec.spec;
    let container = spec.bytes_per_sample;
    if s.channels == 0 {
        return Err(Error::Format("zero channels"));
    }
    if s.sample_rate == 0 {
        return Err(Error::Format("zero sample rate"));
    }
    if s.bits_per_sample == 0 {
        return Err(Error::Format("zero bits per sample"));
    }
    if container == 0 {
        return Err(Error::Format("zero bytes per sample"));
    }
    match s.sample_format {
        SampleFormat::Int => {
            if s.bits_per_sample > 32 {
                return Err(Error::TooWide);
            }
            if container > 4 {
                return Err(Error::Unsupported("PCM containers wider than 4 bytes"));
            }
            if s.bits_per_sample > container * 8 {
                return Err(Error::Format("bits per sample exceed the container width"));
            }
            if container == 1 && s.bits_per_sample != 8 {
                return Err(Error::Format("8-bit PCM must declare 8 valid bits"));
            }
        }
        SampleFormat::Float => {
            let ok = (s.bits_per_sample == 32 && container == 4)
                || (s.bits_per_sample == 64 && container == 8);
            if !ok {
                return Err(Error::Unsupported(
                    "float samples other than 32-bit and 64-bit",
                ));
            }
        }
    }
    Ok(Codec::new(s.sample_format, s.bits_per_sample, container))
}

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

    fn spec16() -> WavSpec {
        WavSpec {
            channels: 1,
            sample_rate: 8000,
            bits_per_sample: 16,
            sample_format: SampleFormat::Int,
        }
    }

    #[test]
    fn too_large_fires_before_the_size_fields_wrap() {
        let mut w = WavWriter::new(Cursor::new(Vec::new()), spec16()).unwrap();
        // riff_size = data_start + data_len + pad - 8 must fit u32.
        assert_eq!(w.max_data, u64::from(u32::MAX) + 7 - 44);
        w.data_len = w.max_data - 2;
        w.write_sample(1i16).unwrap();
        assert!(matches!(w.write_sample(1i16), Err(Error::TooLarge)));
        assert!(matches!(w.write_samples(&[1i16]), Err(Error::TooLarge)));
        assert!(matches!(w.write_samples::<i16>(&[]), Ok(())));
        // The counter was forged, so keep Drop from patching sizes with it.
        w.finalized = true;
    }

    #[test]
    fn max_data_accounts_for_the_fact_chunk() {
        let spec = WavSpec {
            sample_format: SampleFormat::Float,
            bits_per_sample: 32,
            ..spec16()
        };
        let w = WavWriter::new(Cursor::new(Vec::new()), spec).unwrap();
        assert_eq!(w.data_start, 58);
        assert_eq!(w.max_data, u64::from(u32::MAX) + 7 - 58);
    }
}