use crate::bytes::default_controller::{self, NativeAllocationController};
use alloc::boxed::Box;
use alloc::vec::Vec;
use core::alloc::LayoutError;
use core::mem::MaybeUninit;
use core::ops::{Deref, DerefMut};
use core::ptr::NonNull;
pub struct Bytes {
controller: Box<dyn AllocationController>,
len: usize,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AllocationProperty {
File,
Native,
Pinned,
Other,
}
#[derive(Debug, Clone, Copy)]
pub enum SplitError {
InvalidOffset,
Unsupported,
}
pub trait AllocationController {
fn alloc_align(&self) -> usize;
fn property(&self) -> AllocationProperty;
unsafe fn memory_mut(&mut self) -> &mut [MaybeUninit<u8>];
fn memory(&self) -> &[MaybeUninit<u8>];
#[allow(clippy::type_complexity)]
fn split(
&mut self,
_offset: usize,
) -> Result<(Box<dyn AllocationController>, Box<dyn AllocationController>), SplitError> {
Err(SplitError::Unsupported)
}
fn duplicate(&self) -> Option<Box<dyn AllocationController>> {
None
}
unsafe fn copy_into(&self, buf: &mut [u8]) {
let len = buf.len();
let memory = self.memory();
let memory_slice = &memory[0..len];
let data = unsafe {
core::slice::from_raw_parts(memory_slice.as_ptr().cast(), memory_slice.len())
};
buf.copy_from_slice(data);
}
#[allow(unused_variables)]
fn grow(&mut self, size: usize, align: usize) -> Result<(), AllocationError> {
Err(AllocationError::UnsupportedOperation)
}
fn try_detach(&mut self) -> Option<NonNull<u8>> {
None
}
}
#[derive(Debug, Clone, PartialEq)]
pub enum AllocationError {
UnsupportedOperation,
OutOfMemory,
}
impl Bytes {
pub fn split(mut self, offset: usize) -> Result<(Bytes, Bytes), (Bytes, SplitError)> {
let right_len = self.len - offset;
match self.controller.split(offset) {
Ok((left, right)) => unsafe {
Ok((
Bytes::from_controller(left, offset),
Bytes::from_controller(right, right_len),
))
},
Err(err) => match self.try_into_vec() {
Ok(mut left) => {
let right = left.split_off(offset);
Ok((Bytes::from_bytes_vec(left), Bytes::from_bytes_vec(right)))
}
Err(this) => Err((this, err)),
},
}
}
#[cfg(feature = "std")]
pub fn from_file<P: Into<std::path::PathBuf>>(file: P, size: u64, offset: u64) -> Self {
let controller = crate::bytes::file::FileAllocationController::new(file, size, offset);
Self {
controller: Box::new(controller),
len: size as usize,
}
}
#[cfg(feature = "shared-bytes")]
pub fn from_shared(bytes: bytes::Bytes, property: AllocationProperty) -> Self {
let len = bytes.len();
let controller =
crate::bytes::shared::SharedBytesAllocationController::new(bytes, property);
Self {
controller: Box::new(controller),
len,
}
}
#[allow(clippy::len_without_is_empty)]
pub fn len(&self) -> usize {
self.len
}
pub fn copy_into(&self, other: &mut Self) {
unsafe {
self.controller.copy_into(other);
}
}
pub fn property(&self) -> AllocationProperty {
self.controller.property()
}
pub unsafe fn from_controller(controller: Box<dyn AllocationController>, len: usize) -> Self {
debug_assert!(
len <= controller.memory().len(),
"len must not exceed controller memory size"
);
Self { controller, len }
}
pub fn from_bytes_vec(bytes: Vec<u8>) -> Self {
let mut bytes = Self::from_elems(bytes);
bytes
.try_enforce_runtime_align(default_controller::MAX_ALIGN)
.unwrap();
bytes
}
pub fn from_elems<E>(elems: Vec<E>) -> Self
where
E: bytemuck::NoUninit + Send + Sync,
{
let _: () = const {
assert!(
core::mem::align_of::<E>() <= default_controller::MAX_ALIGN,
"element type not supported due to too large alignment"
);
};
let byte_len = elems.len() * core::mem::size_of::<E>();
let controller = NativeAllocationController::from_elems(elems);
Self {
controller: Box::new(controller),
len: byte_len,
}
}
pub fn extend_from_byte_slice(&mut self, bytes: &[u8]) {
self.extend_from_byte_slice_aligned(bytes, default_controller::MAX_ALIGN)
}
pub fn capacity(&self) -> usize {
self.controller.memory().len()
}
pub fn try_into_vec<E: bytemuck::CheckedBitPattern + bytemuck::NoUninit>(
mut self,
) -> Result<Vec<E>, Self> {
let Ok(data) = bytemuck::checked::try_cast_slice::<_, E>(&self) else {
return Err(self);
};
let length = data.len();
let byte_capacity = self.controller.memory().len();
let Some(capacity) = byte_capacity.checked_div(size_of::<E>()) else {
return Err(self);
};
if capacity * size_of::<E>() != byte_capacity {
return Err(self);
};
if self.controller.alloc_align() != align_of::<E>() {
return Err(self);
}
let Some(ptr) = self.controller.try_detach() else {
return Err(self);
};
let vec = unsafe { Vec::from_raw_parts(ptr.as_ptr().cast(), length, capacity) };
Ok(vec)
}
pub fn align(&self) -> usize {
self.controller.alloc_align()
}
pub fn extend_from_byte_slice_aligned(&mut self, bytes: &[u8], align: usize) {
debug_assert!(align.is_power_of_two(), "alignment must be a power of two");
debug_assert!(
align <= default_controller::MAX_ALIGN,
"alignment exceeds maximum supported alignment"
);
let additional = bytes.len();
self.reserve(additional, align);
let len = self.len();
let new_cap = len.wrapping_add(additional); debug_assert!(
new_cap <= self.capacity(),
"new capacity must not exceed allocated capacity"
);
unsafe {
let uninit_spare = &mut self.controller.memory_mut()[len..new_cap];
uninit_spare.copy_from_slice(core::slice::from_raw_parts(
bytes.as_ptr().cast(),
additional,
));
};
self.len = new_cap;
}
fn try_from_data(align: usize, data: &[u8]) -> Result<Self, LayoutError> {
let controller = NativeAllocationController::alloc_with_data(data, align)?;
Ok(Self {
controller: Box::new(controller),
len: data.len(),
})
}
fn try_enforce_runtime_align(&mut self, align: usize) -> Result<(), LayoutError> {
if self.controller.alloc_align() >= align {
return Ok(());
}
*self = Self::try_from_data(align, self)?;
Ok(())
}
fn reserve(&mut self, additional: usize, align: usize) {
debug_assert!(
align <= default_controller::MAX_ALIGN,
"alignment exceeds maximum supported alignment"
);
let needs_to_grow = additional > self.capacity().wrapping_sub(self.len());
if !needs_to_grow {
return;
}
let Some(required_cap) = self.len().checked_add(additional) else {
default_controller::alloc_overflow()
};
let new_cap = required_cap.max(self.capacity() * 2);
let new_cap = new_cap.max(align);
match self.controller.grow(new_cap, align) {
Ok(()) => {}
Err(_err) => {
let new_controller: Box<dyn AllocationController> = Box::new(
NativeAllocationController::alloc_with_capacity(new_cap, align).unwrap(),
);
let mut new_bytes = Self {
controller: new_controller,
len: self.len,
};
new_bytes.copy_from_slice(&*self);
*self = new_bytes;
}
}
}
}
impl Deref for Bytes {
type Target = [u8];
fn deref(&self) -> &Self::Target {
let memory = &self.controller.memory()[0..self.len];
unsafe { core::slice::from_raw_parts(memory.as_ptr().cast(), memory.len()) }
}
}
impl DerefMut for Bytes {
fn deref_mut(&mut self) -> &mut Self::Target {
let len = self.len;
let slice = unsafe { &mut self.controller.memory_mut() };
let memory = &mut slice[0..len];
unsafe { core::slice::from_raw_parts_mut(memory.as_mut_ptr().cast(), memory.len()) }
}
}
unsafe impl Send for Bytes {}
unsafe impl Sync for Bytes {}
fn debug_from_fn<F: Fn(&mut core::fmt::Formatter<'_>) -> core::fmt::Result>(
f: F,
) -> impl core::fmt::Debug {
struct FromFn<F>(F);
impl<F> core::fmt::Debug for FromFn<F>
where
F: Fn(&mut core::fmt::Formatter<'_>) -> core::fmt::Result,
{
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
(self.0)(f)
}
}
FromFn(f)
}
impl core::fmt::Debug for Bytes {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
let data = &**self;
let fmt_data = move |f: &mut core::fmt::Formatter<'_>| {
if data.len() > 3 {
f.debug_list().entries(&data[0..3]).entry(&"...").finish()
} else {
f.debug_list().entries(data).finish()
}
};
f.debug_struct("Bytes")
.field("data", &debug_from_fn(fmt_data))
.field("len", &self.len)
.finish()
}
}
impl serde::Serialize for Bytes {
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where
S: serde::Serializer,
{
serde_bytes::serialize(self.deref(), serializer)
}
}
impl<'de> serde::Deserialize<'de> for Bytes {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: serde::Deserializer<'de>,
{
#[cold]
fn too_large<E: serde::de::Error>(len: usize, align: usize) -> E {
let max_length = (isize::MAX as usize) & !(align - 1);
E::custom(core::format_args!(
"length too large: {len}. Expected at most {max_length} bytes"
))
}
let data: Vec<u8> = serde_bytes::deserialize(deserializer)?;
let align = default_controller::MAX_ALIGN;
let mut bytes = Self::from_elems(data);
bytes
.try_enforce_runtime_align(align)
.map_err(|_| too_large(bytes.len(), align))?;
Ok(bytes)
}
}
impl Clone for Bytes {
fn clone(&self) -> Self {
if let Some(controller) = self.controller.duplicate() {
return Self {
controller,
len: self.len,
};
}
Self::try_from_data(self.align(), self.deref()).unwrap()
}
}
impl PartialEq for Bytes {
fn eq(&self, other: &Self) -> bool {
self.deref() == other.deref()
}
}
impl Eq for Bytes {}
#[cfg(test)]
mod tests {
use super::Bytes;
use alloc::{vec, vec::Vec};
const _CONST_ASSERTS: fn() = || {
fn test_send<T: Send>() {}
fn test_sync<T: Sync>() {}
test_send::<Bytes>();
test_sync::<Bytes>();
};
fn test_serialization_roundtrip(bytes: &Bytes) {
let mut serialized = Vec::new();
ciborium::ser::into_writer(bytes, &mut serialized).expect("serialization to succeed");
let roundtripped: Bytes = ciborium::de::from_reader(&mut serialized.as_slice())
.expect("deserialization to succeed");
assert_eq!(
bytes, &roundtripped,
"roundtripping through serialization didn't lead to equal Bytes"
);
}
#[test_log::test]
fn test_serialization() {
test_serialization_roundtrip(&Bytes::from_elems::<i32>(vec![]));
test_serialization_roundtrip(&Bytes::from_elems(vec![0xdead, 0xbeaf]));
}
#[test_log::test]
fn test_into_vec() {
let mut bytes = Vec::with_capacity(6);
let actual_cap = bytes.capacity();
bytes.extend_from_slice(&[0, 1, 2, 3]);
let mut bytes = Bytes::from_elems::<u8>(bytes);
bytes = bytes
.try_into_vec::<[u8; 0]>()
.expect_err("Conversion should not succeed for a zero-sized type");
if actual_cap % 4 != 0 {
bytes = bytes.try_into_vec::<[u8; 4]>().err().unwrap_or_else(|| {
panic!("Conversion should not succeed due to capacity {actual_cap} not fitting a whole number of elements");
});
}
bytes = bytes
.try_into_vec::<u16>()
.expect_err("Conversion should not succeed due to mismatched alignment");
bytes = bytes.try_into_vec::<[u8; 3]>().expect_err(
"Conversion should not succeed due to size not fitting a whole number of elements",
);
let bytes = bytes.try_into_vec::<[u8; 2]>().expect("Conversion should succeed for bit-convertible types of equal alignment and compatible size");
assert_eq!(bytes, &[[0, 1], [2, 3]]);
}
#[test_log::test]
fn test_grow() {
let mut bytes = Bytes::from_elems::<u8>(vec![]);
bytes.extend_from_byte_slice(&[0, 1, 2, 3]);
assert_eq!(bytes[..], [0, 1, 2, 3][..]);
let mut bytes = Bytes::from_elems(vec![42u8; 4]);
bytes.extend_from_byte_slice(&[0, 1, 2, 3]);
assert_eq!(bytes[..], [42, 42, 42, 42, 0, 1, 2, 3][..]);
}
#[test_log::test]
fn test_large_elems() {
let mut bytes = Bytes::from_elems(vec![42u128]);
const TEST_BYTES: [u8; 16] = [
0x12, 0x90, 0x78, 0x56, 0x34, 0x12, 0x90, 0x78, 0x56, 0x34, 0x12, 0x90, 0x78, 0x56,
0x34, 0x12,
];
bytes.extend_from_byte_slice(&TEST_BYTES);
let vec = bytes.try_into_vec::<u128>().unwrap();
assert_eq!(vec, [42u128, u128::from_ne_bytes(TEST_BYTES)]);
}
#[test_log::test]
fn test_split_and_use() {
let bytes = Bytes::from_elems(vec![0u8, 1, 2, 3, 4, 5, 6, 7]);
let (left, right) = bytes.split(4).unwrap();
assert_eq!(&left[..], &[0, 1, 2, 3]);
assert_eq!(&right[..], &[4, 5, 6, 7]);
let left2 = left.clone();
assert_eq!(&left2[..], &[0, 1, 2, 3]);
}
#[test_log::test]
fn test_split_at_zero() {
let bytes = Bytes::from_elems(vec![10u8, 20, 30, 40]);
let (left, right) = bytes.split(0).unwrap();
assert_eq!(left.len(), 0);
assert_eq!(&right[..], &[10, 20, 30, 40]);
}
#[test_log::test]
fn test_split_at_end() {
let bytes = Bytes::from_elems(vec![10u8, 20, 30, 40]);
let (left, right) = bytes.split(4).unwrap();
assert_eq!(&left[..], &[10, 20, 30, 40]);
assert_eq!(right.len(), 0);
}
#[test_log::test]
fn test_from_bytes_vec_try_into_vec_aligned_type() {
for _ in 0..64 {
let bytes = Bytes::from_bytes_vec(vec![0u8; 16]);
let vec: Vec<u128> = bytes
.try_into_vec::<u128>()
.expect("MAX_ALIGN-aligned bytes must convert to Vec<u128>");
assert_eq!(vec.len(), 1);
}
}
#[test_log::test]
fn test_many_extends_with_growth() {
let mut bytes = Bytes::from_elems::<u8>(vec![]);
for i in 0u8..=255 {
bytes.extend_from_byte_slice(&[i]);
}
assert_eq!(bytes.len(), 256);
assert_eq!(bytes[0], 0);
assert_eq!(bytes[255], 255);
}
}