use std::{
fmt::Debug,
mem::{ManuallyDrop, MaybeUninit},
ops::{Deref, DerefMut},
};
#[derive(Copy, Clone)]
pub(crate) enum Allocator {
Default,
#[cfg(any(target_os = "linux", target_os = "android"))]
Mmap,
}
pub struct Buffer {
pub(crate) ptr: *mut u8,
pub(crate) len: u32,
pub(crate) requested_len: u32,
pub(crate) capacity: u32,
pub(crate) allocator: Allocator,
}
impl Buffer {
#[inline]
pub fn new(requested_len: usize) -> Self {
let len_u32 = requested_len.try_into().expect("length overflow");
let mut vec = ManuallyDrop::new(Vec::with_capacity(requested_len));
Buffer {
ptr: vec.as_mut_ptr(),
len: 0,
requested_len: len_u32,
capacity: vec.capacity().try_into().expect("capacity overflow"),
allocator: Allocator::Default,
}
}
#[cfg(any(target_os = "linux", target_os = "android"))]
pub(crate) fn mmap(
fd: &std::os::unix::prelude::OwnedFd,
len: usize,
) -> Result<Buffer, rustix::io::Errno> {
let len_u32 = len.try_into().expect("length overflow");
let ptr = unsafe {
rustix::mm::mmap(
std::ptr::null_mut(),
len,
rustix::mm::ProtFlags::READ | rustix::mm::ProtFlags::WRITE,
rustix::mm::MapFlags::SHARED,
fd,
0,
)
}?;
Ok(Buffer {
ptr: ptr as *mut u8,
len: 0,
requested_len: len_u32,
capacity: len_u32,
allocator: Allocator::Mmap,
})
}
#[inline]
pub fn len(&self) -> usize {
self.len as usize
}
#[inline]
pub fn is_empty(&self) -> bool {
self.len == 0
}
#[inline]
pub fn requested_len(&self) -> usize {
self.requested_len as usize
}
#[inline]
pub fn set_requested_len(&mut self, len: usize) {
assert!(len <= self.capacity as usize, "length exceeds capacity");
self.requested_len = len.try_into().expect("requested_len overflow");
}
#[inline]
pub fn capacity(&self) -> usize {
self.capacity as usize
}
#[inline]
pub fn remaining_capacity(&self) -> usize {
self.capacity() - self.len()
}
#[inline]
pub fn clear(&mut self) {
self.len = 0;
}
pub fn extend_fill(&mut self, len: usize, value: u8) -> &mut [u8] {
assert!(len <= self.remaining_capacity(), "length exceeds capacity");
unsafe {
std::ptr::write_bytes(self.ptr.add(self.len()), value, len);
}
self.len += len as u32;
unsafe { std::slice::from_raw_parts_mut(self.ptr.add(self.len() - len), len) }
}
pub fn extend_from_slice(&mut self, slice: &[u8]) {
assert!(
slice.len() <= self.remaining_capacity(),
"length exceeds capacity"
);
unsafe {
std::ptr::copy_nonoverlapping(slice.as_ptr(), self.ptr.add(self.len()), slice.len());
}
self.len += slice.len() as u32;
}
pub fn is_zero_copy(&self) -> bool {
!matches!(self.allocator, Allocator::Default)
}
pub fn into_vec(self) -> Vec<u8> {
match self.allocator {
Allocator::Default => {
let buf = ManuallyDrop::new(self);
unsafe { Vec::from_raw_parts(buf.ptr, buf.len as usize, buf.capacity as usize) }
}
#[allow(unreachable_patterns)]
_ => self[..].to_vec(),
}
}
}
unsafe impl Send for Buffer {}
unsafe impl Sync for Buffer {}
impl From<Vec<u8>> for Buffer {
fn from(vec: Vec<u8>) -> Self {
let mut vec = ManuallyDrop::new(vec);
Buffer {
ptr: vec.as_mut_ptr(),
len: vec.len().try_into().expect("len overflow"),
requested_len: vec.len().try_into().expect("len overflow"),
capacity: vec.capacity().try_into().expect("capacity overflow"),
allocator: Allocator::Default,
}
}
}
impl From<&[u8]> for Buffer {
fn from(slice: &[u8]) -> Self {
Self::from(slice.to_vec())
}
}
impl<const N: usize> From<[u8; N]> for Buffer {
fn from(array: [u8; N]) -> Self {
Self::from(array.to_vec())
}
}
impl From<Vec<MaybeUninit<u8>>> for Buffer {
fn from(vec: Vec<MaybeUninit<u8>>) -> Self {
let mut vec = ManuallyDrop::new(vec);
Buffer {
ptr: vec.as_mut_ptr().cast(),
len: 0,
requested_len: vec.len().try_into().expect("len overflow"),
capacity: vec.capacity().try_into().expect("capacity overflow"),
allocator: Allocator::Default,
}
}
}
impl Deref for Buffer {
type Target = [u8];
fn deref(&self) -> &[u8] {
unsafe { std::slice::from_raw_parts(self.ptr, self.len as usize) }
}
}
impl DerefMut for Buffer {
fn deref_mut(&mut self) -> &mut [u8] {
unsafe { std::slice::from_raw_parts_mut(self.ptr, self.len as usize) }
}
}
impl Debug for Buffer {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Buffer")
.field("len", &self.len)
.field("requested_len", &self.requested_len)
.field("data", &format_args!("{:02x?}", &self[..]))
.finish()
}
}
impl Drop for Buffer {
fn drop(&mut self) {
match self.allocator {
Allocator::Default => unsafe {
drop(Vec::from_raw_parts(
self.ptr,
self.len as usize,
self.capacity as usize,
));
},
#[cfg(any(target_os = "linux", target_os = "android"))]
Allocator::Mmap => unsafe {
rustix::mm::munmap(self.ptr as *mut _, self.capacity as usize).unwrap();
},
}
}
}