sharky-arrayvec 0.3.1

use crate::ArrayVec;
use deku::ctx::{Limit, Order, ReadExact};
use deku::prelude::*;
use std::io::prelude::*;
use std::mem;

impl<const N: usize> DekuReader<'_, ReadExact> for ArrayVec<N, u8> {
    fn from_reader_with_ctx<R: Read + Seek>(
        reader: &mut Reader<R>,
        exact: ReadExact,
    ) -> Result<Self, DekuError>
    where
        Self: Sized, {
        assert!(exact.0 <= N);
        let mut bytes = [0x00; N];
        let _ = reader.read_bytes(exact.0, &mut bytes[..exact.0], Order::Lsb0)?;

        Ok(Self::from_slice(&bytes[..exact.0]).unwrap())
    }
}

/// Read `T`s into a vec until a given predicate returns true
/// * `capacity` - an optional capacity to pre-allocate the vector with
/// * `ctx` - The context required by `T`. It will be passed to every `T` when
///   constructing.
/// * `predicate` - the predicate that decides when to stop reading `T`s The
///   predicate takes two parameters: the number of bits that have been read so
///   far, and a borrow of the latest value to have been read. It should return
///   `true` if reading should now stop, and `false` otherwise
fn reader_vec_with_predicate<'a, const N: usize, T, Ctx, Predicate, R: Read + Seek>(
    reader: &mut Reader<R>,
    ctx: Ctx,
    mut predicate: Predicate,
) -> Result<ArrayVec<N, T>, DekuError>
where
    T: DekuReader<'a, Ctx>,
    Ctx: Copy,
    Predicate: FnMut(usize, &T) -> bool, {
    // ZST detected, return empty vec
    if mem::size_of::<T>() == 0 {
        return Ok(ArrayVec::new());
    }

    let mut res = ArrayVec::new();

    let start_read = reader.bits_read;

    loop {
        let val = <T>::from_reader_with_ctx(reader, ctx)?;
        res.push(val);

        // This unwrap is safe as we are pushing to the vec immediately before it,
        // so there will always be a last element
        if predicate(reader.bits_read - start_read, res.last().unwrap()) {
            break;
        }
    }

    Ok(res)
}

fn reader_vec_to_end<'a, const N: usize, T, Ctx, R: Read + Seek>(
    reader: &mut Reader<R>,
    ctx: Ctx,
) -> Result<ArrayVec<N, T>, DekuError>
where
    T: DekuReader<'a, Ctx>,
    Ctx: Copy, {
    // ZST detected, return empty vec
    if mem::size_of::<T>() == 0 {
        return Ok(ArrayVec::new());
    }

    let mut res = ArrayVec::new();
    loop {
        if reader.end() {
            break;
        }
        let val = <T>::from_reader_with_ctx(reader, ctx)?;
        res.push(val);
    }

    Ok(res)
}

impl<'a, const N: usize, T, Ctx, Predicate> DekuReader<'a, (Limit<T, Predicate>, Ctx)>
    for ArrayVec<N, T>
where
    T: DekuReader<'a, Ctx>,
    Ctx: Copy,
    Predicate: FnMut(&T) -> bool,
{
    fn from_reader_with_ctx<R: Read + Seek>(
        reader: &mut Reader<R>,
        (limit, inner_ctx): (Limit<T, Predicate>, Ctx),
    ) -> Result<Self, DekuError>
    where
        Self: Sized, {
        match limit {
            // Read a given count of elements
            Limit::Count(mut count) => {
                assert!(count <= N);

                // Handle the trivial case of reading an empty vector
                if count == 0 {
                    return Ok(ArrayVec::new());
                }

                // Otherwise, read until we have read `count` elements
                reader_vec_with_predicate(reader, inner_ctx, move |_, _| {
                    count -= 1;
                    count == 0
                })
            }

            // Read until a given predicate returns true
            Limit::Until(mut predicate, _) => {
                reader_vec_with_predicate(reader, inner_ctx, move |_, value| predicate(value))
            }

            // Read until a given quantity of bits have been read
            Limit::BitSize(size) => {
                let bit_size = size.0;

                // Handle the trivial case of reading an empty vector
                if bit_size == 0 {
                    return Ok(ArrayVec::new());
                }

                reader_vec_with_predicate(reader, inner_ctx, move |read_bits, _| {
                    read_bits == bit_size
                })
            }

            // Read until a given quantity of bytes have been read
            Limit::ByteSize(size) => {
                let bit_size = size.0 * 8;

                // Handle the trivial case of reading an empty vector
                if bit_size == 0 {
                    return Ok(ArrayVec::new());
                }

                reader_vec_with_predicate(reader, inner_ctx, move |read_bits, _| {
                    read_bits == bit_size
                })
            }

            Limit::End => reader_vec_to_end(reader, inner_ctx),
        }
    }
}

impl<'a, const N: usize, T: DekuReader<'a>, Predicate: FnMut(&T) -> bool>
    DekuReader<'a, Limit<T, Predicate>> for ArrayVec<N, T>
{
    /// Read `T`s until the given limit from input for types which don't require
    /// context.
    fn from_reader_with_ctx<R: Read + Seek>(
        reader: &mut Reader<R>,
        limit: Limit<T, Predicate>,
    ) -> Result<Self, DekuError>
    where
        Self: Sized, {
        ArrayVec::from_reader_with_ctx(reader, (limit, ()))
    }
}

impl<const N: usize, T: DekuWriter<Ctx>, Ctx: Copy> DekuWriter<Ctx> for ArrayVec<N, T> {
    fn to_writer<W: Write + Seek>(
        &self,
        writer: &mut Writer<W>,
        inner_ctx: Ctx,
    ) -> Result<(), DekuError> {
        for v in self.iter() {
            v.to_writer(writer, inner_ctx)?;
        }
        Ok(())
    }
}

#[cfg(test)]
mod tests {

    use super::*;

    #[deku_derive(DekuRead, DekuWrite)]
    #[derive(Debug, Clone, PartialEq, Eq)]
    pub struct Cuesheet {
        #[deku(endian = "big", count = "128")]
        pub media_catalog_number: ArrayVec<128, u8>,
    }
}