use core::mem::MaybeUninit;
use crate::error::{NetpackError, Result};
pub trait CheckedReader: Sized {
fn read_slice_checked(&mut self, size: usize) -> Result<&[u8]>;
#[inline(always)]
fn read_checked<T: ReadDataChecked>(&mut self) -> Result<T> {
T::read_checked(self)
}
}
pub trait CheckedWriter: Sized {
fn write_slice(&mut self, data: &[u8]) -> Result<()>;
#[inline(always)]
fn write<T: WriteDataChecked>(&mut self, value: T) -> Result<()> {
T::write_to_checked(value, self)
}
}
impl<'a> CheckedReader for &'a [u8] {
#[inline(always)]
fn read_slice_checked(&mut self, size: usize) -> Result<&'a [u8]> {
if self.len() < size {
return Err(NetpackError::BufferTooSmall {
provided_size: self.len(),
requried_size: size,
});
}
let (data, remaining) = self.split_at(size);
*self = remaining;
Ok(data)
}
}
impl<'a> CheckedWriter for &'a mut [u8] {
#[inline(always)]
fn write_slice(&mut self, data: &[u8]) -> Result<()> {
if self.len() < data.len() {
return Err(NetpackError::BufferTooSmall {
provided_size: self.len(),
requried_size: data.len(),
});
}
let this = core::mem::take(self);
let (write_buffer, remaining) = this.split_at_mut(data.len());
*self = remaining;
write_buffer.copy_from_slice(data);
Ok(())
}
}
pub trait ReadDataChecked: Sized {
fn read_checked(reader: &mut impl CheckedReader) -> Result<Self>;
}
pub trait WriteDataChecked {
fn write_to_checked(self, writer: &mut impl CheckedWriter) -> Result<()>;
}
impl<const N: usize> ReadDataChecked for [u8; N] {
#[cfg_attr(feature = "fast-rw", inline(always))]
fn read_checked(reader: &mut impl CheckedReader) -> Result<Self> {
let slice = reader.read_slice_checked(N)?;
unsafe { Ok(*slice.as_ptr().cast()) }
}
}
impl<const N: usize> WriteDataChecked for [u8; N] {
fn write_to_checked(self, writer: &mut impl CheckedWriter) -> Result<()> {
writer.write_slice(&self)
}
}
impl<const N: usize> ReadDataChecked for [u16; N] {
fn read_checked(reader: &mut impl CheckedReader) -> Result<Self> {
let mut read_buffer = reader.read_slice_checked(core::mem::size_of::<u16>() * N)?;
let mut write_buffer: [MaybeUninit<u16>; N] =
unsafe { MaybeUninit::uninit().assume_init() };
let read_ptr = &mut read_buffer;
for elem in write_buffer.iter_mut() {
elem.write(u16::read_checked(read_ptr)?);
}
Ok(unsafe { *write_buffer.as_ptr().cast() })
}
}
impl<const N: usize> WriteDataChecked for [u16; N] {
fn write_to_checked(self, writer: &mut impl CheckedWriter) -> Result<()> {
for v in self.iter() {
writer.write(*v)?;
}
Ok(())
}
}
impl ReadDataChecked for u8 {
fn read_checked(reader: &mut impl CheckedReader) -> Result<Self> {
reader.read_slice_checked(1).map(|m| m[0])
}
}
impl WriteDataChecked for u8 {
fn write_to_checked(self, writer: &mut impl CheckedWriter) -> Result<()> {
writer.write_slice(&[self])
}
}
macro_rules! impl_read_write_data {
($($t:ty),+) => {
$(
impl ReadDataChecked for $t {
#[inline(always)]
fn read_checked(reader: &mut impl CheckedReader) -> Result<$t> {
let data = reader.read_checked();
data.map(<$t>::from_be_bytes)
}
}
impl WriteDataChecked for $t {
#[inline(always)]
fn write_to_checked(self, writer: &mut impl CheckedWriter) -> Result<()> {
let data = self.to_be_bytes();
writer.write(data)
}
}
)*
};
}
impl_read_write_data! {
u16,
u32,
u64
}