Add context

7 files changed, 183 insertions, 35 deletions

diff --git a/src/codec/decode.rs b/src/codec/decode.rs
index 1059811..3e03a84 100644
--- a/src/codec/decode.rs
+++ b/src/codec/decode.rs
@@ -1,11 +1,11 @@
 use std::{io::Read, sync::Arc};
 
-use crate::DEFAULT_BUFFER_LEN;
+use crate::{DEFAULT_BUFFER_LEN, context::Context};
 
 macro_rules! impl_decode {
     ($type: ty) => {
         impl Decode for $type {
-            fn decode(reader: &mut dyn Read) -> crate::codec::error::Result<Self>
+            fn decode(reader: &mut dyn Read, _: &dyn Context) -> crate::codec::error::Result<Self>
             where
                 Self: Sized,
             {
@@ -19,13 +19,13 @@ macro_rules! impl_decode {
 }
 
 pub trait Decode {
-    fn decode(reader: &mut dyn Read) -> crate::codec::error::Result<Self>
+    fn decode(reader: &mut dyn Read, ctx: &dyn Context) -> crate::codec::error::Result<Self>
     where
         Self: Sized;
 }
 
 impl Decode for bool {
-    fn decode(reader: &mut dyn Read) -> crate::codec::error::Result<Self>
+    fn decode(reader: &mut dyn Read, _: &dyn Context) -> crate::codec::error::Result<Self>
     where
         Self: Sized,
     {
@@ -35,8 +35,48 @@ impl Decode for bool {
     }
 }
 
+impl<T: Decode> Decode for Option<T> {
+    fn decode(reader: &mut dyn Read, ctx: &dyn Context) -> crate::codec::error::Result<Self>
+    where
+        Self: Sized,
+    {
+        let tag = u8::decode(reader, ctx)?;
+        if tag == 0 {
+            Ok(None)
+        } else {
+            Ok(Some(T::decode(reader, ctx)?))
+        }
+    }
+}
+
+impl<T: Decode> Decode for Vec<T> {
+    fn decode(reader: &mut dyn Read, ctx: &dyn Context) -> crate::codec::error::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)
+    }
+}
+
+impl Decode for String {
+    fn decode(reader: &mut dyn Read, _: &dyn Context) -> crate::codec::error::Result<Self>
+    where
+        Self: Sized,
+    {
+        // assume reader is limited
+        let mut bytes: Vec<u8> = Vec::new();
+        reader.read_to_end(&mut bytes)?;
+        Ok(String::from_utf8_lossy(&bytes).to_string())
+    }
+}
+
 impl Decode for Arc<[u8]> {
-    fn decode(reader: &mut dyn Read) -> crate::codec::error::Result<Self> {
+    fn decode(reader: &mut dyn Read, _: &dyn Context) -> crate::codec::error::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()))
@@ -44,7 +84,7 @@ impl Decode for Arc<[u8]> {
 }
 
 impl<const S: usize> Decode for [u8; S] {
-    fn decode(reader: &mut dyn Read) -> crate::codec::error::Result<Self>
+    fn decode(reader: &mut dyn Read, _: &dyn Context) -> crate::codec::error::Result<Self>
     where
         Self: Sized,
     {
@@ -69,22 +109,22 @@ impl_decode!(i128);
 impl_decode!(isize);
 
 impl Decode for f64 {
-    fn decode(reader: &mut dyn Read) -> crate::codec::error::Result<Self>
+    fn decode(reader: &mut dyn Read, ctx: &dyn Context) -> crate::codec::error::Result<Self>
     where
         Self: Sized,
     {
-        let bits = u64::decode(reader)?;
+        let bits = u64::decode(reader, ctx)?;
         let value = f64::from_bits(bits);
         Ok(value)
     }
 }
 
 impl Decode for f32 {
-    fn decode(reader: &mut dyn Read) -> crate::codec::error::Result<Self>
+    fn decode(reader: &mut dyn Read, ctx: &dyn Context) -> crate::codec::error::Result<Self>
     where
         Self: Sized,
     {
-        let bits = u32::decode(reader)?;
+        let bits = u32::decode(reader, ctx)?;
         let value = f32::from_bits(bits);
         Ok(value)
     }
diff --git a/src/codec/encode.rs b/src/codec/encode.rs
index a40f551..70d9d5c 100644
--- a/src/codec/encode.rs
+++ b/src/codec/encode.rs
@@ -1,10 +1,10 @@
-use crate::codec::error::Error;
+use crate::{codec::error::Error, context::Context};
 use std::{io::Write, sync::Arc};
 
 macro_rules! impl_encode {
     ($t: ty) => {
         impl Encode for $t {
-            fn encode(self, buffer: &mut dyn Write) -> Result<usize, Error> {
+            fn encode(&self, buffer: &mut dyn Write, _: &dyn Context) -> Result<usize, Error> {
                 let size = buffer.write(&self.to_be_bytes())?;
                 Ok(size)
             }
@@ -13,26 +13,57 @@ macro_rules! impl_encode {
 }
 
 pub trait Encode {
-    fn encode(self, buffer: &mut dyn Write) -> Result<usize, Error>;
+    fn encode(&self, buffer: &mut dyn Write, ctx: &dyn Context) -> Result<usize, Error>;
 }
 
 impl Encode for bool {
-    fn encode(self, buffer: &mut dyn Write) -> Result<usize, Error> {
-        let size = buffer.write(&[self as u8])?;
+    fn encode(&self, buffer: &mut dyn Write, _: &dyn Context) -> Result<usize, Error> {
+        let size = buffer.write(&[*self as u8])?;
         Ok(size)
     }
 }
 
+impl<T: Encode> Encode for Option<T> {
+    fn encode(&self, buffer: &mut dyn Write, ctx: &dyn Context) -> Result<usize, Error> {
+        match self {
+            Some(val) => {
+                let mut size = 1u8.encode(buffer, ctx)?;
+                size += val.encode(buffer, ctx)?;
+                Ok(size)
+            }
+            None => 0u8.encode(buffer, ctx),
+        }
+    }
+}
+
+impl<T: Encode> Encode for Vec<T> {
+    fn encode(&self, buffer: &mut dyn Write, ctx: &dyn Context) -> Result<usize, Error> {
+        let mut size = self.len().encode(buffer, ctx)?;
+        for item in self {
+            size += item.encode(buffer, ctx)?;
+        }
+        Ok(size)
+    }
+}
+
+impl Encode for String {
+    fn encode(&self, buffer: &mut dyn Write, _: &dyn Context) -> Result<usize, Error> {
+        let utf8 = self.as_bytes();
+        buffer.write_all(utf8)?;
+        Ok(utf8.len())
+    }
+}
+
 impl Encode for Arc<[u8]> {
-    fn encode(self, buffer: &mut dyn Write) -> Result<usize, Error> {
-        let len = buffer.write(&self)?;
+    fn encode(&self, buffer: &mut dyn Write, _: &dyn Context) -> Result<usize, Error> {
+        let len = buffer.write(self)?;
         Ok(len)
     }
 }
 
 impl<const S: usize> Encode for [u8; S] {
-    fn encode(self, buffer: &mut dyn Write) -> Result<usize, Error> {
-        let len = buffer.write(&self)?;
+    fn encode(&self, buffer: &mut dyn Write, _ctx: &dyn Context) -> Result<usize, Error> {
+        let len = buffer.write(self)?;
         Ok(len)
     }
 }
@@ -53,3 +84,12 @@ impl_encode!(isize);
 
 impl_encode!(f32);
 impl_encode!(f64);
+
+impl<T> Encode for T
+where
+    T: AsRef<dyn Encode>,
+{
+    fn encode(&self, buffer: &mut dyn Write, ctx: &dyn Context) -> Result<usize, Error> {
+        self.as_ref().encode(buffer, ctx)
+    }
+}
diff --git a/src/codec/error.rs b/src/codec/error.rs
index 6639c98..6037a9a 100644
--- a/src/codec/error.rs
+++ b/src/codec/error.rs
@@ -3,6 +3,7 @@ use std::fmt::Display;
 #[derive(Debug)]
 pub enum Error {
     IoError(std::io::Error),
+    Custom(String),
 }
 
 impl std::error::Error for Error {}
@@ -20,3 +21,15 @@ impl From<std::io::Error> for Error {
         Self::IoError(value)
     }
 }
+
+impl From<String> for Error {
+    fn from(value: String) -> Self {
+        Self::Custom(value)
+    }
+}
+
+impl From<&str> for Error {
+    fn from(value: &str) -> Self {
+        Self::Custom(value.to_string())
+    }
+}
diff --git a/src/context.rs b/src/context.rs
new file mode 100644
index 0000000..47661d0
--- /dev/null
+++ b/src/context.rs
@@ -0,0 +1,6 @@
+use std::any::Any;
+
+pub trait Context {
+    fn get(&self, key: &str) -> Option<&dyn Any>;
+    fn set(&mut self, key: &str, value: &dyn Any);
+}
diff --git a/src/lib.rs b/src/lib.rs
index 24c2033..756f9e2 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,4 +1,7 @@
 pub mod codec;
+pub mod context;
 pub mod types;
 
+pub use zr_protocol_macros as macros;
+
 pub(crate) const DEFAULT_BUFFER_LEN: usize = 2048;
diff --git a/src/types/prefix/count.rs b/src/types/prefix/count.rs
index 7495660..cf541ec 100644
--- a/src/types/prefix/count.rs
+++ b/src/types/prefix/count.rs
@@ -2,14 +2,34 @@ use std::{io::Read, marker::PhantomData};
 
 use crate::{
     codec::{Codec, decode::Decode, encode::Encode},
+    context::Context,
     types::size::Size,
 };
 
-pub struct CountPrefix<I: Codec, L: Codec + Size, D: IntoIterator<Item = I>> {
+pub struct CountPrefix<I, L, D>
+where
+    I: Codec,
+    L: Codec + Size,
+    D: IntoIterator<Item = I>,
+{
     data: D,
     _len: PhantomData<L>,
 }
 
+impl<I, L, D> Clone for CountPrefix<I, L, D>
+where
+    I: Codec,
+    L: Codec + Size,
+    D: IntoIterator<Item = I> + Clone,
+{
+    fn clone(&self) -> Self {
+        Self {
+            data: self.data.clone(),
+            _len: PhantomData,
+        }
+    }
+}
+
 impl<I, L, D> CountPrefix<I, L, D>
 where
     I: Codec,
@@ -50,19 +70,20 @@ impl<I, L, D> Encode for CountPrefix<I, L, D>
 where
     I: Codec,
     L: Codec + Size,
-    D: IntoIterator<Item = I>,
+    D: IntoIterator<Item = I> + Clone,
 {
     fn encode(
-        self,
+        &self,
         buffer: &mut dyn std::io::prelude::Write,
+        ctx: &dyn Context,
     ) -> Result<usize, crate::codec::error::Error>
     where
         Self: Sized,
     {
-        let vec = self.data.into_iter().collect::<Vec<I>>();
-        let mut l = L::from_size(vec.len()).encode(buffer)?;
+        let vec = self.data.clone().into_iter().collect::<Vec<I>>();
+        let mut l = L::from_size(vec.len()).encode(buffer, ctx)?;
         for item in vec {
-            l += item.encode(buffer)?;
+            l += item.encode(buffer, ctx)?;
         }
         Ok(l)
     }
@@ -74,14 +95,14 @@ where
     L: Codec + Size,
     D: IntoIterator<Item = I> + FromIterator<I>,
 {
-    fn decode(reader: &mut dyn Read) -> crate::codec::error::Result<Self>
+    fn decode(reader: &mut dyn Read, ctx: &dyn Context) -> crate::codec::error::Result<Self>
     where
         Self: Sized,
     {
-        let len = L::decode(reader)?.as_size();
+        let len = L::decode(reader, ctx)?.as_size();
         let mut data = Vec::with_capacity(len);
         for _ in 0..len {
-            let item = I::decode(reader)?;
+            let item = I::decode(reader, ctx)?;
             data.push(item);
         }
         Ok(Self {
diff --git a/src/types/prefix/length.rs b/src/types/prefix/length.rs
index c825375..fe159fc 100644
--- a/src/types/prefix/length.rs
+++ b/src/types/prefix/length.rs
@@ -3,14 +3,32 @@ use std::{io::Write, marker::PhantomData};
 use crate::{
     DEFAULT_BUFFER_LEN,
     codec::{self, Codec, decode::Decode, encode::Encode},
+    context::Context,
     types::size::Size,
 };
 
-pub struct LenPrefixed<L: Codec + Size, D: Codec> {
+pub struct LenPrefixed<L, D>
+where
+    L: Codec + Size,
+    D: Codec,
+{
     data: D,
     _len: PhantomData<L>,
 }
 
+impl<L, D> Clone for LenPrefixed<L, D>
+where
+    L: Codec + Size,
+    D: Codec + Clone,
+{
+    fn clone(&self) -> Self {
+        Self {
+            data: self.data.clone(),
+            _len: PhantomData,
+        }
+    }
+}
+
 impl<L, D> LenPrefixed<L, D>
 where
     L: Codec + Size,
@@ -49,10 +67,14 @@ where
     L: Codec + Size,
     D: Codec,
 {
-    fn encode(self, writer: &mut dyn Write) -> Result<usize, codec::error::Error> {
+    fn encode(
+        &self,
+        writer: &mut dyn Write,
+        ctx: &dyn Context,
+    ) -> Result<usize, codec::error::Error> {
         let mut buf = Vec::with_capacity(DEFAULT_BUFFER_LEN);
-        self.data.encode(&mut buf)?;
-        let mut l = L::from_size(buf.len()).encode(writer)?;
+        self.data.encode(&mut buf, ctx)?;
+        let mut l = L::from_size(buf.len()).encode(writer, ctx)?;
         l += writer.write(&buf)?;
         Ok(l)
     }
@@ -63,14 +85,17 @@ where
     L: Codec + Size,
     D: Codec,
 {
-    fn decode(reader: &mut dyn std::io::prelude::Read) -> crate::codec::error::Result<Self>
+    fn decode(
+        reader: &mut dyn std::io::prelude::Read,
+        ctx: &dyn Context,
+    ) -> crate::codec::error::Result<Self>
     where
         Self: Sized,
     {
-        let len = L::decode(reader)?.as_size();
+        let len = L::decode(reader, ctx)?.as_size();
         let mut limited = vec![0_u8; len];
         reader.read_exact(&mut limited)?;
-        let data = D::decode(&mut &limited[..])?;
+        let data = D::decode(&mut &limited[..], ctx)?;
         Ok(Self {
             data,
             _len: PhantomData,