use std::ptr::NonNull;
use std::sync::atomic::{AtomicUsize, Ordering};
use crate::err::{Error, Result};
pub const CACHE_LINE: usize = 64;
#[repr(C, align(64))]
pub struct BumpAllocator {
base: NonNull<u8>,
capacity: usize,
offset: AtomicUsize,
_pad: [u8; CACHE_LINE - 24],
}
unsafe impl Send for BumpAllocator {}
unsafe impl Sync for BumpAllocator {}
impl BumpAllocator {
pub unsafe fn init(ptr: *mut u8, capacity: usize) -> Result<&'static Self> {
if ptr.is_null() || capacity < std::mem::size_of::<Self>() * 2 {
return Err(Error::OutOfMemory);
}
if !(ptr as usize).is_multiple_of(CACHE_LINE) {
return Err(Error::Alignment);
}
#[allow(clippy::cast_ptr_alignment)]
let alloc = ptr.cast::<Self>();
let usable_capacity = capacity - std::mem::size_of::<Self>();
let data_start = NonNull::new_unchecked(ptr.add(std::mem::size_of::<Self>()));
std::ptr::write(
alloc,
Self {
base: data_start,
capacity: usable_capacity,
offset: AtomicUsize::new(0),
_pad: [0; CACHE_LINE - 24],
},
);
Ok(&*alloc)
}
pub unsafe fn from_existing(ptr: *mut u8) -> Result<&'static Self> {
if ptr.is_null() {
return Err(Error::Alignment);
}
#[allow(clippy::cast_ptr_alignment)]
Ok(&*ptr.cast::<Self>())
}
pub fn alloc<T>(&self) -> Result<NonNull<T>> {
let size = std::mem::size_of::<T>();
let align = std::mem::align_of::<T>().max(CACHE_LINE);
self.alloc_raw(size, align).map(NonNull::cast)
}
pub fn alloc_slice<T>(&self, count: usize) -> Result<NonNull<T>> {
let size = std::mem::size_of::<T>()
.checked_mul(count)
.ok_or(Error::OutOfMemory)?;
let align = std::mem::align_of::<T>().max(CACHE_LINE);
self.alloc_raw(size, align).map(NonNull::cast)
}
fn alloc_raw(&self, size: usize, align: usize) -> Result<NonNull<u8>> {
loop {
let current = self.offset.load(Ordering::Relaxed);
let base = self.base.as_ptr();
let ptr = unsafe { base.add(current) };
let padding = ptr.align_offset(align);
let next = current
.checked_add(padding)
.and_then(|n| n.checked_add(size))
.ok_or(Error::OutOfMemory)?;
if next > self.capacity {
return Err(Error::OutOfMemory);
}
if self
.offset
.compare_exchange_weak(current, next, Ordering::AcqRel, Ordering::Relaxed)
.is_ok()
{
return Ok(unsafe { NonNull::new_unchecked(ptr.add(padding)) });
}
std::hint::spin_loop();
}
}
pub unsafe fn reset(&self) {
self.offset.store(0, Ordering::Release);
}
#[inline]
pub fn used(&self) -> usize {
self.offset.load(Ordering::Acquire)
}
#[inline]
pub const fn capacity(&self) -> usize {
self.capacity
}
#[inline]
pub fn available(&self) -> usize {
self.capacity.saturating_sub(self.used())
}
}