messagepack-core 0.2.4

//! Array decoding helpers.

use core::marker::PhantomData;

use super::{DecodeBorrowed, Error, NbyteReader};
use crate::{formats::Format, io::IoRead};

/// Decode a MessagePack array of `V` into `Array` collecting iterator.
pub struct ArrayDecoder<Array, V>(PhantomData<(Array, V)>);

impl<'de, Array, V> DecodeBorrowed<'de> for ArrayDecoder<Array, V>
where
    V: DecodeBorrowed<'de>,
    Array: FromIterator<V::Value>,
{
    type Value = Array;

    fn decode_borrowed_with_format<R>(
        format: Format,
        reader: &mut R,
    ) -> core::result::Result<Self::Value, Error<R::Error>>
    where
        R: IoRead<'de>,
    {
        let len = match format {
            Format::FixArray(len) => len.into(),
            Format::Array16 => NbyteReader::<2>::read(reader)?,
            Format::Array32 => NbyteReader::<4>::read(reader)?,
            _ => return Err(Error::UnexpectedFormat),
        };

        let out = (0..len)
            .map(|_| V::decode_borrowed(reader))
            .collect::<core::result::Result<Array, Error<R::Error>>>()?;
        Ok(out)
    }
}

impl<'de, const N: usize, V> DecodeBorrowed<'de> for [V; N]
where
    V: DecodeBorrowed<'de>,
{
    type Value = [V::Value; N];

    fn decode_borrowed_with_format<R>(
        format: Format,
        reader: &mut R,
    ) -> core::result::Result<Self::Value, Error<R::Error>>
    where
        R: IoRead<'de>,
    {
        let len = match format {
            Format::FixArray(len) => len.into(),
            Format::Array16 => NbyteReader::<2>::read(reader)?,
            Format::Array32 => NbyteReader::<4>::read(reader)?,
            _ => return Err(Error::UnexpectedFormat),
        };
        if len != N {
            return Err(Error::InvalidData);
        };

        let mut tmp: [Option<V::Value>; N] = core::array::from_fn(|_| None);
        for item in tmp.iter_mut() {
            *item = Some(V::decode_borrowed(reader)?);
        }
        // NOTE: This `expect` cannot fire given the invariant established above.
        // - We allocate a temporary `[Option<V::Value>; N]` initialized to `None`.
        // - The loop assigns `Some(..)` to every element or returns early on error.
        // - Therefore, reaching here implies all entries are `Some`, making the
        //   unwrap logically safe.
        //
        // Why not avoid `expect` entirely?
        // - On stable Rust, there is no fully-ergonomic way to construct
        //   `[T; N]` from a fallible producer without either allocation or
        //   `unsafe`.
        // - `core::array::try_from_fn` would be a perfect fit but is not yet
        //   stabilized at the time of writing.
        // - Collecting into `Vec<T>` and converting with `try_into::<[T; N]>()`
        //   requires the `alloc` feature, while this crate targets `no_std` by
        //   default.
        //
        // If/when `try_from_fn` is stabilized, or if an `alloc`-enabled fast
        // path is desired, this code can be revisited to remove `expect`.
        let out = core::array::from_fn(|i| tmp[i].take().expect("initialized"));
        Ok(out)
    }
}

macro_rules! tuple_decode_impls {
    ($($len:expr => ($($name:ident)+))+ $(,)?) => {
        $(
            impl<'de, $($name),+> DecodeBorrowed<'de> for ($($name,)+)
            where
                $($name: DecodeBorrowed<'de>,)+
            {
                type Value = ($(<$name as DecodeBorrowed<'de>>::Value,)+);

                fn decode_borrowed_with_format<R>(format: Format, reader: &mut R) -> core::result::Result<Self::Value, Error<R::Error>>
                where
                    R: IoRead<'de>,
                {
                    let len = match format {
                        Format::FixArray(len) => len.into(),
                        Format::Array16 => NbyteReader::<2>::read(reader)?,
                        Format::Array32 => NbyteReader::<4>::read(reader)?,
                        _ => return Err(Error::UnexpectedFormat),
                    };
                    if len != $len {
                        return Err(Error::InvalidData);
                    }

                    let value = (
                        $({
                            let v = <$name as DecodeBorrowed<'de>>::decode_borrowed(reader)?;
                            v
                        },)+
                    );
                    Ok(value)
                }
            }
        )+
    };
}

tuple_decode_impls! {
    1  => (V0)
    2  => (V0 V1)
    3  => (V0 V1 V2)
    4  => (V0 V1 V2 V3)
    5  => (V0 V1 V2 V3 V4)
    6  => (V0 V1 V2 V3 V4 V5)
    7  => (V0 V1 V2 V3 V4 V5 V6)
    8  => (V0 V1 V2 V3 V4 V5 V6 V7)
    9  => (V0 V1 V2 V3 V4 V5 V6 V7 V8)
    10 => (V0 V1 V2 V3 V4 V5 V6 V7 V8 V9)
    11 => (V0 V1 V2 V3 V4 V5 V6 V7 V8 V9 V10)
    12 => (V0 V1 V2 V3 V4 V5 V6 V7 V8 V9 V10 V11)
    13 => (V0 V1 V2 V3 V4 V5 V6 V7 V8 V9 V10 V11 V12)
    14 => (V0 V1 V2 V3 V4 V5 V6 V7 V8 V9 V10 V11 V12 V13)
    15 => (V0 V1 V2 V3 V4 V5 V6 V7 V8 V9 V10 V11 V12 V13 V14)
    16 => (V0 V1 V2 V3 V4 V5 V6 V7 V8 V9 V10 V11 V12 V13 V14 V15)
}

#[cfg(feature = "alloc")]
impl<'de, V> DecodeBorrowed<'de> for alloc::vec::Vec<V>
where
    V: DecodeBorrowed<'de>,
{
    type Value = alloc::vec::Vec<V::Value>;

    fn decode_borrowed_with_format<R>(
        format: Format,
        reader: &mut R,
    ) -> core::result::Result<Self::Value, Error<R::Error>>
    where
        R: IoRead<'de>,
    {
        let len = match format {
            Format::FixArray(len) => len.into(),
            Format::Array16 => NbyteReader::<2>::read(reader)?,
            Format::Array32 => NbyteReader::<4>::read(reader)?,
            _ => return Err(Error::UnexpectedFormat),
        };

        let mut out: alloc::vec::Vec<<V as DecodeBorrowed<'de>>::Value> =
            alloc::vec::Vec::with_capacity(len);
        for _ in 0..len {
            out.push(V::decode_borrowed(reader)?);
        }
        Ok(out)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::decode::Decode;
    use rstest::rstest;

    #[rstest]
    #[case(&[0x92, 0x01, 0x02, 0x01], vec![1u8, 2], &[0x01])]
    #[case(&[0xdc, 0x00, 0x02, 0x2a, 0x2b], vec![42u8, 43], &[])]
    fn array_decode_success(
        #[case] buf: &[u8],
        #[case] expect: Vec<u8>,
        #[case] rest_expect: &[u8],
    ) {
        let mut r = crate::io::SliceReader::new(buf);
        let decoded = ArrayDecoder::<Vec<u8>, u8>::decode(&mut r).unwrap();
        assert_eq!(decoded, expect);
        assert_eq!(r.rest(), rest_expect);
    }

    #[rstest]
    fn array_decoder_unexpected_format() {
        let buf = &[0x81, 0x01, 0x02]; // map(1)
        let mut r = crate::io::SliceReader::new(buf);
        let err = ArrayDecoder::<Vec<u8>, u8>::decode(&mut r).unwrap_err();
        assert!(matches!(err, Error::UnexpectedFormat));
    }

    #[rstest]
    fn fixed_array_len0_success() {
        let buf = &[0x90]; // array(0)
        let mut r = crate::io::SliceReader::new(buf);
        let arr = <[u8; 0] as Decode>::decode(&mut r).unwrap();
        assert_eq!(arr, []);
        assert!(r.rest().is_empty());
    }

    #[rstest]
    fn fixed_array_len3_success() {
        let buf = &[0x93, 0x0a, 0x0b, 0x0c];
        let mut r = crate::io::SliceReader::new(buf);
        let arr = <[u8; 3] as Decode>::decode(&mut r).unwrap();
        assert_eq!(arr, [10u8, 11, 12]);
        assert!(r.rest().is_empty());
    }

    #[rstest]
    #[case(&[0x92, 0x01, 0x02])] // len=2
    #[case(&[0x94, 0x01, 0x02, 0x03, 0x04])] // len=4 
    fn fixed_array_len_mismatch(#[case] buf: &[u8]) {
        let mut r = crate::io::SliceReader::new(buf);
        let err = <[u8; 3] as Decode>::decode(&mut r).unwrap_err();
        assert!(matches!(err, Error::InvalidData));
    }

    #[rstest]
    fn tuple1_success() {
        let buf = &[0x91, 0x2a]; // [42]
        let mut r = crate::io::SliceReader::new(buf);
        let (v0,) = <(u8,) as Decode>::decode(&mut r).unwrap();
        assert_eq!(v0, 42);
        assert!(r.rest().is_empty());
    }

    #[rstest]
    #[case(&[0x92, 0x2a, 0x2b])] // fixarray
    #[case(&[0xdc, 0x00, 0x02, 0x2a, 0x2b])] // array16(2)
    fn tuple2_success(#[case] buf: &[u8]) {
        let mut r = crate::io::SliceReader::new(buf);
        let (a, b) = <(u8, u8) as Decode>::decode(&mut r).unwrap();
        assert_eq!((a, b), (42, 43));
        assert!(r.rest().is_empty());
    }

    #[rstest]
    fn tuple3_success() {
        let buf = &[0x93, 0x01, 0x02, 0x03];
        let mut r = crate::io::SliceReader::new(buf);
        let (a, b, c) = <(u8, u8, u8) as Decode>::decode(&mut r).unwrap();
        assert_eq!((a, b, c), (1, 2, 3));
        assert!(r.rest().is_empty());
    }

    #[rstest]
    #[case(&[0x92, 0x01, 0x02])] // len 2
    #[case(&[0xdc, 0x00, 0x04, 1, 2, 3, 4])] // len 4
    fn tuple_len_mismatch(#[case] buf: &[u8]) {
        let mut r = crate::io::SliceReader::new(buf);
        let err = <(u8, u8, u8) as Decode>::decode(&mut r).unwrap_err();
        assert!(matches!(err, Error::InvalidData));
    }

    #[rstest]
    fn tuple_unexpected_format() {
        let buf = &[0x81, 0x01, 0x02]; // map(1)
        let mut r = crate::io::SliceReader::new(buf);
        let err = <(u8,) as Decode>::decode(&mut r).unwrap_err();
        assert!(matches!(err, Error::UnexpectedFormat));
    }

    #[cfg(feature = "alloc")]
    #[test]
    fn vec_of_u8_success() {
        let buf = &[0x92, 0x2a, 0x2b]; // [42,43]
        let mut r = crate::io::SliceReader::new(buf);
        let v = <alloc::vec::Vec<u8> as Decode>::decode(&mut r).unwrap();
        assert_eq!(v, alloc::vec![42u8, 43]);
        assert!(r.rest().is_empty());
    }
}