wav-codec 0.1.0

use std::fmt;
use std::fs::File;
use std::io::{self, BufReader, Read, Seek, SeekFrom};
use std::iter::FusedIterator;
use std::marker::PhantomData;
use std::path::Path;

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

/// Size of the block buffer the reader refills in bulk.
const BLOCK: usize = 64 * 1024;

/// Streaming WAV decoder over any `io::Read`.
///
/// Construction parses the RIFF header, the `fmt ` chunk, and every chunk
/// up to and including the `data` chunk header. Samples are then decoded
/// from a 64 KiB block buffer that is refilled with bulk reads clamped to
/// the declared `data` boundary, so the underlying reader is never read
/// past the end of the audio payload.
pub struct WavReader<R> {
    inner: R,
    spec: WavSpecEx,
    codec: Codec,
    /// Declared `data` payload length in bytes.
    data_len: u32,
    /// Payload bytes not yet pulled into the block buffer.
    remaining: u32,
    buf: Vec<u8>,
    pos: usize,
    end: usize,
}

impl WavReader<BufReader<File>> {
    /// Open a WAV file from disk.
    pub fn open<P: AsRef<Path>>(path: P) -> Result<Self> {
        let file = File::open(path)?;
        WavReader::new(BufReader::new(file))
    }
}

impl<R: Read> WavReader<R> {
    /// Parse the headers from `reader` and position it at the first sample.
    ///
    /// Fails on a malformed container ([`Error::Format`]), an encoding
    /// outside the crate's scope ([`Error::Unsupported`]), or a stream that
    /// ends inside the headers ([`Error::Io`]).
    pub fn new(mut reader: R) -> Result<Self> {
        let mut riff = [0u8; 12];
        reader.read_exact(&mut riff)?;
        match &riff[0..4] {
            b"RIFF" => {}
            b"RF64" | b"BW64" => return Err(Error::Unsupported("RF64 containers (over 4 GiB)")),
            b"RIFX" => return Err(Error::Unsupported("RIFX big-endian containers")),
            _ => return Err(Error::Format("not a RIFF file")),
        }
        if &riff[8..12] != b"WAVE" {
            return Err(Error::Format("RIFF form type is not WAVE"));
        }

        let mut format: Option<(WavSpecEx, Codec)> = None;
        let (spec, codec, data_len) = loop {
            let mut header = [0u8; 8];
            read_chunk_header(&mut reader, &mut header)?;
            if !header[0..4].iter().all(|b| (0x20..=0x7E).contains(b)) {
                return Err(Error::Format("chunk id is not printable ASCII"));
            }
            let size = u32::from_le_bytes([header[4], header[5], header[6], header[7]]);
            match &header[0..4] {
                b"fmt " => {
                    if format.is_some() {
                        return Err(Error::Format("duplicate fmt chunk"));
                    }
                    format = Some(parse_fmt(&mut reader, size)?);
                }
                b"data" => {
                    let (spec, codec) = match format {
                        Some(f) => f,
                        None => return Err(Error::Format("data chunk before fmt chunk")),
                    };
                    let block_align =
                        u32::from(spec.spec.channels) * u32::from(spec.bytes_per_sample);
                    if size % block_align != 0 {
                        return Err(Error::Format(
                            "data length is not a multiple of the frame size",
                        ));
                    }
                    break (spec, codec, size);
                }
                _ => skip_chunk(&mut reader, size)?,
            }
        };

        Ok(WavReader {
            inner: reader,
            spec,
            codec,
            data_len,
            remaining: data_len,
            buf: vec![0u8; BLOCK],
            pos: 0,
            end: 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
    }

    /// Number of frames (samples per channel) the `data` chunk declares.
    pub fn duration(&self) -> u32 {
        self.data_len / self.block_align()
    }

    /// Total number of samples across all channels the `data` chunk declares.
    ///
    /// For a stream the declared count can exceed what exists. The sample
    /// iterators report that shortfall as an error.
    pub fn len(&self) -> u32 {
        self.data_len / u32::from(self.codec.container)
    }

    /// Whether the `data` chunk declares zero samples.
    pub fn is_empty(&self) -> bool {
        self.data_len == 0
    }

    /// Iterate over the remaining samples, borrowing the reader.
    ///
    /// Items are `Result<S>`. The first error fuses the iterator. A type
    /// that cannot represent the file's samples yields one
    /// [`Error::TooWide`].
    pub fn samples<S: Sample>(&mut self) -> Samples<'_, R, S> {
        let pending = self.check_read::<S>().err();
        Samples {
            reader: self,
            pending,
            done: false,
            _sample: PhantomData,
        }
    }

    /// Iterate over the remaining samples, consuming the reader.
    pub fn into_samples<S: Sample>(self) -> IntoSamples<R, S> {
        let pending = self.check_read::<S>().err();
        IntoSamples {
            reader: self,
            pending,
            done: false,
            _sample: PhantomData,
        }
    }

    /// Decode samples into `buf` and return how many were written.
    ///
    /// The buffer is filled completely unless the `data` chunk ends first,
    /// so a count below `buf.len()` means the end of the audio, and `Ok(0)`
    /// on a later call confirms it. A stream that ends before the declared
    /// `data` size returns `Err` instead of a short count.
    pub fn read_samples<S: Sample>(&mut self, buf: &mut [S]) -> Result<usize> {
        self.check_read::<S>()?;
        let width = self.codec.width();
        let mut filled = 0;
        while filled < buf.len() {
            let available = (self.end - self.pos) / width;
            if available == 0 {
                if self.refill()? {
                    continue;
                }
                break;
            }
            let n = available.min(buf.len() - filled);
            let bytes = n * width;
            decode_block(
                &self.codec,
                &self.buf[self.pos..self.pos + bytes],
                &mut buf[filled..filled + n],
            );
            self.pos += bytes;
            filled += n;
        }
        Ok(filled)
    }

    /// Return the underlying reader.
    ///
    /// Bytes pulled into the block buffer but not yet yielded are lost.
    pub fn into_inner(self) -> R {
        self.inner
    }

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

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

    /// Samples declared but not yet yielded.
    fn samples_remaining(&self) -> usize {
        let buffered = self.end - self.pos;
        let declared = self.remaining as usize;
        (buffered + declared) / self.codec.width()
    }

    /// Decode the next sample, or `None` at the declared end of the data.
    #[inline]
    fn next_sample<S: Sample>(&mut self) -> Option<Result<S>> {
        let width = self.codec.width();
        if self.pos + width > self.end {
            match self.refill() {
                Ok(true) => {}
                Ok(false) => return None,
                Err(e) => return Some(Err(e)),
            }
        }
        let bytes = &self.buf[self.pos..self.pos + width];
        self.pos += width;
        Some(Ok(<S as crate::sample::SampleOps>::decode(
            bytes,
            &self.codec,
        )))
    }

    /// Pull more payload into the block buffer until at least one whole
    /// sample is available.
    ///
    /// Returns `Ok(true)` when a sample can be decoded, `Ok(false)` at the
    /// clean end of the declared payload, and an error when the stream ends
    /// short of the declared size or the reader fails.
    #[inline(never)]
    fn refill(&mut self) -> Result<bool> {
        let width = self.codec.width();
        let leftover = self.end - self.pos;
        if leftover > 0 && self.pos > 0 {
            self.buf.copy_within(self.pos..self.end, 0);
        }
        self.pos = 0;
        self.end = leftover;
        while self.end < width {
            let room = self.buf.len() - self.end;
            let want = room.min(self.remaining as usize);
            if want == 0 {
                break;
            }
            match self.inner.read(&mut self.buf[self.end..self.end + want]) {
                Ok(0) => {
                    return Err(Error::Format(
                        "stream ends before the declared end of the data chunk",
                    ))
                }
                Ok(n) => {
                    let n = n.min(want);
                    self.end += n;
                    self.remaining -= n as u32;
                }
                Err(e) if e.kind() == io::ErrorKind::Interrupted => {}
                Err(e) => return Err(Error::Io(e)),
            }
        }
        if self.end >= width {
            Ok(true)
        } else if self.end == 0 {
            Ok(false)
        } else {
            Err(Error::Format(
                "stream ends before the declared end of the data chunk",
            ))
        }
    }
}

impl<R: Read + Seek> WavReader<R> {
    /// Position the reader at `frame` (a frame is one sample per channel).
    ///
    /// Seeking advances by the container width, so a 24-in-32 file seeks in
    /// 4-byte steps. `frame` may equal [`duration`](Self::duration), which
    /// leaves no samples to read. Anything past that is `InvalidInput`.
    pub fn seek(&mut self, frame: u32) -> io::Result<()> {
        if frame > self.duration() {
            return Err(io::Error::new(
                io::ErrorKind::InvalidInput,
                "seek past the end of the data chunk",
            ));
        }
        let target = u64::from(frame) * u64::from(self.block_align());
        let consumed = u64::from(self.data_len - self.remaining);
        let current = self.inner.stream_position()?;
        let data_start = current - consumed;
        self.inner.seek(SeekFrom::Start(data_start + target))?;
        self.pos = 0;
        self.end = 0;
        self.remaining = self.data_len - target as u32;
        Ok(())
    }
}

/// Read a chunk header, distinguishing a clean end of stream at a chunk
/// boundary (no `data` chunk at all) from a header cut short.
fn read_chunk_header<R: Read>(reader: &mut R, header: &mut [u8; 8]) -> Result<()> {
    let mut got = 0;
    while got < header.len() {
        match reader.read(&mut header[got..]) {
            Ok(0) if got == 0 => return Err(Error::Format("no data chunk")),
            Ok(0) => {
                return Err(Error::Io(io::Error::new(
                    io::ErrorKind::UnexpectedEof,
                    "stream ends inside a chunk header",
                )))
            }
            Ok(n) => got += n,
            Err(e) if e.kind() == io::ErrorKind::Interrupted => {}
            Err(e) => return Err(Error::Io(e)),
        }
    }
    Ok(())
}

fn skip<R: Read>(reader: &mut R, bytes: u64) -> Result<()> {
    let copied = io::copy(&mut reader.by_ref().take(bytes), &mut io::sink())?;
    if copied != bytes {
        return Err(Error::Io(io::Error::new(
            io::ErrorKind::UnexpectedEof,
            "stream ends inside a chunk",
        )));
    }
    Ok(())
}

/// Skip a chunk payload, then require the RIFF pad byte to be zero when
/// `size` is odd.
fn skip_chunk<R: Read>(reader: &mut R, size: u32) -> Result<()> {
    skip(reader, u64::from(size))?;
    if size & 1 == 1 {
        skip_odd_pad(reader)?;
    }
    Ok(())
}

fn skip_odd_pad<R: Read>(reader: &mut R) -> Result<()> {
    let mut pad = [0u8; 1];
    loop {
        match reader.read(&mut pad) {
            Ok(0) => return Err(Error::Format("missing pad byte after odd-sized chunk")),
            Ok(_) if pad[0] == 0 => return Ok(()),
            Ok(_) => return Err(Error::Format("nonzero pad byte after odd-sized chunk")),
            Err(e) if e.kind() == io::ErrorKind::Interrupted => {}
            Err(e) => return Err(Error::Io(e)),
        }
    }
}

fn parse_fmt<R: Read>(reader: &mut R, size: u32) -> Result<(WavSpecEx, Codec)> {
    if size < 16 {
        return Err(Error::Format("fmt chunk too small"));
    }
    let mut head = [0u8; 16];
    reader.read_exact(&mut head)?;
    let mut consumed: u64 = 16;

    let tag = u16::from_le_bytes([head[0], head[1]]);
    let channels = u16::from_le_bytes([head[2], head[3]]);
    let sample_rate = u32::from_le_bytes([head[4], head[5], head[6], head[7]]);
    // head[8..12] is byte_rate. Recalculate it from the other fields.
    let block_align = u16::from_le_bytes([head[12], head[13]]);
    let bits_per_sample = u16::from_le_bytes([head[14], head[15]]);

    let mut extension: Option<[u8; 22]> = None;
    if size >= 18 {
        let mut cb = [0u8; 2];
        reader.read_exact(&mut cb)?;
        consumed += 2;
        let cb_size = u16::from_le_bytes(cb);
        if tag == TAG_EXTENSIBLE {
            if size < 40 || cb_size < 22 {
                return Err(Error::Format(
                    "fmt chunk too small for WAVE_FORMAT_EXTENSIBLE",
                ));
            }
            let mut ext = [0u8; 22];
            reader.read_exact(&mut ext)?;
            consumed += 22;
            extension = Some(ext);
        }
    } else if tag == TAG_EXTENSIBLE {
        return Err(Error::Format(
            "fmt chunk too small for WAVE_FORMAT_EXTENSIBLE",
        ));
    }
    skip(reader, u64::from(size) - consumed)?;
    if size & 1 == 1 {
        skip_odd_pad(reader)?;
    }

    if channels == 0 {
        return Err(Error::Format("zero channels"));
    }
    if sample_rate == 0 {
        return Err(Error::Format("zero sample rate"));
    }
    if block_align == 0 || block_align % channels != 0 {
        return Err(Error::Format(
            "block align is not a nonzero multiple of the channel count",
        ));
    }
    let container = block_align / channels;
    // u32: a hostile block_align of 0xFF01 over one channel overflows u16 here.
    let container_bits = u32::from(container) * 8;

    // The WAVEFORMATEX bits_per_sample field is always bound by the
    // container, including WAVE_FORMAT_EXTENSIBLE when valid_bits is set.
    if u32::from(bits_per_sample) > container_bits {
        return Err(Error::Format("bits per sample exceed the container width"));
    }

    let (format, bits) = match tag {
        TAG_PCM => (SampleFormat::Int, bits_per_sample),
        TAG_IEEE_FLOAT => (SampleFormat::Float, bits_per_sample),
        TAG_EXTENSIBLE => {
            let ext = match extension {
                Some(e) => e,
                None => {
                    return Err(Error::Format(
                        "fmt chunk too small for WAVE_FORMAT_EXTENSIBLE",
                    ))
                }
            };
            let valid_bits = u16::from_le_bytes([ext[0], ext[1]]);
            // ext[2..6] is the channel mask, carried but not interpreted.
            let guid = &ext[6..22];
            let format = match (guid[0], guid[1], guid[2..16] == GUID_PCM[2..]) {
                (0x01, 0x00, true) => SampleFormat::Int,
                (0x03, 0x00, true) => SampleFormat::Float,
                _ => return Err(Error::Unsupported("sub-format GUID")),
            };
            let bits = if valid_bits == 0 {
                bits_per_sample
            } else {
                valid_bits
            };
            (format, bits)
        }
        _ => return Err(Error::Unsupported("compressed or unknown format tag")),
    };

    if bits == 0 {
        return Err(Error::Format("zero bits per sample"));
    }
    if u32::from(bits) > container_bits {
        return Err(Error::Format("bits per sample exceed the container width"));
    }
    match format {
        SampleFormat::Int => {
            if container > 4 {
                return Err(Error::Unsupported("PCM samples wider than 32 bits"));
            }
            // One-byte samples are offset binary, which only has an 8-bit form.
            if container == 1 && bits != 8 {
                return Err(Error::Format("8-bit PCM must declare 8 valid bits"));
            }
        }
        SampleFormat::Float => {
            if !((bits == 32 && container == 4) || (bits == 64 && container == 8)) {
                return Err(Error::Unsupported(
                    "float samples other than 32-bit and 64-bit",
                ));
            }
        }
    }

    let spec = WavSpecEx {
        spec: WavSpec {
            channels,
            sample_rate,
            bits_per_sample: bits,
            sample_format: format,
        },
        bytes_per_sample: container,
    };
    Ok((spec, Codec::new(format, bits, container)))
}

/// Borrowing sample iterator from [`WavReader::samples`].
pub struct Samples<'a, R, S> {
    reader: &'a mut WavReader<R>,
    pending: Option<Error>,
    done: bool,
    _sample: PhantomData<S>,
}

impl<'a, R: Read, S: Sample> Iterator for Samples<'a, R, S> {
    type Item = Result<S>;

    #[inline]
    fn next(&mut self) -> Option<Result<S>> {
        if self.done {
            return None;
        }
        if let Some(e) = self.pending.take() {
            self.done = true;
            return Some(Err(e));
        }
        let item = self.reader.next_sample::<S>();
        if let Some(Err(_)) = item {
            self.done = true;
        }
        item
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        if self.done {
            (0, Some(0))
        } else {
            (0, Some(self.reader.samples_remaining()))
        }
    }
}

impl<'a, R: Read, S: Sample> FusedIterator for Samples<'a, R, S> {}

/// Owning sample iterator from [`WavReader::into_samples`].
pub struct IntoSamples<R, S> {
    reader: WavReader<R>,
    pending: Option<Error>,
    done: bool,
    _sample: PhantomData<S>,
}

impl<R, S> IntoSamples<R, S> {
    /// Return the reader, positioned after the last sample yielded.
    pub fn into_inner(self) -> WavReader<R> {
        self.reader
    }
}

impl<R: Read, S: Sample> Iterator for IntoSamples<R, S> {
    type Item = Result<S>;

    #[inline]
    fn next(&mut self) -> Option<Result<S>> {
        if self.done {
            return None;
        }
        if let Some(e) = self.pending.take() {
            self.done = true;
            return Some(Err(e));
        }
        let item = self.reader.next_sample::<S>();
        if let Some(Err(_)) = item {
            self.done = true;
        }
        item
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        if self.done {
            (0, Some(0))
        } else {
            (0, Some(self.reader.samples_remaining()))
        }
    }
}

impl<R: Read, S: Sample> FusedIterator for IntoSamples<R, S> {}

impl<R> fmt::Debug for WavReader<R> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("WavReader")
            .field("spec", &self.spec)
            .field("data_len", &self.data_len)
            .field("remaining", &self.remaining)
            .finish_non_exhaustive()
    }
}

impl<R, S> fmt::Debug for Samples<'_, R, S> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("Samples")
            .field("pending", &self.pending.is_some())
            .field("done", &self.done)
            .finish_non_exhaustive()
    }
}

impl<R, S> fmt::Debug for IntoSamples<R, S> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("IntoSamples")
            .field("pending", &self.pending.is_some())
            .field("done", &self.done)
            .finish_non_exhaustive()
    }
}