add encode/decode implementation

4 files changed, 281 insertions, 38 deletions

diff --git a/src/codec/decode.rs b/src/codec/decode.rs
index 39253a2..529c448 100644
--- a/src/codec/decode.rs
+++ b/src/codec/decode.rs
@@ -1,10 +1,36 @@
 //! Decoding: reading protocol values from a byte stream.
-use crate::context::Context;
-use std::{io::Read, sync::Arc};
+use std::{
+    collections::{BTreeMap, BTreeSet, BinaryHeap, HashMap, HashSet},
+    hash::Hash,
+    io::Read,
+    sync::Arc,
+};
 
 use crate::{DEFAULT_BUFFER_LEN, codec::error::Result};
+use crate::{codec::error::CodecError, context::Context};
 
 macro_rules! impl_decode {
+    (u8) => {
+        impl<Data> Decode<Data> for u8 {
+            fn decode(reader: &mut dyn Read, _: &Context<Data>) -> Result<Self>
+            where
+                Self: Sized,
+            {
+                let mut byte = [0u8; 1];
+                reader.read_exact(&mut byte)?;
+                Ok(byte[0])
+            }
+
+            fn decode_slice(reader: &mut dyn Read, _: &Context<Data>) -> Result<Vec<Self>>
+            where
+                Self: Sized,
+            {
+                let mut buf = Vec::new();
+                reader.read_to_end(&mut buf)?;
+                Ok(buf)
+            }
+        }
+    };
     ($type: ty) => {
         impl<Data> Decode<Data> for $type {
             // decode number using big endian
@@ -17,86 +43,225 @@ macro_rules! impl_decode {
                 reader.read_exact(&mut bytes)?;
                 Ok(<$type>::from_be_bytes(bytes))
             }
+
+            fn decode_slice(reader: &mut dyn Read, _: &Context<Data>) -> Result<Vec<Self>>
+            where
+                Self: Sized,
+            {
+                let mut buf = Vec::with_capacity(DEFAULT_BUFFER_LEN);
+                reader.read_to_end(&mut buf);
+                const BYTES: usize = <$type>::BITS as usize / 8;
+                Ok(buf
+                    .chunks(BYTES)
+                    .map(|slice| {
+                        let Some(bytes): Option<&[u8; BYTES]> = slice.as_array() else {
+                            unreachable!()
+                        };
+                        <$type>::from_be_bytes(*bytes)
+                    })
+                    .collect())
+            }
         }
     };
 }
 
+macro_rules! impl_decode_tuples {
+    ($($generic: ident),+) => {
+            impl<Data, $($generic),+> Decode<Data> for ($($generic,)+) where $($generic: Decode<Data>,)+  {
+                #[allow(non_snake_case)]
+                fn decode(reader: &mut dyn Read, ctx: &Context<Data>) -> Result<Self>
+                where
+                    Self: Sized {
+                        $(
+                            let $generic = $generic::decode(reader, ctx)?;
+                        )+
+                        Ok(($($generic,)+))
+                }
+            }
+    };
+}
+
 /// Read a value from a byte stream
 pub trait Decode<Data> {
     fn decode(reader: &mut dyn Read, ctx: &Context<Data>) -> Result<Self>
     where
         Self: Sized;
+
+    // TODO : change for anything else than Vec
+    /// assume reader contains only the slice to decode
+    fn decode_slice(reader: &mut dyn Read, ctx: &Context<Data>) -> Result<Vec<Self>>
+    where
+        Self: Sized,
+    {
+        let mut buf = Vec::new();
+        loop {
+            match Self::decode(reader, ctx) {
+                Ok(value) => buf.push(value),
+                Err(CodecError::IoError(err))
+                    if let std::io::ErrorKind::UnexpectedEof = err.kind() =>
+                {
+                    return Ok(buf);
+                }
+                Err(err) => return Err(err),
+            }
+        }
+    }
 }
 
-impl<Data> Decode<Data> for bool {
-    fn decode(reader: &mut dyn Read, _: &Context<Data>) -> Result<Self>
+// ~ Decode arrays
+
+impl<Data, T: Decode<Data>> Decode<Data> for Vec<T> {
+    fn decode(reader: &mut dyn Read, ctx: &Context<Data>) -> Result<Self>
     where
         Self: Sized,
     {
-        let mut byte = [0_u8];
-        reader.read_exact(&mut byte)?;
-        Ok(byte[0] != 0)
+        let len = usize::decode(reader, ctx)?;
+        let mut vec = Vec::with_capacity(len);
+        for _ in 0..len {
+            vec.push(T::decode(reader, ctx)?);
+        }
+        Ok(vec)
     }
 }
 
 impl<Data, T: Decode<Data>> Decode<Data> for Option<T> {
+    /// Assume is always some
     fn decode(reader: &mut dyn Read, ctx: &Context<Data>) -> Result<Self>
     where
         Self: Sized,
     {
-        let tag = u8::decode(reader, ctx)?;
-        if tag == 0 {
-            Ok(None)
-        } else {
-            Ok(Some(T::decode(reader, ctx)?))
-        }
+        T::decode(reader, ctx).map(Some)
     }
 }
 
-impl<Data, T: Decode<Data>> Decode<Data> for Vec<T> {
+// ~ Decode slices
+
+impl<Data, T, const S: usize> Decode<Data> for [T; S]
+where
+    T: Decode<Data> + Default + Copy,
+{
     fn decode(reader: &mut dyn Read, ctx: &Context<Data>) -> Result<Self>
     where
         Self: Sized,
     {
-        let len = usize::decode(reader, ctx)?;
-        let mut vec = Vec::with_capacity(len);
-        for _ in 0..len {
-            vec.push(T::decode(reader, ctx)?);
-        }
-        Ok(vec)
+        let slice = T::decode_slice(reader, ctx)?;
+        let got = slice.len();
+        slice
+            .as_array()
+            .cloned()
+            .ok_or(CodecError::WrongSize { expected: S, got })
     }
 }
 
-impl<Data> Decode<Data> for String {
-    fn decode(reader: &mut dyn Read, _: &Context<Data>) -> Result<Self>
+impl<Data, T> Decode<Data> for &[T]
+where
+    T: Decode<Data> + Default + Copy,
+{
+    fn decode(_: &mut dyn Read, _: &Context<Data>) -> Result<Self>
     where
         Self: Sized,
     {
-        let mut bytes: Vec<u8> = Vec::new();
-        reader.read_to_end(&mut bytes)?;
-        Ok(String::from_utf8_lossy(&bytes).to_string())
+        unimplemented!()
     }
 }
 
-impl<Data> Decode<Data> for Arc<[u8]> {
-    fn decode(reader: &mut dyn Read, _: &Context<Data>) -> Result<Self> {
-        let mut buf = Vec::with_capacity(DEFAULT_BUFFER_LEN);
-        reader.read_to_end(&mut buf)?;
-        Ok(Arc::<[u8]>::from(buf.into_boxed_slice()))
+// ~ Decode set
+
+impl<Data, T> Decode<Data> for HashSet<T>
+where
+    T: Decode<Data> + Eq + Hash,
+{
+    fn decode(reader: &mut dyn Read, ctx: &Context<Data>) -> Result<Self>
+    where
+        Self: Sized,
+    {
+        let slice = T::decode_slice(reader, ctx)?;
+        Ok(slice.into_iter().collect())
     }
 }
 
-impl<Data, const S: usize> Decode<Data> for [u8; S] {
+impl<Data, T> Decode<Data> for BTreeSet<T>
+where
+    T: Decode<Data> + Ord,
+{
+    fn decode(reader: &mut dyn Read, ctx: &Context<Data>) -> Result<Self>
+    where
+        Self: Sized,
+    {
+        let slice = T::decode_slice(reader, ctx)?;
+        Ok(slice.into_iter().collect())
+    }
+}
+
+// ~ Decode misc
+
+impl<Data, T> Decode<Data> for BinaryHeap<T>
+where
+    T: Decode<Data> + Ord,
+{
+    fn decode(reader: &mut dyn Read, ctx: &Context<Data>) -> Result<Self>
+    where
+        Self: Sized,
+    {
+        let slice = T::decode_slice(reader, ctx)?;
+        Ok(BinaryHeap::from_iter(slice))
+    }
+}
+
+// ~ Decode maps
+
+impl<Data, K, V> Decode<Data> for HashMap<K, V>
+where
+    K: Decode<Data> + Eq + Hash,
+    V: Decode<Data>,
+{
+    fn decode(reader: &mut dyn Read, ctx: &Context<Data>) -> Result<Self>
+    where
+        Self: Sized,
+    {
+        let slice = <(K, V) as Decode<Data>>::decode_slice(reader, ctx)?;
+        Ok(slice.into_iter().collect())
+    }
+}
+
+impl<Data, K, V> Decode<Data> for BTreeMap<K, V>
+where
+    K: Decode<Data> + Ord,
+    V: Decode<Data>,
+{
+    fn decode(reader: &mut dyn Read, ctx: &Context<Data>) -> Result<Self>
+    where
+        Self: Sized,
+    {
+        let slice = <(K, V) as Decode<Data>>::decode_slice(reader, ctx)?;
+        Ok(slice.into_iter().collect())
+    }
+}
+
+// ~ Decode string
+
+impl<Data> Decode<Data> for &str {
+    fn decode(_: &mut dyn Read, _: &Context<Data>) -> Result<Self>
+    where
+        Self: Sized,
+    {
+        unimplemented!()
+    }
+}
+
+impl<Data> Decode<Data> for String {
     fn decode(reader: &mut dyn Read, _: &Context<Data>) -> Result<Self>
     where
         Self: Sized,
     {
-        let mut slice = [0_u8; S];
-        reader.read_exact(&mut slice)?;
-        Ok(slice)
+        let mut bytes: Vec<u8> = Vec::new();
+        reader.read_to_end(&mut bytes)?;
+        Ok(String::from_utf8_lossy(&bytes).to_string())
     }
 }
 
+// ~ Decode primitive
+
 impl_decode!(u8);
 impl_decode!(u16);
 impl_decode!(u32);
@@ -132,3 +297,58 @@ impl<Data> Decode<Data> for f32 {
         Ok(value)
     }
 }
+
+impl<Data> Decode<Data> for bool {
+    fn decode(reader: &mut dyn Read, _: &Context<Data>) -> Result<Self>
+    where
+        Self: Sized,
+    {
+        let mut byte = [0_u8];
+        reader.read_exact(&mut byte)?;
+        Ok(byte[0] != 0)
+    }
+}
+
+// ~ Decode tuple
+
+impl_decode_tuples!(A);
+impl_decode_tuples!(A, B);
+impl_decode_tuples!(A, B, C);
+impl_decode_tuples!(A, B, C, D);
+impl_decode_tuples!(A, B, C, D, E);
+impl_decode_tuples!(A, B, C, D, E, F);
+impl_decode_tuples!(A, B, C, D, E, F, G);
+impl_decode_tuples!(A, B, C, D, E, F, G, H);
+impl_decode_tuples!(A, B, C, D, E, F, G, H, I);
+impl_decode_tuples!(A, B, C, D, E, F, G, H, I, J);
+impl_decode_tuples!(A, B, C, D, E, F, G, H, I, J, K);
+impl_decode_tuples!(A, B, C, D, E, F, G, H, I, J, K, L);
+impl_decode_tuples!(A, B, C, D, E, F, G, H, I, J, K, L, M);
+impl_decode_tuples!(A, B, C, D, E, F, G, H, I, J, K, L, M, N);
+impl_decode_tuples!(A, B, C, D, E, F, G, H, I, J, K, L, M, N, O);
+impl_decode_tuples!(A, B, C, D, E, F, G, H, I, J, K, L, M, N, O, P);
+impl_decode_tuples!(A, B, C, D, E, F, G, H, I, J, K, L, M, N, O, P, Q);
+
+// ~ Decode pointer
+
+impl<Data, T> Decode<Data> for std::sync::Arc<T>
+where
+    T: Decode<Data>,
+{
+    fn decode(reader: &mut dyn Read, ctx: &Context<Data>) -> Result<Self> {
+        T::decode(reader, ctx).map(Arc::new)
+    }
+}
+
+impl<Data, T> Decode<Data> for Arc<[T]>
+where
+    T: Decode<Data>,
+{
+    fn decode(reader: &mut dyn Read, ctx: &Context<Data>) -> Result<Self>
+    where
+        Self: Sized,
+    {
+        let slice = T::decode_slice(reader, ctx)?;
+        Ok(slice.into())
+    }
+}
diff --git a/src/codec/encode.rs b/src/codec/encode.rs
index ae5eb4e..7463995 100644
--- a/src/codec/encode.rs
+++ b/src/codec/encode.rs
@@ -269,6 +269,28 @@ impl<Data> Encode<Data> for bool {
     }
 }
 
+// ~ Encode pointer
+
+impl<Data, T> Encode<Data> for std::sync::Arc<T>
+where
+    T: Encode<Data>,
+{
+    fn encode(&self, buffer: &mut dyn Write, ctx: &Context<Data>) -> Result<usize> {
+        let data = self.as_ref();
+        data.encode(buffer, ctx)
+    }
+}
+
+impl<Data, T> Encode<Data> for std::sync::Arc<[T]>
+where
+    T: Encode<Data>,
+{
+    fn encode(&self, buffer: &mut dyn Write, ctx: &Context<Data>) -> Result<usize> {
+        let data = self.as_ref();
+        data.encode(buffer, ctx)
+    }
+}
+
 impl_encode!(u8);
 impl_encode!(u16);
 impl_encode!(u32);
diff --git a/src/codec/error.rs b/src/codec/error.rs
index 7d88d55..43da488 100644
--- a/src/codec/error.rs
+++ b/src/codec/error.rs
@@ -4,7 +4,8 @@ use thiserror::Error;
 pub enum CodecError {
     #[error("io error: {0}")]
     IoError(#[from] std::io::Error),
-
+    #[error("WrongSize expected : {expected} ; got : {got}")]
+    WrongSize { expected: usize, got: usize },
     #[error("{0}")]
     Custom(String),
 }
diff --git a/src/types/prefix/length.rs b/src/types/prefix/length.rs
index 8704d01..f30eea0 100644
--- a/src/types/prefix/length.rs
+++ b/src/types/prefix/length.rs
@@ -62,9 +62,9 @@ where
         let mut buf = Vec::with_capacity(DEFAULT_BUFFER_LEN);
         self.data.encode(&mut buf, ctx)?;
         let len = L::from_size(buf.len());
-        let mut l = len.encode(writer, ctx)?;
-        l += writer.write(&buf)?;
-        Ok(l)
+        let mut len = len.encode(writer, ctx)?;
+        len += writer.write(&buf)?;
+        Ok(len)
     }
 }