use core::alloc::{GlobalAlloc, Layout};
use core::cell::UnsafeCell;
use core::mem::{self, MaybeUninit};
use core::ptr::{NonNull, null_mut};
#[cfg(not(feature = "polyfill"))]
use core::sync::atomic::{AtomicUsize, Ordering};
#[cfg(feature = "polyfill")]
use portable_atomic::{AtomicUsize, Ordering};
use crate::leaked::LeakBox;
use alloc_traits::{AllocTime, LocalAlloc, NonZeroLayout};
#[repr(C)]
pub struct Bump<T> {
header: Header,
storage: UnsafeCell<MaybeUninit<T>>,
}
#[repr(C)]
pub struct BumpSlice {
header: Header,
storage: UnsafeCell<[MaybeUninit<u8>]>,
}
#[derive(Clone, Copy)]
struct BumpView<'lt> {
header: &'lt Header,
storage: &'lt UnsafeCell<[MaybeUninit<u8>]>,
}
#[repr(C)]
struct Header {
consumed: AtomicUsize,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub struct LeakError<T> {
val: T,
failure: Failure,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct Level(pub(crate) usize);
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
pub struct Allocation<'a, T: ?Sized = u8> {
pub ptr: NonNull<T>,
pub lifetime: AllocTime<'a>,
pub level: Level,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum Failure {
Exhausted,
Mismatch {
observed: Level,
},
}
impl<T> Bump<T> {
pub const fn uninit() -> Self {
Bump {
header: Header::empty(),
storage: UnsafeCell::new(MaybeUninit::uninit()),
}
}
pub fn zeroed() -> Self {
Bump {
header: Header::empty(),
storage: UnsafeCell::new(MaybeUninit::zeroed()),
}
}
pub const fn new(storage: T) -> Self {
Bump {
header: Header::empty(),
storage: UnsafeCell::new(MaybeUninit::new(storage)),
}
}
pub const fn as_bump_slice(&self) -> Option<&BumpSlice> {
if mem::offset_of!(Self, storage) != mem::size_of::<Header>() {
return None;
}
let data_len = mem::size_of::<T>();
let ptr = (self as *const Self).cast::<MaybeUninit<u8>>();
let mem: *const [MaybeUninit<u8>] = core::ptr::slice_from_raw_parts(ptr, data_len);
Some(unsafe { &*(mem as *const BumpSlice) })
}
pub const fn as_mut_bump_slice(&mut self) -> Option<&mut BumpSlice> {
if mem::offset_of!(Self, storage) != mem::size_of::<Header>() {
return None;
}
let data_len = mem::size_of::<T>();
let ptr = (self as *mut Self).cast::<MaybeUninit<u8>>();
let mem: *mut [MaybeUninit<u8>] = core::ptr::slice_from_raw_parts_mut(ptr, data_len);
Some(unsafe { &mut *(mem as *mut BumpSlice) })
}
pub fn reset(&mut self) {
self.header = Header::empty();
}
fn as_view(&self) -> BumpView<'_> {
BumpView {
header: &self.header,
storage: {
let base = self.storage.get();
let len = mem::size_of::<T>();
let data = core::ptr::slice_from_raw_parts(base as *const _, len);
unsafe { &*(data as *const UnsafeCell<[_]>) }
},
}
}
pub fn alloc(&self, layout: Layout) -> Option<NonNull<u8>> {
self.as_view().alloc(layout)
}
pub fn alloc_at(&self, layout: Layout, level: Level) -> Result<Allocation<'_>, Failure> {
self.as_view().alloc_at(layout, level)
}
pub fn get_layout(&self, layout: Layout) -> Option<Allocation<'_>> {
self.as_view().get_layout(layout)
}
pub fn get_layout_at(&self, layout: Layout, at: Level) -> Result<Allocation<'_>, Failure> {
self.as_view().get_layout_at(layout, at)
}
pub fn get<V>(&self) -> Option<Allocation<'_, V>> {
self.as_view().get()
}
pub fn get_at<V>(&self, level: Level) -> Result<Allocation<'_, V>, Failure> {
self.as_view().get_at(level)
}
pub fn get_slice<V>(&self, len: usize) -> Option<Allocation<'_, [V]>> {
self.as_view().get_slice(len)
}
pub fn leak_box<V>(&self, val: V) -> Option<LeakBox<'_, V>> {
self.as_view().leak_box(val)
}
pub fn leak_box_at<V>(&self, val: V, level: Level) -> Result<LeakBox<'_, V>, Failure> {
self.as_view().leak_box_at(val, level)
}
pub fn level(&self) -> Level {
self.as_view().level()
}
pub unsafe fn get_unchecked<V>(&self, level: Level) -> Allocation<'_, V> {
unsafe { self.as_view().get_unchecked(level) }
}
#[expect(clippy::mut_from_ref)] pub fn leak<V>(&self, val: V) -> Result<&mut V, LeakError<V>> {
match self.get::<V>() {
Some(alloc) => Ok(unsafe { alloc.leak(val) }),
None => Err(LeakError::new(val, Failure::Exhausted)),
}
}
#[expect(clippy::mut_from_ref)] pub fn leak_at<V>(&self, val: V, level: Level) -> Result<(&mut V, Level), LeakError<V>> {
let alloc = match self.get_at::<V>(level) {
Ok(alloc) => alloc,
Err(err) => return Err(LeakError::new(val, err)),
};
let level = alloc.level;
let mutref = unsafe { alloc.leak(val) };
Ok((mutref, level))
}
}
impl BumpSlice {
pub fn reset(&mut self) {
self.header = Header::empty();
}
fn as_view(&self) -> BumpView<'_> {
BumpView {
header: &self.header,
storage: &self.storage,
}
}
pub fn alloc(&self, layout: Layout) -> Option<NonNull<u8>> {
self.as_view().alloc(layout)
}
pub fn alloc_at(&self, layout: Layout, level: Level) -> Result<Allocation<'_>, Failure> {
self.as_view().alloc_at(layout, level)
}
pub fn get_layout(&self, layout: Layout) -> Option<Allocation<'_>> {
self.as_view().get_layout(layout)
}
pub fn get_layout_at(&self, layout: Layout, at: Level) -> Result<Allocation<'_>, Failure> {
self.as_view().get_layout_at(layout, at)
}
pub fn get<V>(&self) -> Option<Allocation<'_, V>> {
self.as_view().get()
}
pub fn get_at<V>(&self, level: Level) -> Result<Allocation<'_, V>, Failure> {
self.as_view().get_at(level)
}
pub fn get_slice<V>(&self, len: usize) -> Option<Allocation<'_, [V]>> {
self.as_view().get_slice(len)
}
pub fn leak_box<V>(&self, val: V) -> Option<LeakBox<'_, V>> {
self.as_view().leak_box(val)
}
pub fn leak_box_at<V>(&self, val: V, level: Level) -> Result<LeakBox<'_, V>, Failure> {
self.as_view().leak_box_at(val, level)
}
pub fn level(&self) -> Level {
self.as_view().level()
}
pub unsafe fn get_unchecked<V>(&self, level: Level) -> Allocation<'_, V> {
unsafe { self.as_view().get_unchecked(level) }
}
#[expect(clippy::mut_from_ref)] pub fn leak<V>(&self, val: V) -> Result<&mut V, LeakError<V>> {
match self.get::<V>() {
Some(alloc) => Ok(unsafe { alloc.leak(val) }),
None => Err(LeakError::new(val, Failure::Exhausted)),
}
}
#[expect(clippy::mut_from_ref)] pub fn leak_at<V>(&self, val: V, level: Level) -> Result<(&mut V, Level), LeakError<V>> {
let alloc = match self.get_at::<V>(level) {
Ok(alloc) => alloc,
Err(err) => return Err(LeakError::new(val, err)),
};
let level = alloc.level;
let mutref = unsafe { alloc.leak(val) };
Ok((mutref, level))
}
}
impl<'lt> BumpView<'lt> {
pub fn alloc(self, layout: Layout) -> Option<NonNull<u8>> {
Some(self.try_alloc(layout)?.ptr)
}
pub fn alloc_at(self, layout: Layout, level: Level) -> Result<Allocation<'lt>, Failure> {
let Allocation {
ptr,
lifetime,
level,
} = self.try_alloc_at(layout, level.0)?;
Ok(Allocation {
ptr: ptr.cast(),
lifetime,
level,
})
}
pub fn get_layout(self, layout: Layout) -> Option<Allocation<'lt>> {
self.try_alloc(layout)
}
pub fn get_layout_at(self, layout: Layout, at: Level) -> Result<Allocation<'lt>, Failure> {
self.try_alloc_at(layout, at.0)
}
pub fn get<V>(self) -> Option<Allocation<'lt, V>> {
if mem::size_of::<V>() == 0 {
return Some(self.zst_fake_alloc());
}
let layout = Layout::new::<V>();
let Allocation {
ptr,
lifetime,
level,
} = self.try_alloc(layout)?;
Some(Allocation {
ptr: ptr.cast(),
lifetime,
level,
})
}
pub fn get_at<V>(self, level: Level) -> Result<Allocation<'lt, V>, Failure> {
if mem::size_of::<V>() == 0 {
let fake = self.zst_fake_alloc();
if fake.level != level {
return Err(Failure::Mismatch {
observed: fake.level,
});
}
return Ok(fake);
}
let layout = Layout::new::<V>();
let Allocation {
ptr,
lifetime,
level,
} = self.try_alloc_at(layout, level.0)?;
Ok(Allocation {
ptr: ptr.cast(),
lifetime,
level,
})
}
pub fn get_slice<V>(&self, len: usize) -> Option<Allocation<'lt, [V]>> {
if len == 0 {
return Some(Allocation::for_empty_slice(self.level()));
}
let (layout, _) = Layout::new::<V>().repeat(len).ok()?;
if layout.size() == 0 {
return Some(Allocation::for_zst_slice(len, self.level()));
};
let alloc = self.get_layout(layout)?;
Some(Allocation {
ptr: NonNull::slice_from_raw_parts(alloc.ptr.cast(), len),
lifetime: alloc.lifetime,
level: alloc.level,
})
}
pub fn leak_box<V>(self, val: V) -> Option<LeakBox<'lt, V>> {
let Allocation { ptr, lifetime, .. } = self.get::<V>()?;
Some(unsafe { LeakBox::new_from_raw_non_null(ptr, val, lifetime) })
}
pub fn leak_box_at<V>(self, val: V, level: Level) -> Result<LeakBox<'lt, V>, Failure> {
let Allocation { ptr, lifetime, .. } = self.get_at::<V>(level)?;
Ok(unsafe { LeakBox::new_from_raw_non_null(ptr, val, lifetime) })
}
pub fn level(&self) -> Level {
Level(self.header.consumed.load(Ordering::SeqCst))
}
pub unsafe fn get_unchecked<V>(self, level: Level) -> Allocation<'lt, V> {
debug_assert!(level.0 <= mem::size_of_val(self.storage));
debug_assert!(
level <= self.level(),
"Tried to access an allocation that does not yet exist"
);
let base_ptr = self.storage.get().cast::<u8>();
let alloc = unsafe { base_ptr.add(level.0) };
let ptr = NonNull::new(alloc).unwrap().cast::<V>();
debug_assert!(
ptr.as_ptr().is_aligned(),
"Tried to access an allocation with improper type"
);
Allocation {
level,
lifetime: AllocTime::default(),
ptr,
}
}
fn try_alloc(self, layout: Layout) -> Option<Allocation<'lt>> {
let mut consumed = 0;
loop {
match self.try_alloc_at(layout, consumed) {
Ok(alloc) => return Some(alloc),
Err(Failure::Exhausted) => return None,
Err(Failure::Mismatch { observed }) => consumed = observed.0,
}
}
}
fn try_alloc_at(
self,
layout: Layout,
expect_consumed: usize,
) -> Result<Allocation<'lt>, Failure> {
assert!(layout.size() > 0);
let length = self.storage.get().len();
let base_ptr = self.storage.get().cast::<u8>();
let alignment = layout.align();
let requested = layout.size();
assert!(expect_consumed <= length);
let available = length.checked_sub(expect_consumed).unwrap();
let ptr_to = base_ptr.wrapping_add(expect_consumed);
let offset = ptr_to.align_offset(alignment);
if requested > available.saturating_sub(offset) {
return Err(Failure::Exhausted); }
assert!(offset < available);
let at_aligned = expect_consumed.checked_add(offset).unwrap();
let new_consumed = at_aligned.checked_add(requested).unwrap();
assert!(new_consumed <= length);
assert!(at_aligned < length);
match self.bump(expect_consumed, new_consumed) {
Ok(()) => (),
Err(observed) => {
return Err(Failure::Mismatch {
observed: Level(observed),
});
}
}
let aligned = unsafe {
base_ptr.byte_add(at_aligned)
};
Ok(Allocation {
ptr: NonNull::new(aligned).unwrap(),
lifetime: AllocTime::default(),
level: Level(new_consumed),
})
}
fn zst_fake_alloc<Z>(&self) -> Allocation<'lt, Z> {
Allocation::for_zst(self.level())
}
fn bump(&self, expect_consumed: usize, new_consumed: usize) -> Result<(), usize> {
assert!(expect_consumed <= new_consumed);
assert!(new_consumed <= self.storage.get().len());
self.header.bump(expect_consumed, new_consumed)
}
}
impl Header {
const fn empty() -> Self {
Header {
consumed: AtomicUsize::new(0),
}
}
fn bump(&self, expect_consumed: usize, new_consumed: usize) -> Result<(), usize> {
self.consumed
.compare_exchange(
expect_consumed,
new_consumed,
Ordering::SeqCst,
Ordering::SeqCst,
)
.map(drop)
}
}
impl<'alloc, T> Allocation<'alloc, T> {
pub unsafe fn leak(self, val: T) -> &'alloc mut T {
unsafe { core::ptr::write(self.ptr.as_ptr(), val) };
unsafe { &mut *self.ptr.as_ptr() }
}
pub unsafe fn boxed(self, val: T) -> LeakBox<'alloc, T> {
unsafe { core::ptr::write(self.ptr.as_ptr(), val) };
unsafe { LeakBox::from_raw(self.ptr.as_ptr()) }
}
pub unsafe fn uninit(self) -> &'alloc mut MaybeUninit<T> {
unsafe { &mut *self.ptr.cast().as_ptr() }
}
pub(crate) fn for_zst(level: Level) -> Self {
assert!(mem::size_of::<T>() == 0);
let alloc: &[T; 0] = &[];
Allocation {
ptr: NonNull::from(alloc).cast(),
lifetime: AllocTime::default(),
level,
}
}
pub(crate) fn for_zst_slice(len: usize, level: Level) -> Allocation<'alloc, [T]> {
assert!(mem::size_of::<T>() == 0);
let alloc: &[T; 0] = &[];
Allocation {
ptr: NonNull::slice_from_raw_parts(NonNull::from(alloc).cast(), len),
lifetime: AllocTime::default(),
level,
}
}
pub(crate) fn for_empty_slice(level: Level) -> Allocation<'alloc, [T]> {
let alloc: &[T; 0] = &[];
Allocation {
ptr: NonNull::from(alloc),
lifetime: AllocTime::default(),
level,
}
}
}
impl<T> LeakError<T> {
fn new(val: T, failure: Failure) -> Self {
LeakError { val, failure }
}
pub fn kind(&self) -> Failure {
self.failure
}
pub fn into_inner(self) -> T {
self.val
}
}
unsafe impl<T> Sync for Bump<T> {}
unsafe impl Sync for BumpView<'_> {}
unsafe impl Send for BumpView<'_> {}
unsafe impl<T> GlobalAlloc for Bump<T> {
unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
unsafe { GlobalAlloc::alloc(&self.as_view(), layout) }
}
unsafe fn realloc(&self, ptr: *mut u8, current: Layout, new_size: usize) -> *mut u8 {
unsafe { GlobalAlloc::realloc(&self.as_view(), ptr, current, new_size) }
}
unsafe fn dealloc(&self, _ptr: *mut u8, _layout: Layout) {
}
}
unsafe impl GlobalAlloc for &'static BumpSlice {
unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
unsafe { GlobalAlloc::alloc(&self.as_view(), layout) }
}
unsafe fn realloc(&self, ptr: *mut u8, current: Layout, new_size: usize) -> *mut u8 {
unsafe { GlobalAlloc::realloc(&self.as_view(), ptr, current, new_size) }
}
unsafe fn dealloc(&self, _ptr: *mut u8, _layout: Layout) {
}
}
unsafe impl GlobalAlloc for BumpView<'_> {
unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
BumpView::alloc(*self, layout)
.map(NonNull::as_ptr)
.unwrap_or_else(null_mut)
}
unsafe fn realloc(&self, ptr: *mut u8, current: Layout, new_size: usize) -> *mut u8 {
let current = NonZeroLayout::from_layout(current.into()).unwrap();
let new_size = unsafe { core::num::NonZeroUsize::new_unchecked(new_size) };
let target = match layout_reallocated(current, new_size) {
Some(target) => target,
None => return core::ptr::null_mut(),
};
let reconstructed = alloc_traits::Allocation {
ptr: unsafe { NonNull::new_unchecked(ptr) },
layout: current,
lifetime: AllocTime::default(),
};
unsafe { alloc_traits::LocalAlloc::realloc(self, reconstructed, target) }
.map(|alloc| alloc.ptr.as_ptr())
.unwrap_or_else(core::ptr::null_mut)
}
unsafe fn dealloc(&self, _ptr: *mut u8, _layout: Layout) {
}
}
fn layout_reallocated(
layout: NonZeroLayout,
target: core::num::NonZeroUsize,
) -> Option<NonZeroLayout> {
let layout = Layout::from_size_align(target.get(), layout.align()).ok()?;
Some(NonZeroLayout::from_layout(layout.into()).unwrap())
}
unsafe impl<'alloc, T> LocalAlloc<'alloc> for Bump<T> {
fn alloc(&'alloc self, layout: NonZeroLayout) -> Option<alloc_traits::Allocation<'alloc>> {
let raw_alloc = self.get_layout(layout.into())?;
Some(alloc_traits::Allocation {
ptr: raw_alloc.ptr,
layout,
lifetime: AllocTime::default(),
})
}
unsafe fn realloc(
&'alloc self,
alloc: alloc_traits::Allocation<'alloc>,
layout: NonZeroLayout,
) -> Option<alloc_traits::Allocation<'alloc>> {
if alloc.ptr.as_ptr() as usize % layout.align() == 0 && alloc.layout.size() >= layout.size()
{
return Some(alloc_traits::Allocation {
ptr: alloc.ptr,
layout,
lifetime: alloc.lifetime,
});
}
let new_alloc = LocalAlloc::alloc(self, layout)?;
unsafe {
core::ptr::copy_nonoverlapping(
alloc.ptr.as_ptr(),
new_alloc.ptr.as_ptr(),
layout.size().min(alloc.layout.size()).into(),
);
}
Some(new_alloc)
}
unsafe fn dealloc(&'alloc self, _: alloc_traits::Allocation<'alloc>) {
}
}
unsafe impl<'alloc> LocalAlloc<'alloc> for BumpView<'alloc> {
fn alloc(&'alloc self, layout: NonZeroLayout) -> Option<alloc_traits::Allocation<'alloc>> {
let raw_alloc = self.get_layout(layout.into())?;
Some(alloc_traits::Allocation {
ptr: raw_alloc.ptr,
layout,
lifetime: AllocTime::default(),
})
}
unsafe fn realloc(
&'alloc self,
alloc: alloc_traits::Allocation<'alloc>,
layout: NonZeroLayout,
) -> Option<alloc_traits::Allocation<'alloc>> {
if alloc.ptr.as_ptr() as usize % layout.align() == 0 && alloc.layout.size() >= layout.size()
{
return Some(alloc_traits::Allocation {
ptr: alloc.ptr,
layout,
lifetime: alloc.lifetime,
});
}
let new_alloc = LocalAlloc::alloc(self, layout)?;
unsafe {
core::ptr::copy_nonoverlapping(
alloc.ptr.as_ptr(),
new_alloc.ptr.as_ptr(),
layout.size().min(alloc.layout.size()).into(),
);
}
Some(new_alloc)
}
unsafe fn dealloc(&'alloc self, _: alloc_traits::Allocation<'alloc>) {
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn zst_no_drop() {
#[derive(Debug)]
struct PanicOnDrop;
impl Drop for PanicOnDrop {
fn drop(&mut self) {
panic!("No instance of this should ever get dropped");
}
}
let alloc = Bump::<()>::uninit();
let _ = alloc.leak(PanicOnDrop).unwrap();
}
}