use bytes::{Buf, BufMut, Bytes, BytesMut};
use std::{
collections::VecDeque,
io::{Error, IoSlice},
};
use crate::interfaces::{Reader, Writer};
pub const ERR_EOB: &str = "No more bytes left to be read in buffer";
pub const ERR_EOM: &str = "Buffer is full, cannot write more bytes";
pub const ERR_VARINT_TOO_LONG: &str = "Varint is too long to be written to buffer";
macro_rules! can_read {
($self: ident, $size: expr) => {
$self.buf.remaining() >= $size
};
}
macro_rules! can_write {
($self: ident, $size: expr) => {
$self.buf.remaining_mut() >= $size
};
}
macro_rules! read_fn {
($name: ident, $typ: ident, $fn_name: ident, $byte_size: literal) => {
#[inline]
pub fn $name(&mut self) -> Result<$typ, std::io::Error> {
if can_read!(self, $byte_size) {
return Ok(self.buf.$fn_name());
} else {
return Err(Error::new(std::io::ErrorKind::UnexpectedEof, ERR_EOB));
}
}
};
}
macro_rules! write_fn {
($name: ident, $typ: ident, $fn_name: ident, $byte_size: literal) => {
#[inline]
pub fn $name(&mut self, num: $typ) -> Result<(), std::io::Error> {
if can_write!(self, $byte_size) {
self.buf.$fn_name(num);
return Ok(());
} else {
return Err(Error::new(std::io::ErrorKind::OutOfMemory, ERR_EOM));
}
}
};
}
#[derive(Debug, Clone)]
pub struct ByteReader {
pub(crate) buf: Bytes,
}
impl From<ByteWriter> for ByteReader {
fn from(writer: ByteWriter) -> Self {
Self {
buf: writer.buf.freeze(),
}
}
}
impl Into<Bytes> for ByteReader {
fn into(self) -> Bytes {
self.buf
}
}
impl Into<Vec<u8>> for ByteReader {
fn into(self) -> Vec<u8> {
self.buf.to_vec()
}
}
impl Into<VecDeque<u8>> for ByteReader {
fn into(self) -> VecDeque<u8> {
self.buf.to_vec().into()
}
}
impl From<Bytes> for ByteReader {
fn from(buf: Bytes) -> Self {
Self { buf }
}
}
impl From<Vec<u8>> for ByteReader {
fn from(buf: Vec<u8>) -> Self {
Self { buf: buf.into() }
}
}
impl From<&[u8]> for ByteReader {
fn from(buf: &[u8]) -> Self {
Self {
buf: Bytes::from(buf.to_vec()),
}
}
}
impl ByteReader {
pub fn peek_ahead(&mut self, pos: usize) -> Result<u8, std::io::Error> {
if can_read!(self, pos) {
return Ok(self.buf.chunk()[pos]);
} else {
return Err(Error::new(std::io::ErrorKind::UnexpectedEof, ERR_EOB));
}
}
read_fn!(read_u8, u8, get_u8, 1);
read_fn!(read_i8, i8, get_i8, 1);
read_fn!(read_u16, u16, get_u16, 2);
read_fn!(read_u16_le, u16, get_u16_le, 2);
read_fn!(read_i16, i16, get_i16, 2);
read_fn!(read_i16_le, i16, get_i16_le, 2);
pub fn read_u24(&mut self) -> Result<u32, std::io::Error> {
if can_read!(self, 3) {
if let Ok(num) = self.read_uint(3) {
dbg!(num);
return Ok(num as u32);
} else {
return Err(Error::new(std::io::ErrorKind::UnexpectedEof, ERR_EOB));
}
} else {
return Err(Error::new(std::io::ErrorKind::UnexpectedEof, ERR_EOB));
}
}
pub fn read_u24_le(&mut self) -> Result<u32, std::io::Error> {
if can_read!(self, 3) {
if let Ok(num) = self.read_uint_le(3) {
return Ok(num as u32);
} else {
return Err(Error::new(std::io::ErrorKind::UnexpectedEof, ERR_EOB));
}
} else {
return Err(Error::new(std::io::ErrorKind::UnexpectedEof, ERR_EOB));
}
}
pub fn read_i24(&mut self) -> Result<i32, std::io::Error> {
if can_read!(self, 3) {
if let Ok(num) = self.read_int(3) {
return Ok(num as i32);
} else {
return Err(Error::new(std::io::ErrorKind::UnexpectedEof, ERR_EOB));
}
} else {
return Err(Error::new(std::io::ErrorKind::UnexpectedEof, ERR_EOB));
}
}
pub fn read_i24_le(&mut self) -> Result<i32, std::io::Error> {
if can_read!(self, 3) {
if let Ok(num) = self.read_int_le(3) {
return Ok(num as i32);
} else {
return Err(Error::new(std::io::ErrorKind::UnexpectedEof, ERR_EOB));
}
} else {
return Err(Error::new(std::io::ErrorKind::UnexpectedEof, ERR_EOB));
}
}
read_fn!(read_u32, u32, get_u32, 4);
read_fn!(read_u32_le, u32, get_u32_le, 4);
read_fn!(read_f32, f32, get_f32, 4);
read_fn!(read_f32_le, f32, get_f32_le, 4);
#[inline]
pub fn read_var_u32(&mut self) -> Result<u32, std::io::Error> {
let mut num = 0u32;
let mut interval = 0_usize;
for i in (0..35).step_by(7) {
let byte = self.peek_ahead(interval)?;
num |= ((byte & 0x7F) as u32) << i;
interval += 1;
if byte & 0x80 == 0 {
self.buf.advance(interval);
return Ok(num);
}
}
return Err(Error::new(
std::io::ErrorKind::Other,
"Varint overflow's 32-bit integer",
));
}
read_fn!(read_i32, i32, get_i32, 4);
read_fn!(read_i32_le, i32, get_i32_le, 4);
pub fn read_var_i32(&mut self) -> Result<i32, std::io::Error> {
let num = self.read_var_u32()?;
Ok((num >> 1) as i32 ^ -((num & 1) as i32))
}
read_fn!(read_u64, u64, get_u64, 8);
read_fn!(read_u64_le, u64, get_u64_le, 8);
read_fn!(read_i64, i64, get_i64, 8);
read_fn!(read_i64_le, i64, get_i64_le, 8);
read_fn!(read_f64, f64, get_f64, 8);
read_fn!(read_f64_le, f64, get_f64_le, 8);
#[inline]
pub fn read_var_u64(&mut self) -> Result<u64, std::io::Error> {
let mut num = 0u64;
let mut interval = 0_usize;
for i in (0..70).step_by(7) {
let byte = self.peek_ahead(interval)?;
num |= ((byte & 0x7F) as u64) << i;
interval += 1;
if byte & 0x80 == 0 {
self.buf.advance(interval);
return Ok(num);
}
}
return Err(Error::new(
std::io::ErrorKind::Other,
"Varint overflow's 64-bit integer",
));
}
#[inline]
pub fn read_var_i64(&mut self) -> Result<i64, std::io::Error> {
let num = self.read_var_u64()?;
Ok((num >> 1) as i64 ^ -((num & 1) as i64))
}
read_fn!(read_u128, u128, get_u128, 16);
read_fn!(read_u128_le, u128, get_u128_le, 16);
read_fn!(read_i128, i128, get_i128, 16);
read_fn!(read_i128_le, i128, get_i128_le, 16);
pub fn read_uint(&mut self, size: usize) -> Result<u64, std::io::Error> {
if can_read!(self, size) {
return Ok(self.buf.get_uint(size));
} else {
return Err(Error::new(std::io::ErrorKind::UnexpectedEof, ERR_EOB));
}
}
pub fn read_uint_le(&mut self, size: usize) -> Result<u64, std::io::Error> {
if can_read!(self, size) {
return Ok(self.buf.get_uint_le(size));
} else {
return Err(Error::new(std::io::ErrorKind::UnexpectedEof, ERR_EOB));
}
}
pub fn read_int(&mut self, size: usize) -> Result<i64, std::io::Error> {
if can_read!(self, size) {
return Ok(self.buf.get_int(size));
} else {
return Err(Error::new(std::io::ErrorKind::UnexpectedEof, ERR_EOB));
}
}
pub fn read_int_le(&mut self, size: usize) -> Result<i64, std::io::Error> {
if can_read!(self, size) {
return Ok(self.buf.get_int_le(size));
} else {
return Err(Error::new(std::io::ErrorKind::UnexpectedEof, ERR_EOB));
}
}
pub fn read_char(&mut self) -> Result<char, std::io::Error> {
let c = self.read_u32()?;
if let Some(c) = char::from_u32(c) {
return Ok(c);
} else {
return Err(Error::new(std::io::ErrorKind::InvalidData, "Invalid char"));
}
}
pub fn read_bool(&mut self) -> Result<bool, std::io::Error> {
if can_read!(self, 1) {
return Ok(self.buf.get_u8() != 0);
} else {
return Err(Error::new(std::io::ErrorKind::UnexpectedEof, ERR_EOB));
}
}
pub fn read_string(&mut self) -> Result<String, std::io::Error> {
let len = self.read_var_u64()?;
if can_read!(self, len as usize) {
let mut string = String::with_capacity(len as usize);
unsafe {
let v = string.as_mut_vec();
v.set_len(len as usize);
self.buf.copy_to_slice(&mut v[..]);
}
return Ok(string);
} else {
return Err(Error::new(std::io::ErrorKind::UnexpectedEof, ERR_EOB));
}
}
pub fn read_option<T: Reader<T>>(&mut self) -> Result<Option<T>, std::io::Error> {
if self.read_bool()? {
return Ok(Some(T::read(self)?));
} else {
return Ok(None);
}
}
pub fn read_sized_slice(&mut self) -> Result<Bytes, std::io::Error> {
let len = self.read_var_u32()?;
if can_read!(self, len as usize) {
let b = self.buf.slice(..len as usize);
self.buf.advance(len as usize);
return Ok(b);
} else {
return Err(Error::new(std::io::ErrorKind::UnexpectedEof, ERR_EOB));
}
}
pub fn read(&mut self, buffer: &mut [u8]) -> Result<(), std::io::Error> {
if can_read!(self, buffer.len()) {
self.buf.copy_to_slice(buffer);
return Ok(());
} else {
return Err(Error::new(std::io::ErrorKind::UnexpectedEof, ERR_EOB));
}
}
#[deprecated(note = "Use `read_type` instead")]
pub fn read_struct<T: Reader<T>>(&mut self) -> Result<T, std::io::Error> {
return self.read_type::<T>();
}
pub fn read_type<T: Reader<T>>(&mut self) -> Result<T, std::io::Error> {
return T::read(self);
}
pub fn as_slice(&self) -> &[u8] {
self.buf.chunk()
}
}
#[derive(Debug, Clone)]
pub struct ByteWriter {
pub(crate) buf: BytesMut,
}
impl Into<BytesMut> for ByteWriter {
fn into(self) -> BytesMut {
self.buf
}
}
impl Into<Bytes> for ByteWriter {
fn into(self) -> Bytes {
self.buf.freeze()
}
}
impl Into<Vec<u8>> for ByteWriter {
fn into(self) -> Vec<u8> {
self.buf.to_vec()
}
}
impl Into<VecDeque<u8>> for ByteWriter {
fn into(self) -> VecDeque<u8> {
self.buf.to_vec().into()
}
}
impl From<IoSlice<'_>> for ByteWriter {
fn from(slice: IoSlice) -> Self {
let mut buf = BytesMut::with_capacity(slice.len());
buf.put_slice(&slice);
return Self { buf };
}
}
impl From<&[u8]> for ByteWriter {
fn from(slice: &[u8]) -> Self {
let mut buf = BytesMut::with_capacity(slice.len());
buf.put_slice(slice);
return Self { buf };
}
}
impl From<ByteReader> for ByteWriter {
fn from(reader: ByteReader) -> Self {
Self {
buf: reader.buf.chunk().into(),
}
}
}
impl ByteWriter {
pub fn new() -> Self {
return Self {
buf: BytesMut::new(),
};
}
write_fn!(write_u8, u8, put_u8, 1);
write_fn!(write_i8, i8, put_i8, 1);
write_fn!(write_u16, u16, put_u16, 2);
write_fn!(write_u16_le, u16, put_u16_le, 2);
write_fn!(write_i16, i16, put_i16, 2);
write_fn!(write_i16_le, i16, put_i16_le, 2);
pub fn write_u24<I: Into<u32>>(&mut self, num: I) -> Result<(), std::io::Error> {
return self.write_uint(num.into().into(), 3);
}
pub fn write_u24_le<I: Into<u32>>(&mut self, num: I) -> Result<(), std::io::Error> {
return self.write_uint_le(num.into().into(), 3);
}
pub fn write_i24<I: Into<i32>>(&mut self, num: I) -> Result<(), std::io::Error> {
return self.write_int(num.into().into(), 3);
}
pub fn write_i24_le<I: Into<i32>>(&mut self, num: I) -> Result<(), std::io::Error> {
return self.write_int_le(num.into().into(), 3);
}
write_fn!(write_u32, u32, put_u32, 4);
write_fn!(write_u32_le, u32, put_u32_le, 4);
write_fn!(write_i32, i32, put_i32, 4);
write_fn!(write_i32_le, i32, put_i32_le, 4);
write_fn!(write_f32, f32, put_f32, 4);
write_fn!(write_f32_le, f32, put_f32_le, 4);
pub fn write_var_u32(&mut self, num: u32) -> Result<(), std::io::Error> {
let mut x = num;
while x >= 0x80 {
self.write_u8((x as u8) | 0x80)?;
x >>= 7;
}
self.write_u8(x as u8)?;
return Ok(());
}
pub fn write_var_i32(&mut self, num: i32) -> Result<(), std::io::Error> {
return if num < 0 {
let num = num as u32;
self.write_var_u32(!(num << 1))
} else {
let num = num as u32;
self.write_var_u32(num << 1)
};
}
write_fn!(write_u64, u64, put_u64, 8);
write_fn!(write_u64_le, u64, put_u64_le, 8);
write_fn!(write_i64, i64, put_i64, 8);
write_fn!(write_i64_le, i64, put_i64_le, 8);
write_fn!(write_f64, f64, put_f64, 8);
write_fn!(write_f64_le, f64, put_f64_le, 8);
pub fn write_var_u64(&mut self, num: u64) -> Result<(), std::io::Error> {
let mut x = (num as u64) & u64::MAX;
for _ in (0..70).step_by(7) {
if x >> 7 == 0 {
self.write_u8(x as u8)?;
return Ok(());
} else {
self.write_u8(((x & 0x7F) | 0x80) as u8)?;
x >>= 7;
}
}
return Err(Error::new(
std::io::ErrorKind::InvalidData,
ERR_VARINT_TOO_LONG,
));
}
pub fn write_var_i64(&mut self, num: i64) -> Result<(), std::io::Error> {
return if num < 0 {
let num = num as u64;
self.write_var_u64(!(num << 1))
} else {
let num = num as u64;
self.write_var_u64(num << 1)
};
}
write_fn!(write_u128, u128, put_u128, 16);
write_fn!(write_u128_le, u128, put_u128_le, 16);
write_fn!(write_i128, i128, put_i128, 16);
write_fn!(write_i128_le, i128, put_i128_le, 16);
pub fn write_uint(&mut self, num: u64, size: usize) -> Result<(), std::io::Error> {
if can_write!(self, size) {
self.buf.put_uint(num, size);
return Ok(());
} else {
return Err(Error::new(std::io::ErrorKind::OutOfMemory, ERR_EOM));
}
}
pub fn write_uint_le(&mut self, num: u64, size: usize) -> Result<(), std::io::Error> {
if can_write!(self, size) {
self.buf.put_uint_le(num, size);
return Ok(());
} else {
return Err(Error::new(std::io::ErrorKind::OutOfMemory, ERR_EOM));
}
}
pub fn write_int(&mut self, num: i64, size: usize) -> Result<(), std::io::Error> {
if can_write!(self, size) {
self.buf.put_int(num, size);
return Ok(());
} else {
return Err(Error::new(std::io::ErrorKind::OutOfMemory, ERR_EOM));
}
}
pub fn write_int_le(&mut self, num: i64, size: usize) -> Result<(), std::io::Error> {
if can_write!(self, size) {
self.buf.put_int_le(num, size);
return Ok(());
} else {
return Err(Error::new(std::io::ErrorKind::OutOfMemory, ERR_EOM));
}
}
pub fn write_char(&mut self, c: char) -> Result<(), std::io::Error> {
self.write_u32(c as u32)
}
pub fn write_bool(&mut self, b: bool) -> Result<(), std::io::Error> {
if can_write!(self, 1) {
self.buf.put_u8(b as u8);
return Ok(());
} else {
return Err(Error::new(std::io::ErrorKind::OutOfMemory, ERR_EOM));
}
}
pub fn write_string(&mut self, string: &str) -> Result<(), std::io::Error> {
if can_write!(self, string.len()) {
self.write_var_u32(string.len() as u32)?;
self.buf.put_slice(string.as_bytes());
return Ok(());
} else {
return Err(Error::new(std::io::ErrorKind::OutOfMemory, ERR_EOM));
}
}
pub fn write_option(&mut self, option: &Option<impl Writer>) -> Result<(), std::io::Error> {
if let Some(option) = option {
self.write_bool(true)?;
option.write(self)?;
} else {
self.write_bool(false)?;
}
return Ok(());
}
pub fn write_slice(&mut self, slice: &[u8]) -> Result<(), std::io::Error> {
if can_write!(self, slice.len()) {
self.write_var_u32(slice.len() as u32)?;
self.buf.put_slice(slice);
return Ok(());
} else {
return Err(Error::new(std::io::ErrorKind::OutOfMemory, ERR_EOM));
}
}
pub fn write(&mut self, buf: &[u8]) -> Result<(), std::io::Error> {
if can_write!(self, buf.len()) {
self.buf.put_slice(buf);
return Ok(());
} else {
return Err(Error::new(std::io::ErrorKind::OutOfMemory, ERR_EOM));
}
}
pub fn write_type<T: Writer>(&mut self, t: &T) -> Result<(), std::io::Error> {
t.write(self)
}
pub fn as_slice(&self) -> &[u8] {
self.buf.chunk()
}
pub fn clear(&mut self) {
self.buf.clear();
}
}