#![allow(clippy::too_many_arguments,non_camel_case_types,non_upper_case_globals,dead_code,unused, non_snake_case)]
use bytes::{Buf, BufMut, Bytes, BytesMut};
use std::borrow::Cow;
use std::collections::hash_map::RandomState;
use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet, LinkedList, VecDeque};
use std::hash::Hash;
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr};
use std::time::{Duration, SystemTime, UNIX_EPOCH};
#[macro_export]
#[macro_use]
pub use wyre_derive::{Pack, Unpack};
#[derive(Debug)]
pub enum WyreError {
Io(std::io::Error),
InvalidFormat(String),
Utf8Error(std::string::FromUtf8Error),
OutOfRange(&'static str),
}
pub trait Pack {
fn pack<B: BufMut>(&self, buf: &mut B);
}
pub trait Unpack: Sized {
fn unpack(buf: &mut Bytes) -> Result<Self, WyreError>;
}
#[inline(always)]
pub fn write_varint<B: BufMut>(buf: &mut B, mut value: u64) {
loop {
let mut byte = (value & 0x7F) as u8;
value >>= 7;
if value != 0 {
byte |= 0x80;
}
buf.put_u8(byte);
if value == 0 {
break;
}
}
}
#[inline(always)]
pub fn read_varint(buf: &mut Bytes) -> Result<u64, WyreError> {
let mut value = 0u64;
let mut shift = 0u32;
loop {
if buf.is_empty() {
return Err(WyreError::InvalidFormat("Incomplete varint".into()));
}
let byte = buf.get_u8();
value |= ((byte & 0x7F) as u64) << shift;
if byte & 0x80 == 0 {
return Ok(value);
}
shift += 7;
if shift >= 64 {
return Err(WyreError::InvalidFormat("Varint too long".into()));
}
}
}
macro_rules! impl_pack_for_int {
($t:ty, $zero_optimization:expr) => {
impl Pack for $t {
#[inline]
fn pack<B: BufMut>(&self, buf: &mut B) {
if $zero_optimization && *self == 0 {
buf.put_u8(0);
} else {
write_varint(
buf,
(*self as u64).wrapping_shl(1) ^ (*self as i64 >> 63) as u64,
);
}
}
}
impl Unpack for $t {
#[inline]
fn unpack(buf: &mut Bytes) -> Result<Self, WyreError> {
let value = read_varint(buf)?;
if $zero_optimization && value == 0 {
return Ok(0);
}
Ok(((value >> 1) as $t) ^ (-((value & 1) as i64) as $t))
}
}
};
}
impl_pack_for_int!(i8, false);
impl_pack_for_int!(i16, false);
impl_pack_for_int!(i32, true);
impl_pack_for_int!(i64, true);
impl_pack_for_int!(u8, false);
impl_pack_for_int!(u16, false);
impl_pack_for_int!(u32, true);
impl_pack_for_int!(u64, true);
impl Pack for bool {
#[inline]
fn pack<B: BufMut>(&self, buf: &mut B) {
buf.put_u8(*self as u8);
}
}
impl Unpack for bool {
#[inline]
fn unpack(buf: &mut Bytes) -> Result<Self, WyreError> {
Ok(buf.get_u8() != 0)
}
}
impl Pack for f32 {
#[inline]
fn pack<B: BufMut>(&self, buf: &mut B) {
buf.put_f32_le(*self);
}
}
impl Unpack for f32 {
#[inline]
fn unpack(buf: &mut Bytes) -> Result<Self, WyreError> {
Ok(buf.get_f32_le())
}
}
impl Pack for f64 {
#[inline]
fn pack<B: BufMut>(&self, buf: &mut B) {
buf.put_f64_le(*self);
}
}
impl Unpack for f64 {
#[inline]
fn unpack(buf: &mut Bytes) -> Result<Self, WyreError> {
Ok(buf.get_f64_le())
}
}
impl Pack for char {
#[inline]
fn pack<B: BufMut>(&self, buf: &mut B) {
write_varint(buf, *self as u64);
}
}
impl Unpack for char {
#[inline]
fn unpack(buf: &mut Bytes) -> Result<Self, WyreError> {
let value = read_varint(buf)?;
char::from_u32(value as u32).ok_or_else(|| WyreError::InvalidFormat("Invalid char".into()))
}
}
impl Pack for String {
#[inline]
fn pack<B: BufMut>(&self, buf: &mut B) {
write_varint(buf, self.len() as u64);
buf.put_slice(self.as_bytes());
}
}
impl Unpack for String {
#[inline]
fn unpack(buf: &mut Bytes) -> Result<Self, WyreError> {
let len = read_varint(buf)? as usize;
if buf.remaining() < len {
return Err(WyreError::InvalidFormat(
"Not enough bytes for string".into(),
));
}
let bytes = buf.split_to(len);
String::from_utf8(bytes.to_vec()).map_err(WyreError::Utf8Error)
}
}
impl<'a> Pack for &'a str {
#[inline]
fn pack<B: BufMut>(&self, buf: &mut B) {
write_varint(buf, self.len() as u64);
buf.put_slice(self.as_bytes());
}
}
impl<'a> Pack for Cow<'a, str> {
#[inline]
fn pack<B: BufMut>(&self, buf: &mut B) {
write_varint(buf, self.len() as u64);
buf.put_slice(self.as_bytes());
}
}
impl<'a> Unpack for Cow<'a, str> {
#[inline]
fn unpack(buf: &mut Bytes) -> Result<Self, WyreError> {
let len = read_varint(buf)? as usize;
if buf.remaining() < len {
return Err(WyreError::InvalidFormat(
"Not enough bytes for Cow<str>".into(),
));
}
let bytes = buf.split_to(len);
String::from_utf8(bytes.to_vec())
.map(Cow::Owned)
.map_err(WyreError::Utf8Error)
}
}
impl<T: Pack> Pack for Vec<T> {
#[inline]
fn pack<B: BufMut>(&self, buf: &mut B) {
write_varint(buf, self.len() as u64);
for item in self {
item.pack(buf);
}
}
}
impl<T: Unpack> Unpack for Vec<T> {
#[inline]
fn unpack(buf: &mut Bytes) -> Result<Self, WyreError> {
let len = read_varint(buf)? as usize;
let mut vec = Vec::with_capacity(len);
for _ in 0..len {
vec.push(T::unpack(buf)?);
}
Ok(vec)
}
}
impl<T: Pack> Pack for VecDeque<T> {
#[inline]
fn pack<B: BufMut>(&self, buf: &mut B) {
write_varint(buf, self.len() as u64);
for item in self {
item.pack(buf);
}
}
}
impl<T: Unpack> Unpack for VecDeque<T> {
#[inline]
fn unpack(buf: &mut Bytes) -> Result<Self, WyreError> {
let len = read_varint(buf)? as usize;
let mut deque = VecDeque::with_capacity(len);
for _ in 0..len {
deque.push_back(T::unpack(buf)?);
}
Ok(deque)
}
}
impl<T: Pack> Pack for LinkedList<T> {
#[inline]
fn pack<B: BufMut>(&self, buf: &mut B) {
write_varint(buf, self.len() as u64);
for item in self {
item.pack(buf);
}
}
}
impl<T: Unpack> Unpack for LinkedList<T> {
#[inline]
fn unpack(buf: &mut Bytes) -> Result<Self, WyreError> {
let len = read_varint(buf)? as usize;
let mut list = LinkedList::new();
for _ in 0..len {
list.push_back(T::unpack(buf)?);
}
Ok(list)
}
}
impl<K: Pack + Eq + Hash, V: Pack> Pack for HashMap<K, V> {
#[inline]
fn pack<B: BufMut>(&self, buf: &mut B) {
write_varint(buf, self.len() as u64);
for (key, value) in self {
key.pack(buf);
value.pack(buf);
}
}
}
impl<K: Unpack + Eq + Hash, V: Unpack> Unpack for HashMap<K, V> {
#[inline]
fn unpack(buf: &mut Bytes) -> Result<Self, WyreError> {
let len = read_varint(buf)? as usize;
let mut map = HashMap::with_capacity_and_hasher(len, RandomState::new());
for _ in 0..len {
let key = K::unpack(buf)?;
let value = V::unpack(buf)?;
map.insert(key, value);
}
Ok(map)
}
}
impl<K: Pack + Ord, V: Pack> Pack for BTreeMap<K, V> {
#[inline]
fn pack<B: BufMut>(&self, buf: &mut B) {
write_varint(buf, self.len() as u64);
for (key, value) in self {
key.pack(buf);
value.pack(buf);
}
}
}
impl<K: Unpack + Ord, V: Unpack> Unpack for BTreeMap<K, V> {
#[inline]
fn unpack(buf: &mut Bytes) -> Result<Self, WyreError> {
let len = read_varint(buf)? as usize;
let mut map = BTreeMap::new();
for _ in 0..len {
let key = K::unpack(buf)?;
let value = V::unpack(buf)?;
map.insert(key, value);
}
Ok(map)
}
}
impl<T: Pack + Eq + Hash> Pack for HashSet<T> {
#[inline]
fn pack<B: BufMut>(&self, buf: &mut B) {
write_varint(buf, self.len() as u64);
for item in self {
item.pack(buf);
}
}
}
impl<T: Unpack + Eq + Hash> Unpack for HashSet<T> {
#[inline]
fn unpack(buf: &mut Bytes) -> Result<Self, WyreError> {
let len = read_varint(buf)? as usize;
let mut set = HashSet::with_capacity_and_hasher(len, RandomState::new());
for _ in 0..len {
set.insert(T::unpack(buf)?);
}
Ok(set)
}
}
impl<T: Pack + Ord> Pack for BTreeSet<T> {
#[inline]
fn pack<B: BufMut>(&self, buf: &mut B) {
write_varint(buf, self.len() as u64);
for item in self {
item.pack(buf);
}
}
}
impl<T: Unpack + Ord> Unpack for BTreeSet<T> {
#[inline]
fn unpack(buf: &mut Bytes) -> Result<Self, WyreError> {
let len = read_varint(buf)? as usize;
let mut set = BTreeSet::new();
for _ in 0..len {
set.insert(T::unpack(buf)?);
}
Ok(set)
}
}
impl<T: Pack> Pack for Option<T> {
#[inline]
fn pack<B: BufMut>(&self, buf: &mut B) {
match self {
Some(value) => {
buf.put_u8(1);
value.pack(buf);
}
None => buf.put_u8(0),
}
}
}
impl<T: Unpack> Unpack for Option<T> {
#[inline]
fn unpack(buf: &mut Bytes) -> Result<Self, WyreError> {
match buf.get_u8() {
0 => Ok(None),
1 => Ok(Some(T::unpack(buf)?)),
_ => Err(WyreError::InvalidFormat("Invalid Option tag".into())),
}
}
}
impl<T: Pack, E: Pack> Pack for Result<T, E> {
#[inline]
fn pack<B: BufMut>(&self, buf: &mut B) {
match self {
Ok(value) => {
buf.put_u8(0);
value.pack(buf);
}
Err(err) => {
buf.put_u8(1);
err.pack(buf);
}
}
}
}
impl<T: Unpack, E: Unpack> Unpack for Result<T, E> {
#[inline]
fn unpack(buf: &mut Bytes) -> Result<Self, WyreError> {
match buf.get_u8() {
0 => Ok(Ok(T::unpack(buf)?)),
1 => Ok(Err(E::unpack(buf)?)),
_ => Err(WyreError::InvalidFormat("Invalid Result tag".into())),
}
}
}
impl Pack for () {
#[inline]
fn pack<B: BufMut>(&self, _buf: &mut B) {}
}
impl Unpack for () {
#[inline]
fn unpack(_buf: &mut Bytes) -> Result<Self, WyreError> {
Ok(())
}
}
impl<T: Pack> Pack for Box<T> {
#[inline]
fn pack<B: BufMut>(&self, buf: &mut B) {
(**self).pack(buf)
}
}
impl<T: Unpack> Unpack for Box<T> {
#[inline]
fn unpack(buf: &mut Bytes) -> Result<Self, WyreError> {
Ok(Box::new(T::unpack(buf)?))
}
}
impl<T: Pack> Pack for [T; 1] {
#[inline]
fn pack<B: BufMut>(&self, buf: &mut B) {
self[0].pack(buf);
}
}
impl<T: Unpack> Unpack for [T; 1] {
#[inline]
fn unpack(buf: &mut Bytes) -> Result<Self, WyreError> {
Ok([T::unpack(buf)?])
}
}
macro_rules! impl_tuple {
($($t:ident),+) => {
impl<$($t: Pack),+> Pack for ($($t,)+) {
#[inline]
fn pack<B: BufMut>(&self, buf: &mut B) {
let ($($t,)+) = self;
$(
$t.pack(buf);
)+
}
}
impl<$($t: Unpack),+> Unpack for ($($t,)+) {
#[inline]
fn unpack(buf: &mut Bytes) -> Result<Self, WyreError> {
Ok((
$(
$t::unpack(buf)?,
)+
))
}
}
}
}
impl_tuple!(T1, T2);
impl_tuple!(T1, T2, T3);
impl_tuple!(T1, T2, T3, T4);
impl_tuple!(T1, T2, T3, T4, T5);
impl_tuple!(T1, T2, T3, T4, T5, T6);
impl_tuple!(T1, T2, T3, T4, T5, T6, T7);
impl_tuple!(T1, T2, T3, T4, T5, T6, T7, T8);
impl_tuple!(T1, T2, T3, T4, T5, T6, T7, T8, T9);
impl_tuple!(T1, T2, T3, T4, T5, T6, T7, T8, T9, T10);
impl_tuple!(T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11);
impl_tuple!(T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12);
impl Pack for Ipv4Addr {
#[inline]
fn pack<B: BufMut>(&self, buf: &mut B) {
buf.put_slice(&self.octets());
}
}
impl Unpack for Ipv4Addr {
#[inline]
fn unpack(buf: &mut Bytes) -> Result<Self, WyreError> {
if buf.remaining() < 4 {
return Err(WyreError::InvalidFormat(
"Not enough bytes for Ipv4Addr".into(),
));
}
let octets = [buf.get_u8(), buf.get_u8(), buf.get_u8(), buf.get_u8()];
Ok(Ipv4Addr::from(octets))
}
}
impl Pack for Ipv6Addr {
#[inline]
fn pack<B: BufMut>(&self, buf: &mut B) {
buf.put_slice(&self.octets());
}
}
impl Unpack for Ipv6Addr {
#[inline]
fn unpack(buf: &mut Bytes) -> Result<Self, WyreError> {
if buf.remaining() < 16 {
return Err(WyreError::InvalidFormat(
"Not enough bytes for Ipv6Addr".into(),
));
}
let mut octets = [0u8; 16];
buf.copy_to_slice(&mut octets);
Ok(Ipv6Addr::from(octets))
}
}
impl Pack for IpAddr {
#[inline]
fn pack<B: BufMut>(&self, buf: &mut B) {
match self {
IpAddr::V4(addr) => {
buf.put_u8(4);
addr.pack(buf);
}
IpAddr::V6(addr) => {
buf.put_u8(6);
addr.pack(buf);
}
}
}
}
impl Unpack for IpAddr {
#[inline]
fn unpack(buf: &mut Bytes) -> Result<Self, WyreError> {
match buf.get_u8() {
4 => Ok(IpAddr::V4(Ipv4Addr::unpack(buf)?)),
6 => Ok(IpAddr::V6(Ipv6Addr::unpack(buf)?)),
_ => Err(WyreError::InvalidFormat("Invalid IpAddr tag".into())),
}
}
}
impl Pack for SocketAddr {
#[inline]
fn pack<B: BufMut>(&self, buf: &mut B) {
match self {
SocketAddr::V4(addr) => {
buf.put_u8(4);
addr.ip().pack(buf);
buf.put_u16_le(addr.port());
}
SocketAddr::V6(addr) => {
buf.put_u8(6);
addr.ip().pack(buf);
buf.put_u16_le(addr.port());
}
}
}
}
impl Unpack for SocketAddr {
#[inline]
fn unpack(buf: &mut Bytes) -> Result<Self, WyreError> {
match buf.get_u8() {
4 => {
let ip = Ipv4Addr::unpack(buf)?;
let port = buf.get_u16_le();
Ok(SocketAddr::V4(std::net::SocketAddrV4::new(ip, port)))
}
6 => {
let ip = Ipv6Addr::unpack(buf)?;
let port = buf.get_u16_le();
Ok(SocketAddr::V6(std::net::SocketAddrV6::new(ip, port, 0, 0)))
}
_ => Err(WyreError::InvalidFormat("Invalid SocketAddr tag".into())),
}
}
}
impl Pack for Duration {
#[inline]
fn pack<B: BufMut>(&self, buf: &mut B) {
self.as_secs().pack(buf);
self.subsec_nanos().pack(buf);
}
}
impl Unpack for Duration {
#[inline]
fn unpack(buf: &mut Bytes) -> Result<Self, WyreError> {
let secs = u64::unpack(buf)?;
let nanos = u32::unpack(buf)?;
Ok(Duration::new(secs, nanos))
}
}
impl Pack for SystemTime {
#[inline]
fn pack<B: BufMut>(&self, buf: &mut B) {
self.duration_since(UNIX_EPOCH)
.expect("SystemTime before UNIX EPOCH")
.pack(buf);
}
}
impl Unpack for SystemTime {
#[inline]
fn unpack(buf: &mut Bytes) -> Result<Self, WyreError> {
let duration = Duration::unpack(buf)?;
Ok(UNIX_EPOCH + duration)
}
}
pub fn to_bytes<T: Pack>(value: &T) -> BytesMut {
let mut buf = BytesMut::new();
value.pack(&mut buf);
buf
}
pub fn from_bytes<T: Unpack>(bytes: &[u8]) -> Result<T, WyreError> {
let mut buf = Bytes::copy_from_slice(bytes);
T::unpack(&mut buf)
}
#[cfg(test)]
mod tests {
use super::*;
use std::collections::BTreeSet;
use std::net::{Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6};
use std::time::{Duration, SystemTime, UNIX_EPOCH};
#[test]
fn test_primitives() {
let values = vec![
0i8,
-1i8,
127i8,
-128i8,
0i16 as i8,
-1i16 as i8,
32767i16 as i8,
];
for &value in &values {
let packed = to_bytes(&value);
let unpacked = from_bytes(&packed).unwrap();
assert_eq!(value, unpacked);
}
}
#[test]
fn test_floats_f32() {
let values = vec![
0.0f32,
1.0f32,
-1.0f32,
f32::MIN,
f32::MAX,
f32::INFINITY,
f32::NEG_INFINITY,
f32::NAN,
];
for &value in &values {
let packed = to_bytes(&value);
let unpacked: f32 = from_bytes(&packed).unwrap();
if value.is_nan() {
assert!(unpacked.is_nan());
} else {
assert_eq!(value, unpacked);
}
}
}
#[test]
fn test_floats_f64() {
let values = vec![
0.0f64,
1.0f64,
-1.0f64,
f64::MIN,
f64::MAX,
f64::INFINITY,
f64::NEG_INFINITY,
f64::NAN,
];
for &value in &values {
let packed = to_bytes(&value);
let unpacked: f64 = from_bytes(&packed).unwrap();
if value.is_nan() {
assert!(unpacked.is_nan());
} else {
assert_eq!(value, unpacked);
}
}
}
#[test]
fn test_bool() {
assert_eq!(from_bytes::<bool>(&to_bytes(&true)).unwrap(), true);
assert_eq!(from_bytes::<bool>(&to_bytes(&false)).unwrap(), false);
}
#[test]
fn test_char() {
let chars = vec!['a', 'Z', '0', '🦀', '한'];
for &c in &chars {
assert_eq!(from_bytes::<char>(&to_bytes(&c)).unwrap(), c);
}
}
#[test]
fn test_string() {
let strings = vec!["", "hello", "world", "こんにちは", "🌍🌎🌏"];
for s in &strings {
assert_eq!(from_bytes::<String>(&to_bytes(s)).unwrap(), *s);
}
}
#[test]
fn test_option() {
let values: Vec<Option<i32>> = vec![None, Some(0), Some(42), Some(-42)];
for &value in &values {
assert_eq!(from_bytes::<Option<i32>>(&to_bytes(&value)).unwrap(), value);
}
}
#[test]
fn test_result() {
let values: Vec<Result<i32, String>> =
vec![Ok(0), Ok(42), Ok(-42), Err("error".to_string())];
for value in &values {
assert_eq!(
from_bytes::<Result<i32, String>>(&to_bytes(value)).unwrap(),
*value
);
}
}
#[test]
fn test_vec() {
let vec: Vec<i32> = vec![1, 2, 3, 4, 5];
assert_eq!(from_bytes::<Vec<i32>>(&to_bytes(&vec)).unwrap(), vec);
}
#[test]
fn test_hashmap() {
let mut map = HashMap::new();
map.insert("one".to_string(), 1);
map.insert("two".to_string(), 2);
map.insert("three".to_string(), 3);
let unpacked: HashMap<String, i32> = from_bytes(&to_bytes(&map)).unwrap();
assert_eq!(map, unpacked);
}
#[test]
fn test_btreemap() {
let mut map = BTreeMap::new();
map.insert("one".to_string(), 1);
map.insert("two".to_string(), 2);
map.insert("three".to_string(), 3);
let unpacked: BTreeMap<String, i32> = from_bytes(&to_bytes(&map)).unwrap();
assert_eq!(map, unpacked);
}
#[test]
fn test_hashset() {
let mut set = HashSet::new();
set.insert("one".to_string());
set.insert("two".to_string());
set.insert("three".to_string());
let unpacked: HashSet<String> = from_bytes(&to_bytes(&set)).unwrap();
assert_eq!(set, unpacked);
}
#[test]
fn test_btreeset() {
let mut set = BTreeSet::new();
set.insert("one".to_string());
set.insert("two".to_string());
set.insert("three".to_string());
let unpacked: BTreeSet<String> = from_bytes(&to_bytes(&set)).unwrap();
assert_eq!(set, unpacked);
}
#[test]
fn test_derive() {
use crate as wyre;
#[derive(Debug, PartialEq, Eq, Pack, Unpack)]
struct Test {
a: u8,
b: u16,
c: u32,
d: u64,
e: String,
f: Vec<u8>,
g: HashMap<String, u32>,
h: BTreeMap<String, u32>,
i: HashSet<String>,
j: BTreeSet<String>,
k: Option<u32>,
l: Result<u32, String>,
}
let test = Test {
a: 1,
b: 2,
c: 3,
d: 4,
e: "hello".to_string(),
f: vec![1, 2, 3, 4, 5],
g: HashMap::from_iter(vec![("one".to_string(), 1), ("two".to_string(), 2), ("three".to_string(), 3)]),
h: BTreeMap::from_iter(vec![("one".to_string(), 1), ("two".to_string(), 2), ("three".to_string(), 3)]),
i: HashSet::from_iter(vec!["one".to_string(), "two".to_string(), "three".to_string()]),
j: BTreeSet::from_iter(vec!["one".to_string(), "two".to_string(), "three".to_string()]),
k: None,
l: Ok(42),
};
let packed = to_bytes(&test);
let unpacked: Test = from_bytes(&packed).unwrap();
assert_eq!(test, unpacked);
}
}