sharkyflac 0.1.0

use std::io::{self, Read, Seek};

use crate::Parse;
use crate::bit_reader::BitReader;

#[derive(Debug, thiserror::Error)]
pub enum Error {
    #[error("IO: {0}")]
    Io(#[from] io::Error),
    #[error("Effective bit depth cannot be 0")]
    BitDepthZero,
    #[error("A unary encoded value was too long")]
    OverlongUnary,
    #[error("Reserved leading bit was set in subframe header")]
    ReservedBit,
    #[error("Reserved subframe type")]
    ReservedType,
    #[error("Reserved Rice coding method")]
    ReservedRiceCoding,
    #[error("Reserved QLP precision (0b1111)")]
    ReservedQlpPrecision,
    #[error("Invalid predictor order")]
    InvalidPredictorOrder,
}

/// order 1: constant velocity
/// order 2: constant acceleration
/// order 3: constant jerk
/// order 4: constant snap
#[rustfmt::skip]
const FIXED_COEFFS: [&[i32]; 5] = [
    &[],
    &[1],
    &[2, -1],
    &[3, -3, 1],
    &[4, -6, 4, -1],
];

fn read_unary<R: Read>(r: &mut BitReader<R>) -> Result<u64, Error> {
    let mut n = 0u64;
    while !r.read_bit()? {
        n = n.checked_add(1).ok_or(Error::OverlongUnary)?;
    }
    Ok(n)
}

/// Uninterleave a [zigzag](https://en.wikipedia.org/wiki/Variable-length_quantity#Zigzag_encoding) coded number.
#[inline(always)]
fn decode_zigzag(n: u64) -> i64 {
    if n & 1 == 0 {
        (n >> 1) as i64
    } else {
        -((n >> 1) as i64) - 1
    }
}

/// The Rice coding method
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum RiceMethod {
    /// Partitioned Rice code with 4-bit parameters
    P4,
    /// Partitioned Rice code with 5-bit parameters
    P5,
}

impl RiceMethod {
    /// The number of parameter bits.
    const fn bits(self) -> u8 {
        match self {
            Self::P4 => 4,
            Self::P5 => 5,
        }
    }

    const fn sentinel(self) -> u8 {
        match self {
            Self::P4 => 0b1111,
            Self::P5 => 0b1_1111,
        }
    }
}

impl Parse<()> for RiceMethod {
    type Error = Error;

    fn parse_from_reader<R: Read + io::Seek>(
        reader: &mut BitReader<R>,
        _opt: (),
    ) -> Result<Self, Self::Error> {
        Ok(match reader.read_bits(2)? {
            0b00 => Self::P4,
            0b01 => Self::P5,
            0b10..=0b11 => return Err(Error::ReservedRiceCoding),
            _ => unreachable!("Two bits can't encode any more information"),
        })
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Type {
    /// In a constant subframe, only a single sample is stored. This sample is
    /// stored as an integer number coded big-endian, signed two's complement.
    ///
    /// See <https://www.ietf.org/rfc/rfc9639.html#name-constant-subframe>
    Constant,
    /// A verbatim subframe stores all samples unencoded in sequential order.
    /// See <https://www.ietf.org/rfc/rfc9639.html#name-verbatim-subframe>
    Verbatim,
    FixedPredictor(u8),
    LinearPredictor(u8),
}

impl Parse<()> for Type {
    type Error = Error;

    fn parse_from_reader<R: Read + Seek>(
        reader: &mut BitReader<R>,
        _opt: (),
    ) -> Result<Self, Self::Error> {
        match reader.read_bits(6)? as u8 {
            0b00_0000 => Ok(Self::Constant),
            0b00_0001 => Ok(Self::Verbatim),
            v @ 0b00_1000..=0b00_1100 => {
                let v = v - 8;
                if v > 4 {
                    Err(Error::InvalidPredictorOrder)
                } else {
                    Ok(Self::FixedPredictor(v))
                }
            }
            v @ 0b10_0000..=0b11_1111 => {
                let v = v - 31;
                if !(1..=32).contains(&v) {
                    Err(Error::InvalidPredictorOrder)
                } else {
                    Ok(Self::LinearPredictor(v))
                }
            }
            0b00_0010..=0b00_0111 | 0b00_1101..=0b01_1111 => Err(Error::ReservedType),

            _ => {
                unreachable!("Six bits can not encode any more information")
            }
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Header {
    kind:        Type,
    /// See <https://www.ietf.org/rfc/rfc9639.html#section-9.2.2>
    wasted_bits: u8,
}

impl Header {
    #[inline]
    pub const fn effective_bit_depth(self, frame_bit_depth: u8, is_side_channel: bool) -> u8 {
        (frame_bit_depth + is_side_channel as u8).saturating_sub(self.wasted_bits)
    }
}

impl Parse<()> for Header {
    type Error = Error;

    fn parse_from_reader<R: Read + Seek>(
        reader: &mut BitReader<R>,
        _opt: (),
    ) -> Result<Self, Self::Error> {
        if reader.read_bit()? {
            return Err(Error::ReservedBit);
        }

        let kind = Type::parse_from_reader(reader, ())?;

        let wasted_bits = if reader.read_bit()? {
            read_unary(reader)
                .ok()
                .and_then(|x| u8::try_from(x).ok())
                .and_then(|x| x.checked_add(1))
                .ok_or(Error::OverlongUnary)?
        } else {
            0
        };

        Ok(Self { kind, wasted_bits })
    }
}

/// Read and decode residuals into `buf[order..]`. `buf` is expected to contain
/// the warmup samples.
fn decode_residuals<R: Read + Seek>(
    reader: &mut BitReader<R>,
    buf: &mut Vec<i32>,
    predictor_order: usize,
    block_size: usize,
) -> Result<(), Error> {
    // The first two bits in a coded residual indicate which coding method is
    // used. See https://www.ietf.org/rfc/rfc9639.html#table-21
    // and https://www.ietf.org/rfc/rfc9639.html#coded-residual

    let rice_method = RiceMethod::parse_from_reader(reader, ())?;
    let partition_order = reader.read_bits(4)? as u8;
    let num_partitions = 1usize << partition_order;
    let parameter_bits = rice_method.bits();
    let sentinel = rice_method.sentinel();

    for partition in 0..num_partitions {
        let n = if partition_order == 0 {
            block_size
                .checked_sub(predictor_order)
                .ok_or(Error::InvalidPredictorOrder)?
        } else if partition == 0 {
            // subtract the number of warm-up samples in the first partition
            (block_size >> partition_order)
                .checked_sub(predictor_order)
                .ok_or(Error::InvalidPredictorOrder)?
        } else {
            block_size >> partition_order
        };

        let rice_param = reader.read_bits(parameter_bits.into())? as u8;

        if rice_param == sentinel {
           // This is an escape partition, so samples are stored verbatim.
           // 
           // > Directly following the escape code are 5 bits containing the number of bits with which each residual sample is stored, as an unsigned number.
           let bit_count = reader.read_bits(5)? as u8;
           for _ in 0..n {
               buf.push(if bit_count == 0 { 0} else {
                   reader.read_signed(bit_count)? as i32
               });
           }
        } else {
            let k = rice_param;
            for _ in 0..n {
                let q= read_unary(reader)?;
                let lo = reader.read_bits(k.into())?;
                buf.push(decode_zigzag(q << k | lo) as i32);
            }
        }
    }

    Ok(())
}

/// Decode samples in-place using a fixed predictor.
fn reconstruct_fixed(buf: &mut [i32], order: usize) {
    let coefficents = FIXED_COEFFS[order];

    // Use a sliding window of length `order`. Multiply each of the last `j`
    // elements by their coeficents. The sum is the sample.
    for i in order..buf.len() {
        let p: i64 = coefficents
            .iter()
            .zip((0..order).map(|j| buf[i - 1 - j]))
            .map(|(&c, s)| c as i64 * s as i64)
            .sum();
        buf[i] = buf[i].wrapping_add(p as i32);
    }
}

#[inline(always)]
fn lpc_dot_product(coefficients: &[i32], history: &[i32]) -> i64 {
    debug_assert_eq!(history.len(), coefficients.len());
    coefficients
        .iter()
        .zip(history.iter().rev())
        .map(|(&c, &s)| c as i64 * s as i64)
        .sum()
}

// TODO: use SIMD for the dot product. Perhaps conditionally with a feature
// gate. fn lpc_dot_product_simd(coefficients: &[i32], history: &[i32]) -> i64

fn reconstruct_lpc(buf: &mut [i32], coefficients: &[i32], shift: i8) {
    let order = coefficients.len();
    for i in order..buf.len() {
        let mut p = lpc_dot_product(coefficients, &buf[i - order..i]);
        if shift.is_positive() {
            p >>= shift;
        } else {
            p <<= -shift;
        }

        buf[i] = buf[i].wrapping_add(p as i32);
    }
}

/// Each sample is stored in its original bit depth, but sign extended. To get
/// samples scaled to i32 range, use [`scale_block`].
pub type Block = Box<[i32]>;

/// Scale a block's samples into `i32` range from the range specified by
/// `bit_depth`.
#[inline]
pub fn scale_block(block: &mut [i32], bit_depth: u8) {
    let shift = 32u32.wrapping_sub(bit_depth as u32);
    for s in block {
        *s = s.wrapping_shl(shift);
    }
}
/// Decode a single subframe block and return its samples.
pub fn decode_block<R: Read + Seek>(
    reader: &mut BitReader<R>,
    block_size: usize,
    frame_bit_depth: u8,
    is_side_channel: bool,
) -> Result<Box<[i32]>, Error> {
    let header = Header::parse_from_reader(reader, ())?;

    let effective_depth = header.effective_bit_depth(frame_bit_depth, is_side_channel);
    if effective_depth == 0 {
        return Err(Error::BitDepthZero);
    }

    let mut buf: Vec<i32> = Vec::with_capacity(block_size);

    match header.kind {
        Type::Constant => {
            let sample = reader.read_signed(effective_depth)?;
            buf.resize(block_size, sample as i32);
        }
        Type::Verbatim => {
            for _ in 0..block_size {
                buf.push(reader.read_signed(effective_depth)? as i32);
            }
        }
        Type::FixedPredictor(order) => {
            let order = order as usize;
            // Read the warm-up samples (the number of warm-up samples is the same as the
            // predictor order).
            for _ in 0..order {
                buf.push(reader.read_signed(effective_depth)? as i32);
            }

            decode_residuals(reader, &mut buf, order, block_size)?;
            reconstruct_fixed(&mut buf, order);
        }
        Type::LinearPredictor(order) => {
            let order = order as usize;

            for _ in 0..order {
                buf.push(reader.read_signed(effective_depth)? as i32);
            }

            let coeff_precision = {
                let prec = reader.read_bits(4)? as u8;
                if prec == 0b1111 {
                    return Err(Error::ReservedQlpPrecision);
                }
                prec + 1
            };

            let shift = reader.read_signed(5)? as i8;

            let coefficients = {
                let mut coeffs = [0i32; 32];
                for c in &mut coeffs[..order] {
                    *c = reader.read_signed(coeff_precision)? as i32;
                }
                coeffs
            };

            decode_residuals(reader, &mut buf, order, block_size)?;
            reconstruct_lpc(&mut buf, &coefficients[0..order], shift);
        }
    }

    debug_assert_eq!(buf.len(), block_size);

    if header.wasted_bits > 0 {
        for s in &mut buf {
            *s <<= header.wasted_bits;
        }
    }

    Ok(buf.into_boxed_slice())
}

/// Iterate over subframes of a single frame. Each `Box<[i32]>` contains samples
/// for one channel, which are produced in stream order.
pub struct BlockIter<'a, R: Read + Seek> {
    reader:          &'a mut BitReader<R>,
    block_size:      usize,
    frame_bit_depth: u8,
    side_flags:      [bool; 8],
    channel_idx:     usize,
    num_channels:    usize,
}

impl<'a, R: Read + Seek> BlockIter<'a, R> {
    /// Create a new [`BlockIter`]. `side_flags` mark whether each
    /// channel channel carries a side-channel. `side_flags.len()` *must* equal
    /// the number of channels in the frame.
    pub fn new(
        reader: &'a mut BitReader<R>,
        block_size: usize,
        frame_bit_depth: u8,
        side_flags: &[bool],
    ) -> Self {
        let n = side_flags.len().min(8);
        let mut flags = [false; 8];
        flags[..n].copy_from_slice(&side_flags[..n]);

        Self {
            reader,
            block_size,
            frame_bit_depth,
            side_flags: flags,
            channel_idx: 0,
            num_channels: n,
        }
    }
}

impl<R: Read + Seek> Iterator for BlockIter<'_, R> {
    type Item = Result<Box<[i32]>, Error>;

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        if self.channel_idx >= self.num_channels {
            return None;
        }
        self.channel_idx += 1;

        Some(decode_block(
            self.reader,
            self.block_size,
            self.frame_bit_depth,
            self.side_flags[self.channel_idx - 1],
        ))
    }

    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        let remaining = self.num_channels - self.channel_idx;
        (remaining, Some(remaining))
    }
}

impl<R: Read + Seek> ExactSizeIterator for BlockIter<'_, R> {}