mirage-engine 0.1.1

//! Turns one `.ogg` source into a clip and the index its seeks start from.
//! No public type or method names the format.

use std::io::{Cursor, Read, Seek};
use std::sync::Arc;

use lewton::inside_ogg::OggStreamReader;

use crate::Error;
use crate::sound::{Channels, ClipFrame, Encoded, Resampler, SampleRate};

/// The clip `bytes` holds, decoded whole to prove it decodes at all.
///
/// Errors are part of "the asset source `x` …", so it names which source
/// failed.
pub(crate) fn decode(bytes: &[u8]) -> Result<Encoded, Error> {
    let scanned = scan(Cursor::new(bytes))?;
    Ok(Encoded::new(
        Arc::from(bytes),
        scanned.rate,
        scanned.channels,
        scanned.frames,
        scanned.seeks,
    ))
}

/// Every sample of `clip`, interleaved, at `rate`.
pub(crate) fn samples(clip: Arc<Encoded>, rate: SampleRate) -> Result<Vec<f32>, Error> {
    let frames = clip.rate().frames_at(clip.frames(), rate) as usize;
    let mut samples = Vec::with_capacity(frames * clip.channels().count());
    Stream::open(clip, rate)?.read(frames, &mut samples)?;

    Ok(samples)
}

/// One clip decoded as it plays, a packet at a time, at the rate the mix
/// runs at.
///
/// Positions are frames of that rate; the clip's position at its own rate,
/// `source`, is private to the type.
pub(crate) struct Stream {
    reader: OggStreamReader<Cursor<Arc<[u8]>>>,
    clip: Arc<Encoded>,
    width: usize,
    /// The packet being read out, already interleaved, and how many of its
    /// frames are gone.
    packet: Vec<f32>,
    taken: usize,
    /// The clip frame the rest of that packet starts at, at the clip's own
    /// rate.
    source: ClipFrame,
    resampler: Resampler,
    /// Whether the clip's last packet has been decoded.
    spent: bool,
    position: ClipFrame,
    length: ClipFrame,
}

impl Stream {
    pub(crate) fn open(clip: Arc<Encoded>, rate: SampleRate) -> Result<Self, Error> {
        let reader = reader(&clip)?;
        let resampler = Resampler::new(clip.rate(), rate, clip.channels());
        let length = ClipFrame::new(clip.rate().frames_at(clip.frames(), rate));

        Ok(Self {
            width: usize::from(reader.ident_hdr.audio_channels),
            reader,
            clip,
            packet: Vec::new(),
            taken: 0,
            source: ClipFrame::ZERO,
            resampler,
            spent: false,
            position: ClipFrame::ZERO,
            length,
        })
    }

    /// The clip frame that comes out next.
    pub(crate) fn position(&self) -> ClipFrame {
        self.position
    }

    /// Moves to `frame`, starting from the last indexed position before it.
    pub(crate) fn seek(&mut self, frame: ClipFrame) -> Result<(), Error> {
        let from = self.resampler.restart(frame);
        self.packet.clear();
        self.taken = 0;
        self.spent = false;
        self.position = frame;
        self.source = match self.clip.seek_before(from).get() {
            // The clip's first page holds its headers, not audio, so a
            // target the index finds nothing before is read again.
            0 => {
                self.reader = reader(&self.clip)?;
                ClipFrame::ZERO
            }
            raw => {
                self.reader
                    .seek_absgp_pg(raw)
                    .map_err(|error| Error::msg(format!("a sound did not seek: {error}")))?;
                self.landed()?
            }
        };

        self.skip(from)
    }

    /// Appends up to `frames` more frames to `into`, returning how many it
    /// could: fewer only at the end of the clip.
    pub(crate) fn read(
        &mut self,
        frames: usize,
        into: &mut impl Extend<f32>,
    ) -> Result<usize, Error> {
        let frames = frames.min(self.length.saturating_sub(self.position) as usize);
        let mut read = 0;
        while read < frames {
            let took = self.resampler.take(frames - read, into);
            self.position += took as u64;
            read += took;
            if read == frames || self.spent {
                break;
            }
            self.spent = !self.more()?;
        }

        Ok(read)
    }

    /// Position a seek ended up at, which the reader can only report once
    /// it has read a packet off the page it landed on.
    fn landed(&mut self) -> Result<ClipFrame, Error> {
        loop {
            if !self.fill()? {
                return Ok(ClipFrame::new(self.clip.frames()));
            }
            if let Some(landed) = self.reader.get_last_absgp().map(ClipFrame::new) {
                return Ok(landed.back(self.left() as u64));
            }
        }
    }

    /// Drops the clip's frames up to `to`, at the clip's own rate.
    fn skip(&mut self, to: ClipFrame) -> Result<(), Error> {
        while self.source < to {
            if self.left() == 0 && !self.fill()? {
                return Ok(());
            }
            let dropped = (to.saturating_sub(self.source) as usize).min(self.left());
            self.taken += dropped;
            self.source += dropped as u64;
        }

        Ok(())
    }

    /// Feeds the resampler what the clip holds next; false at the end of the
    /// clip.
    fn more(&mut self) -> Result<bool, Error> {
        if self.left() == 0 && !self.fill()? {
            self.resampler.end();
            return Ok(false);
        }

        let channels = self.clip.channels().count();
        self.resampler.feed(&self.packet[self.taken * channels..]);
        self.source += self.left() as u64;
        self.taken = self.packet.len() / channels;

        Ok(true)
    }

    /// Frames of the decoded packet that have not been read out.
    fn left(&self) -> usize {
        self.packet.len() / self.clip.channels().count() - self.taken
    }

    /// Decodes the next packet, or returns false at the end of the clip.
    fn fill(&mut self) -> Result<bool, Error> {
        let packet = self
            .reader
            .read_dec_packet_itl()
            .map_err(|error| Error::msg(format!("a sound did not decode: {error}")))?;
        let Some(packet) = packet else {
            return Ok(false);
        };

        self.packet.clear();
        self.taken = 0;
        interleave(&packet, self.width, self.clip.channels(), &mut self.packet);
        Ok(true)
    }
}

/// Report of decoding a clip whole.
struct Scanned {
    rate: SampleRate,
    channels: Channels,
    frames: u64,
    seeks: Vec<ClipFrame>,
}

/// A reader over the clip's own bytes, at the first of its audio.
fn reader(clip: &Encoded) -> Result<OggStreamReader<Cursor<Arc<[u8]>>>, Error> {
    OggStreamReader::new(Cursor::new(clip.bytes().clone()))
        .map_err(|error| Error::msg(format!("a sound did not decode: {error}")))
}

/// Decodes every packet to compute the clip's shape: its rate, channels and
/// frame count.
fn scan<T: Read + Seek>(source: T) -> Result<Scanned, Error> {
    let mut reader = OggStreamReader::new(source)
        .map_err(|error| Error::msg(format!("did not decode: {error}")))?;
    let rate = reader.ident_hdr.audio_sample_rate;
    let width = usize::from(reader.ident_hdr.audio_channels);
    if rate == 0 || width == 0 {
        return Err(Error::msg("holds a sound of no rate or no channels"));
    }
    let rate = SampleRate::new(rate);

    let mut scanned = Scanned {
        rate,
        channels: Channels::of(reader.ident_hdr.audio_channels),
        frames: 0,
        seeks: Vec::new(),
    };
    loop {
        let packet = reader
            .read_dec_packet_itl()
            .map_err(|error| Error::msg(format!("did not decode: {error}")))?;
        let Some(packet) = packet else {
            break;
        };

        scanned.frames += (packet.len() / width) as u64;

        // The index holds page positions, the only ones a seek can land on.
        if let Some(page) = reader.get_last_absgp().map(ClipFrame::new)
            && scanned
                .seeks
                .last()
                .is_none_or(|&indexed| page.saturating_sub(indexed) >= u64::from(rate))
        {
            scanned.seeks.push(page);
        }
    }

    // An encoder pads its last packet out to a whole block, and the last page
    // states which frame the audio itself ends at.
    if let Some(ends) = reader
        .get_last_absgp()
        .filter(|&ends| ends < scanned.frames)
    {
        scanned.frames = ends;
    }

    Ok(scanned)
}

/// Appends `packet`, which is `width` channels wide, as `channels`: each
/// frame keeps only its first `channels` samples.
fn interleave(packet: &[i16], width: usize, channels: Channels, into: &mut Vec<f32>) {
    into.extend(packet.chunks_exact(width).flat_map(|frame| {
        frame[..channels.count()]
            .iter()
            .map(|&sample| f32::from(sample) / f32::from(i16::MAX))
    }));
}

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

    /// A two-second sweep, the fixture every sound test decodes.
    const SWEEP: &[u8] = include_bytes!("../../tests/assets/sweep.ogg");

    /// A short tone, all of whose audio lies on one page.
    const TONE: &[u8] = include_bytes!("../../tests/assets/tone.ogg");

    /// The fixture `bytes` hold, at its own rate.
    fn clip(bytes: &[u8]) -> Arc<Encoded> {
        Arc::new(decode(bytes).expect("the fixture decodes"))
    }

    #[test]
    fn a_clip_lasts_as_long_as_its_source_states_and_no_longer() {
        let clip = clip(TONE);
        let recorded = 44_100 / 20;

        assert_eq!(clip.frames(), recorded, "a twentieth of a second of it");
        assert_eq!(
            samples(Arc::clone(&clip), clip.rate())
                .expect("it decodes again")
                .len(),
            recorded as usize,
            "and its samples stop where its frames do"
        );

        let mut stream = Stream::open(Arc::clone(&clip), clip.rate()).expect("it opens");
        let mut read = Vec::new();
        assert_eq!(
            stream
                .read(recorded as usize + 1_000, &mut read)
                .expect("it reads"),
            recorded as usize,
            "and a stream stops there as well"
        );
    }

    #[test]
    fn a_source_decodes_to_a_clip_that_knows_its_own_shape() {
        let clip = clip(SWEEP);

        assert_eq!(clip.rate(), SampleRate::new(44_100));
        assert_eq!(clip.channels(), Channels::Mono);
        assert_eq!(clip.frames(), 88_200, "two seconds of it");
        assert_eq!(
            samples(Arc::clone(&clip), clip.rate())
                .expect("it decodes again")
                .len(),
            88_200
        );
    }

    #[test]
    fn the_index_names_a_position_about_every_second() {
        let clip = clip(SWEEP);

        assert_eq!(clip.seek_before(ClipFrame::new(0)).get(), 0);
        assert!(
            clip.seek_before(ClipFrame::new(88_000)).get() >= 44_100,
            "a seek to the end starts past the middle"
        );
    }

    #[test]
    fn a_stream_reads_the_samples_the_whole_decode_does() {
        let clip = clip(SWEEP);
        let whole = samples(Arc::clone(&clip), clip.rate()).expect("it decodes");
        let mut stream = Stream::open(Arc::clone(&clip), clip.rate()).expect("it opens");

        let mut read = Vec::new();
        stream.read(1_000, &mut read).expect("it reads");
        assert_eq!(stream.position().get(), 1_000);
        assert!(
            read.iter().zip(&whole).all(|(from, to)| from == to),
            "a stream and a whole decode agree sample for sample"
        );
    }

    #[test]
    fn a_seek_lands_where_it_was_asked_to() {
        let clip = clip(SWEEP);
        let whole = samples(Arc::clone(&clip), clip.rate()).expect("it decodes");
        let mut stream = Stream::open(Arc::clone(&clip), clip.rate()).expect("it opens");

        stream.seek(ClipFrame::new(60_000)).expect("it seeks");
        assert_eq!(stream.position().get(), 60_000);

        let mut read = Vec::new();
        stream.read(2_000, &mut read).expect("it reads");
        let off = read
            .iter()
            .zip(&whole[60_000..])
            .map(|(from, to)| (from - to).abs())
            .fold(0.0f32, f32::max);
        assert!(off < 0.05, "a seek reads the same sweep, off by {off}");
    }

    /// The `CRC` an `.ogg` page holds, over `page` with its own read as
    /// zero.
    fn checksum(page: &[u8]) -> u32 {
        let mut sum = 0u32;
        for &byte in page {
            sum ^= u32::from(byte) << 24;
            for _ in 0..8 {
                sum = match sum & 0x8000_0000 != 0 {
                    true => (sum << 1) ^ 0x04c1_1db7,
                    false => sum << 1,
                };
            }
        }

        sum
    }

    /// The fixture whose last page ends at `at`, its `CRC` written again so
    /// that a reader still takes it.
    fn ending_at(at: u64) -> Vec<u8> {
        let mut out = SWEEP.to_vec();
        let page = out
            .windows(4)
            .rposition(|window| window == b"OggS")
            .expect("the fixture holds pages");
        out[page + 6..page + 14].copy_from_slice(&at.to_le_bytes());
        out[page + 22..page + 26].copy_from_slice(&[0; 4]);
        let sum = checksum(&out[page..]);
        out[page + 22..page + 26].copy_from_slice(&sum.to_le_bytes());

        out
    }

    /// Frames `bytes` hold, or nothing where they do not decode or a stream
    /// over them does not read.
    fn played(bytes: &[u8]) -> Option<u64> {
        let clip = Arc::new(decode(bytes).ok()?);
        let mut stream = Stream::open(Arc::clone(&clip), clip.rate()).ok()?;
        let mut read = Vec::new();
        stream.seek(ClipFrame::new(1_000)).ok()?;
        stream.read(1_000, &mut read).ok()?;

        Some(clip.frames())
    }

    #[test]
    fn a_source_claiming_a_position_past_what_the_index_counts_in_still_reads() {
        let claimed = ending_at(u64::MAX - 1_000);
        let clip = Arc::new(decode(&claimed).expect("its pages still decode"));

        let mut stream = Stream::open(Arc::clone(&clip), clip.rate()).expect("it opens");
        let mut read = Vec::new();
        stream.seek(ClipFrame::new(60_000)).expect("it seeks");
        stream.read(1_000, &mut read).expect("it reads");

        assert_eq!(read.len(), 1_000);
        assert!(
            clip.seek_before(ClipFrame::new(60_000)).get() <= 60_000,
            "a seek starts at a position the clip reaches, not the one it claims"
        );
    }

    #[test]
    fn a_source_cut_off_or_broken_anywhere_reads_no_further_than_itself() {
        let whole = decode(SWEEP).expect("the fixture decodes").frames();
        assert_eq!(played(SWEEP), Some(whole), "the fixture itself plays");

        for at in (0..SWEEP.len()).step_by(128) {
            if let Some(frames) = played(&SWEEP[..at]) {
                assert!(frames <= whole, "a cut at {at} read past its own end");
            }
        }
        for at in (0..SWEEP.len()).step_by(251) {
            let mut broken = SWEEP.to_vec();
            broken[at] ^= u8::MAX;
            if let Some(frames) = played(&broken) {
                assert!(frames <= whole, "a flip at {at} read past its own end");
            }
        }
    }

    #[test]
    fn a_source_that_is_not_a_sound_fails_to_decode() {
        let error = decode(b"not a sound at all").expect_err("nothing decodes that");

        assert!(error.to_string().starts_with("did not decode"), "{error}");
    }
}