orchid/orchid-api-traits/src/coding.rs

use std::collections::HashMap;
use std::future::Future;
use std::hash::Hash;
use std::io;
use std::num::NonZero;
use std::ops::{Range, RangeInclusive};
use std::pin::Pin;
use std::rc::Rc;
use std::sync::Arc;

use futures::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
use never::Never;
use ordered_float::NotNan;

use crate::{decode_err, decode_err_for, encode_enum, spin_on};

pub trait Decode: 'static + Sized {
	/// Decode an instance from the beginning of the buffer. Return the decoded
	/// data and the remaining buffer.
	fn decode<R: AsyncRead + ?Sized>(
		read: Pin<&mut R>,
	) -> impl Future<Output = io::Result<Self>> + '_;
	fn decode_slice(slc: &mut &[u8]) -> Self {
		spin_on(Self::decode(Pin::new(slc) as Pin<&mut _>)).expect("Decode from slice cannot fail")
	}
}
pub trait Encode {
	/// Append an instance of the struct to the buffer
	fn encode<W: AsyncWrite + ?Sized>(
		&self,
		write: Pin<&mut W>,
	) -> impl Future<Output = io::Result<()>>;
	fn encode_vec(&self, vec: &mut Vec<u8>) {
		spin_on(self.encode(Pin::new(vec) as Pin<&mut _>)).expect("Encode to vector cannot fail")
	}
}
pub trait Coding: Encode + Decode + Clone {
	fn get_decoder<T: 'static>(
		map: impl AsyncFn(Self) -> T + Clone + 'static,
	) -> impl AsyncFn(Pin<&mut dyn AsyncRead>) -> io::Result<T> {
		async move |r| Ok(map(Self::decode(r).await?).await)
	}
}
impl<T: Encode + Decode + Clone> Coding for T {}

macro_rules! num_impl {
	($number:ty) => {
		impl Decode for $number {
			async fn decode<R: AsyncRead + ?Sized>(mut read: Pin<&mut R>) -> io::Result<Self> {
				let mut bytes = [0u8; (<$number>::BITS / 8) as usize];
				read.read_exact(&mut bytes).await?;
				Ok(<$number>::from_be_bytes(bytes))
			}
		}
		impl Encode for $number {
			async fn encode<W: AsyncWrite + ?Sized>(&self, mut write: Pin<&mut W>) -> io::Result<()> {
				write.write_all(&self.to_be_bytes()).await
			}
		}
	};
}
num_impl!(u128);
num_impl!(u64);
num_impl!(u32);
num_impl!(u16);
num_impl!(u8);
num_impl!(i128);
num_impl!(i64);
num_impl!(i32);
num_impl!(i16);
num_impl!(i8);

macro_rules! nonzero_impl {
	($name:ty) => {
		impl Decode for NonZero<$name> {
			async fn decode<R: AsyncRead + ?Sized>(read: Pin<&mut R>) -> io::Result<Self> {
				Self::new(<$name as Decode>::decode(read).await?).ok_or_else(decode_err)
			}
		}
		impl Encode for NonZero<$name> {
			async fn encode<W: AsyncWrite + ?Sized>(&self, write: Pin<&mut W>) -> io::Result<()> {
				self.get().encode(write).await
			}
		}
	};
}

nonzero_impl!(u8);
nonzero_impl!(u16);
nonzero_impl!(u32);
nonzero_impl!(u64);
nonzero_impl!(u128);
nonzero_impl!(i8);
nonzero_impl!(i16);
nonzero_impl!(i32);
nonzero_impl!(i64);
nonzero_impl!(i128);

impl<T: Encode + ?Sized> Encode for &T {
	async fn encode<W: AsyncWrite + ?Sized>(&self, write: Pin<&mut W>) -> io::Result<()> {
		(**self).encode(write).await
	}
}
macro_rules! float_impl {
	($t:ty, $size:expr) => {
		impl Decode for NotNan<$t> {
			async fn decode<R: AsyncRead + ?Sized>(mut read: Pin<&mut R>) -> io::Result<Self> {
				let mut bytes = [0u8; $size];
				read.read_exact(&mut bytes).await?;
				NotNan::new(<$t>::from_be_bytes(bytes)).map_err(|_| decode_err())
			}
		}
		impl Encode for NotNan<$t> {
			async fn encode<W: AsyncWrite + ?Sized>(&self, mut write: Pin<&mut W>) -> io::Result<()> {
				write.write_all(&self.as_ref().to_be_bytes()).await
			}
		}
	};
}

float_impl!(f64, 8);
float_impl!(f32, 4);

impl Decode for String {
	async fn decode<R: AsyncRead + ?Sized>(mut read: Pin<&mut R>) -> io::Result<Self> {
		let len: usize = u64::decode(read.as_mut()).await?.try_into().map_err(decode_err_for)?;
		let mut data = vec![0u8; len];
		read.read_exact(&mut data).await?;
		Ok(std::str::from_utf8(&data).map_err(decode_err_for)?.to_owned())
	}
}
impl Encode for String {
	async fn encode<W: AsyncWrite + ?Sized>(&self, mut write: Pin<&mut W>) -> io::Result<()> {
		u64::try_from(self.len()).map_err(decode_err_for)?.encode(write.as_mut()).await?;
		write.write_all(self.as_bytes()).await
	}
}
impl Encode for str {
	async fn encode<W: AsyncWrite + ?Sized>(&self, mut write: Pin<&mut W>) -> io::Result<()> {
		u64::try_from(self.len()).map_err(decode_err_for)?.encode(write.as_mut()).await?;
		write.write_all(self.as_bytes()).await
	}
}
impl<T: Decode> Decode for Vec<T> {
	async fn decode<R: AsyncRead + ?Sized>(mut read: Pin<&mut R>) -> io::Result<Self> {
		let len = u64::decode(read.as_mut()).await?;
		let mut values = Vec::with_capacity(len.try_into().map_err(decode_err_for)?);
		for _ in 0..len {
			values.push(T::decode(read.as_mut()).await?);
		}
		Ok(values)
	}
}
impl<T: Encode> Encode for Vec<T> {
	async fn encode<W: AsyncWrite + ?Sized>(&self, write: Pin<&mut W>) -> io::Result<()> {
		self.as_slice().encode(write).await
	}
}
impl<T: Encode> Encode for [T] {
	async fn encode<W: AsyncWrite + ?Sized>(&self, mut write: Pin<&mut W>) -> io::Result<()> {
		u64::try_from(self.len()).unwrap().encode(write.as_mut()).await?;
		for t in self.iter() {
			t.encode(write.as_mut()).await?
		}
		Ok(())
	}
}
impl<T: Decode> Decode for Option<T> {
	async fn decode<R: AsyncRead + ?Sized>(mut read: Pin<&mut R>) -> io::Result<Self> {
		Ok(match bool::decode(read.as_mut()).await? {
			false => None,
			true => Some(T::decode(read).await?),
		})
	}
}
impl<T: Encode> Encode for Option<T> {
	async fn encode<W: AsyncWrite + ?Sized>(&self, mut write: Pin<&mut W>) -> io::Result<()> {
		self.is_some().encode(write.as_mut()).await?;
		if let Some(t) = self {
			t.encode(write).await?
		}
		Ok(())
	}
}
impl<T: Decode, E: Decode> Decode for Result<T, E> {
	async fn decode<R: AsyncRead + ?Sized>(mut read: Pin<&mut R>) -> io::Result<Self> {
		Ok(match bool::decode(read.as_mut()).await? {
			false => Self::Ok(T::decode(read).await?),
			true => Self::Err(E::decode(read).await?),
		})
	}
}

impl<T: Encode, E: Encode> Encode for Result<T, E> {
	async fn encode<W: AsyncWrite + ?Sized>(&self, write: Pin<&mut W>) -> io::Result<()> {
		match self {
			Ok(t) => encode_enum(write, 0, |w| t.encode(w)).await,
			Err(e) => encode_enum(write, 1, |w| e.encode(w)).await,
		}
	}
}
impl<K: Decode + Eq + Hash, V: Decode> Decode for HashMap<K, V> {
	async fn decode<R: AsyncRead + ?Sized>(mut read: Pin<&mut R>) -> io::Result<Self> {
		let len = u64::decode(read.as_mut()).await?;
		let mut map = HashMap::with_capacity(len.try_into().map_err(decode_err_for)?);
		for _ in 0..len {
			map.insert(K::decode(read.as_mut()).await?, V::decode(read.as_mut()).await?);
		}
		Ok(map)
	}
}
impl<K: Encode + Eq + Hash, V: Encode> Encode for HashMap<K, V> {
	async fn encode<W: AsyncWrite + ?Sized>(&self, mut write: Pin<&mut W>) -> io::Result<()> {
		u64::try_from(self.len()).unwrap().encode(write.as_mut()).await?;
		for (key, value) in self.iter() {
			key.encode(write.as_mut()).await?;
			value.encode(write.as_mut()).await?;
		}
		Ok(())
	}
}
macro_rules! tuple {
  (($($t:ident)*) ($($T:ident)*)) => {
    impl<$($T: Decode),*> Decode for ($($T,)*) {
			async fn decode<R: AsyncRead + ?Sized>(mut read: Pin<&mut R>) -> io::Result<Self> {
				Ok(($($T::decode(read.as_mut()).await?,)*))
			}
    }
    impl<$($T: Encode),*> Encode for ($($T,)*) {
			async fn encode<W: AsyncWrite + ?Sized>(&self, mut write: Pin<&mut W>) -> io::Result<()> {
        let ($($t,)*) = self;
        $( $t.encode(write.as_mut()).await?; )*
				Ok(())
      }
    }
  };
}

tuple!((t)(T));
tuple!((t u) (T U));
tuple!((t u v) (T U V));
tuple!((t u v x) (T U V X)); // 4
tuple!((t u v x y) (T U V X Y));
tuple!((t u v x y z) (T U V X Y Z));
tuple!((t u v x y z a) (T U V X Y Z A));
tuple!((t u v x y z a b) (T U V X Y Z A B)); // 8
tuple!((t u v x y z a b c) (T U V X Y Z A B C));
tuple!((t u v x y z a b c d) (T U V X Y Z A B C D));
tuple!((t u v x y z a b c d e) (T U V X Y Z A B C D E));
tuple!((t u v x y z a b c d e f) (T U V X Y Z A B C D E F)); // 12
tuple!((t u v x y z a b c d e f g) (T U V X Y Z A B C D E F G));
tuple!((t u v x y z a b c d e f g h) (T U V X Y Z A B C D E F G H));
tuple!((t u v x y z a b c d e f g h i) (T U V X Y Z A B C D E F G H I));
tuple!((t u v x y z a b c d e f g h i j) (T U V X Y Z A B C D E F G H I J)); // 16

impl Decode for () {
	async fn decode<R: AsyncRead + ?Sized>(_: Pin<&mut R>) -> io::Result<Self> { Ok(()) }
}
impl Encode for () {
	async fn encode<W: AsyncWrite + ?Sized>(&self, _: Pin<&mut W>) -> io::Result<()> { Ok(()) }
}
impl Decode for Never {
	async fn decode<R: AsyncRead + ?Sized>(_: Pin<&mut R>) -> io::Result<Self> {
		unreachable!("A value of Never cannot exist so it can't have been serialized");
	}
}
impl Encode for Never {
	async fn encode<W: AsyncWrite + ?Sized>(&self, _: Pin<&mut W>) -> io::Result<()> {
		match *self {}
	}
}
impl Decode for bool {
	async fn decode<R: AsyncRead + ?Sized>(mut read: Pin<&mut R>) -> io::Result<Self> {
		let mut buf = [0];
		read.read_exact(&mut buf).await?;
		Ok(buf[0] != 0)
	}
}
impl Encode for bool {
	async fn encode<W: AsyncWrite + ?Sized>(&self, mut write: Pin<&mut W>) -> io::Result<()> {
		write.write_all(&[if *self { 0xffu8 } else { 0u8 }]).await
	}
}
impl<T: Decode, const N: usize> Decode for [T; N] {
	async fn decode<R: AsyncRead + ?Sized>(mut read: Pin<&mut R>) -> io::Result<Self> {
		let mut v = Vec::with_capacity(N);
		for _ in 0..N {
			v.push(T::decode(read.as_mut()).await?);
		}
		match v.try_into() {
			Err(_) => unreachable!("The length of this stream is statically known"),
			Ok(arr) => Ok(arr),
		}
	}
}
impl<T: Encode, const N: usize> Encode for [T; N] {
	async fn encode<W: AsyncWrite + ?Sized>(&self, mut write: Pin<&mut W>) -> io::Result<()> {
		for t in self.iter() {
			t.encode(write.as_mut()).await?
		}
		Ok(())
	}
}

macro_rules! two_end_range {
  ($this:ident, $name:tt, $op:tt, $start:expr, $end:expr) => {
    impl<T: Decode> Decode for $name<T> {
			async fn decode<R: AsyncRead + ?Sized>(mut read: Pin<&mut R>) -> io::Result<Self> {
				Ok(T::decode(read.as_mut()).await? $op T::decode(read).await?)
			}
    }
    impl<T: Encode> Encode for $name<T> {
			async fn encode<W: AsyncWrite + ?Sized>(&self, mut write: Pin<&mut W>) -> io::Result<()> {
        let $this = self;
        ($start).encode(write.as_mut()).await?;
        ($end).encode(write).await?;
				Ok(())
      }
    }
  }
}

two_end_range!(x, Range, .., x.start, x.end);
two_end_range!(x, RangeInclusive, ..=, x.start(), x.end());

macro_rules! smart_ptr {
	($name:tt) => {
		impl<T: Decode> Decode for $name<T> {
			async fn decode<R: AsyncRead + ?Sized>(read: Pin<&mut R>) -> io::Result<Self> {
				Ok($name::new(T::decode(read).await?))
			}
		}
		impl<T: Encode> Encode for $name<T> {
			async fn encode<W: AsyncWrite + ?Sized>(&self, write: Pin<&mut W>) -> io::Result<()> {
				(**self).encode(write).await
			}
		}
	};
}

smart_ptr!(Arc);
smart_ptr!(Rc);
smart_ptr!(Box);

impl Decode for char {
	async fn decode<R: AsyncRead + ?Sized>(read: Pin<&mut R>) -> io::Result<Self> {
		char::from_u32(u32::decode(read).await?).ok_or_else(decode_err)
	}
}
impl Encode for char {
	async fn encode<W: AsyncWrite + ?Sized>(&self, write: Pin<&mut W>) -> io::Result<()> {
		(*self as u32).encode(write).await
	}
}