use crate::byte_buffer_read::ByteBufferRead;
use crate::byte_buffer_write::ByteBufferWrite;
use std::{
alloc::{self, Layout},
cmp::max,
ptr, slice,
};
use crate::error::{ByteBufferError, Result};
use core::ptr::NonNull;
#[derive(Debug)]
pub struct ByteBuffer {
cap: usize,
length: usize,
cursor: usize,
pointer: NonNull<u8>,
}
unsafe impl Send for ByteBuffer {}
unsafe impl Sync for ByteBuffer {}
impl ByteBuffer {
pub const MAX_SIZE: usize = isize::MAX as usize;
pub const MIN_SIZE: usize = 8;
pub fn new() -> Result<Self> {
Self::with_capacity(Self::MIN_SIZE)
}
pub fn with_capacity(capacity: usize) -> Result<Self> {
if capacity == 0 {
return Err(ByteBufferError::MinCapacity);
} else if capacity > Self::MAX_SIZE {
return Err(ByteBufferError::MaxCapacity);
}
let layout = alloc::Layout::from_size_align(capacity, 1)
.map_err(|_| ByteBufferError::LayoutFailure { size: capacity })?;
let new_ptr = unsafe { alloc::alloc(layout) };
let pointer = match NonNull::new(new_ptr) {
Some(p) => p,
None => Err(ByteBufferError::AllocationFailure { size: capacity })?,
};
Ok(Self {
cap: layout.size(),
length: 0,
cursor: 0,
pointer,
})
}
pub fn resize(&mut self, capacity: usize) -> Result<&mut Self> {
if capacity == 0 {
return Err(ByteBufferError::MinCapacity);
} else if capacity > Self::MAX_SIZE {
return Err(ByteBufferError::MaxCapacity);
}
let new_layout = alloc::Layout::from_size_align(capacity, 1)
.map_err(|_| ByteBufferError::LayoutFailure { size: capacity })?;
let old_layout = alloc::Layout::from_size_align(self.cap, 1)
.map_err(|_| ByteBufferError::LayoutFailure { size: self.cap })?;
let new_ptr =
unsafe { alloc::realloc(self.pointer.as_ptr(), old_layout, new_layout.size()) };
let pointer = match NonNull::new(new_ptr) {
Some(p) => p,
None => Err(ByteBufferError::AllocationFailure { size: capacity })?,
};
if self.length >= capacity {
self.length = capacity;
if self.cursor >= self.length {
self.cursor = self.length;
}
}
self.cap = capacity;
self.pointer = pointer;
Ok(self)
}
pub fn expand(&mut self, amount: usize) -> Result<&mut Self> {
self.resize(
self.cap
.checked_add(amount)
.ok_or(ByteBufferError::MaxCapacity)?,
)
}
pub fn shrink_to(&mut self, min_capacity: usize) -> Result<&mut Self> {
if self.cap > min_capacity {
self.resize(max(min_capacity, self.length))?;
}
Ok(self)
}
pub fn shrink_to_fit(&mut self) -> Result<&mut Self> {
if self.capacity() > self.length {
self.resize(max(self.length, Self::MIN_SIZE))?;
}
Ok(self)
}
pub unsafe fn write_slice_unchecked(&mut self, source: &[u8]) -> &mut Self {
let source_length = source.len();
ptr::copy_nonoverlapping(
source.as_ptr(),
self.pointer.as_ptr().add(self.cursor),
source_length,
);
self.cursor += source.len();
self
}
pub fn write_slice(&mut self, source: &[u8]) -> Result<&mut Self> {
if self.cursor + source.len() > self.cap {
let capacity = (self.cursor + source.len())
.checked_next_power_of_two()
.ok_or(ByteBufferError::MaxCapacity)?;
self.resize(capacity)?;
}
unsafe {
self.write_slice_unchecked(source);
}
if self.cursor > self.length {
self.length += self.cursor - self.length
}
Ok(self)
}
pub fn write<T: ByteBufferWrite>(&mut self, source: T) -> Result<&mut Self> {
source.write_to_bytey_buffer(self)?;
Ok(self)
}
pub fn write_le<T: ByteBufferWrite>(&mut self, source: T) -> Result<&mut Self> {
source.write_to_bytey_buffer_le(self)?;
Ok(self)
}
pub fn write_be<T: ByteBufferWrite>(&mut self, source: T) -> Result<&mut Self> {
source.write_to_bytey_buffer_be(self)?;
Ok(self)
}
pub unsafe fn read_slice_unchecked(&mut self, size: usize) -> &[u8] {
let ret = slice::from_raw_parts(self.pointer.as_ptr().add(self.cursor), size);
self.cursor += size;
ret
}
pub fn read_slice(&mut self, size: usize) -> Result<&[u8]> {
if self.cursor + size > self.length {
return Err(ByteBufferError::ReadOutOfBounds {
length: self.length,
start: self.cursor,
end: self.cursor + size,
});
}
Ok(unsafe { self.read_slice_unchecked(size) })
}
pub fn read<T: ByteBufferRead>(&mut self) -> Result<T> {
T::read_from_bytey_buffer(self)
}
pub fn read_le<T: ByteBufferRead>(&mut self) -> Result<T> {
T::read_from_bytey_buffer_le(self)
}
pub fn read_be<T: ByteBufferRead>(&mut self) -> Result<T> {
T::read_from_bytey_buffer_be(self)
}
pub unsafe fn move_cursor_unchecked(&mut self, location: usize) -> &mut Self {
self.cursor = location;
self
}
pub fn move_cursor(&mut self, location: usize) -> Result<&mut Self> {
if location > self.length {
return Err(ByteBufferError::CursorOutOfBounds {
length: self.length,
cursor: location,
});
}
self.cursor = location;
Ok(self)
}
pub fn move_cursor_to_end(&mut self) -> &mut Self {
self.cursor = self.length;
self
}
pub fn move_cursor_to_start(&mut self) -> &mut Self {
self.cursor = 0;
self
}
pub fn truncate(&mut self, length: usize) -> Result<&mut Self> {
if length > self.length {
return Err(ByteBufferError::LengthOutOfBounds {
current: self.length,
new: length,
});
}
self.length = length;
if self.cursor > length {
self.cursor = length;
}
Ok(self)
}
pub fn length(&self) -> usize {
self.length
}
pub fn capacity(&self) -> usize {
self.cap
}
pub fn cursor(&self) -> usize {
self.cursor
}
pub unsafe fn pointer(&self) -> *const u8 {
self.pointer.as_ptr()
}
pub unsafe fn mut_pointer(&self) -> *mut u8 {
self.pointer.as_ptr()
}
#[inline]
pub fn is_empty(&self) -> bool {
self.length() == 0
}
#[inline]
pub fn read_to_buffer(&mut self, len: usize) -> Result<Self> {
let mut buffer = ByteBuffer::with_capacity(len)?;
let bytes = self.read_slice(len)?;
buffer.write_slice(bytes)?;
buffer.cursor = 0;
Ok(buffer)
}
pub fn as_slice(&mut self) -> &[u8] {
self.cursor = 0;
unsafe { self.read_slice_unchecked(self.length()) }
}
pub fn slice_from(&mut self, cursor: usize, size: usize) -> Result<&[u8]> {
if cursor + size > self.length {
return Err(ByteBufferError::ReadOutOfBounds {
length: self.length,
start: cursor,
end: cursor + size,
});
}
self.cursor = cursor;
Ok(unsafe { self.read_slice_unchecked(size) })
}
}
impl Drop for ByteBuffer {
fn drop(&mut self) {
unsafe {
let layout = Layout::array::<u8>(self.cap).unwrap();
alloc::dealloc(self.pointer.as_ptr(), layout);
}
}
}
impl Clone for ByteBuffer {
fn clone(&self) -> Self {
let layout = alloc::Layout::from_size_align(self.cap, 1).unwrap();
let pointer = unsafe { alloc::alloc(layout) };
unsafe {
ptr::copy(self.pointer.as_ptr(), pointer, self.length);
}
Self {
cap: self.cap,
length: self.length,
cursor: self.cursor,
pointer: NonNull::new(pointer).unwrap(),
}
}
}