use core::alloc::{GlobalAlloc, Layout};
use core::cell::UnsafeCell;
use core::mem::{self, MaybeUninit};
use core::ptr::{NonNull, null_mut};
use core::sync::atomic::{AtomicUsize, Ordering};
use super::{Box, FixedVec, Uninit};
use crate::rc::Rc;
pub struct Slab<T> {
consumed: AtomicUsize,
storage: UnsafeCell<MaybeUninit<T>>,
}
#[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(usize);
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
pub struct Allocation {
pub ptr: NonNull<u8>,
pub level: Level,
}
#[derive(Debug)]
pub struct UninitAllocation<'a, T=()> {
pub uninit: Uninit<'a, T>,
pub level: Level,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum Failure {
Exhausted,
Mismatch {
observed: Level,
},
}
impl<T> Slab<T> {
pub const fn uninit() -> Self {
Slab {
consumed: AtomicUsize::new(0),
storage: UnsafeCell::new(MaybeUninit::uninit()),
}
}
pub fn zeroed() -> Self {
Slab {
consumed: AtomicUsize::new(0),
storage: UnsafeCell::new(MaybeUninit::zeroed()),
}
}
pub const fn new(storage: T) -> Self {
Slab {
consumed: AtomicUsize::new(0),
storage: UnsafeCell::new(MaybeUninit::new(storage)),
}
}
pub fn alloc(&self, layout: Layout) -> Option<NonNull<u8>> {
Some(self.try_alloc(layout)?.uninit.as_non_null().cast())
}
pub fn alloc_at(&self, layout: Layout, level: Level)
-> Result<Allocation, Failure>
{
let UninitAllocation { uninit, level } = self.try_alloc_at(layout, level.0)?;
Ok(Allocation {
ptr: uninit.as_non_null().cast(),
level,
})
}
pub fn get_layout(&self, layout: Layout)
-> Option<UninitAllocation<'_>>
{
self.try_alloc(layout)
}
pub fn get_layout_at(&self, layout: Layout, at: Level)
-> Result<UninitAllocation<'_>, Failure>
{
self.try_alloc_at(layout, at.0)
}
pub fn get<V>(&self) -> Option<UninitAllocation<V>> {
if mem::size_of::<V>() == 0 {
return Some(self.zst_fake_alloc());
}
let layout = Layout::new::<V>();
let UninitAllocation { uninit, level, } = self.try_alloc(layout)?;
Some(UninitAllocation {
uninit: uninit.cast().ok().unwrap(),
level,
})
}
pub fn get_at<V>(&self, level: Level) -> Result<UninitAllocation<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 UninitAllocation { uninit, level, } = self.try_alloc_at(layout, level.0)?;
Ok(UninitAllocation {
uninit: uninit.cast().ok().unwrap(),
level,
})
}
pub fn boxed<V>(&self, val: V) -> Option<Box<'_, V>> {
let alloc = self.get::<V>()?;
Some(Box::new(val, alloc.uninit))
}
pub fn fixed_vec<V>(&self, capacity: usize) -> Option<FixedVec<'_, V>> {
let size = mem::size_of::<V>().checked_mul(capacity)?;
let layout = Layout::from_size_align(size, mem::align_of::<V>()).ok()?;
let uninit = if layout.size() == 0 {
Uninit::empty()
} else {
self.get_layout(layout)?
.uninit
.cast_slice()
.unwrap()
};
Some(FixedVec::new(uninit))
}
pub fn rc<V>(&self, val: V) -> Option<Rc<'_, V>> {
let alloc = self.get_layout(Rc::<V>::layout())?;
Some(Rc::new(val, alloc.uninit))
}
pub fn level(&self) -> Level {
Level(self.consumed.load(Ordering::SeqCst))
}
fn try_alloc(&self, layout: Layout)
-> Option<UninitAllocation<'_>>
{
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<UninitAllocation<'_>, Failure>
{
assert!(layout.size() > 0);
let length = mem::size_of::<T>();
let base_ptr = self.storage.get()
as *mut T
as *mut 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 as *mut u8).add(at_aligned)
};
let ptr = NonNull::new(aligned).unwrap();
let uninit = unsafe {
Uninit::from_memory(ptr.cast(), requested)
};
Ok(UninitAllocation {
uninit,
level: Level(new_consumed),
})
}
pub fn leak<V>(&self, val: V) -> Result<&mut V, LeakError<V>> {
match self.get::<V>() {
Some(alloc) => Ok(alloc.uninit.init(val)),
None => Err(LeakError::new(val, Failure::Exhausted)),
}
}
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 mutref = alloc.uninit.init(val);
Ok((mutref, alloc.level))
}
fn zst_fake_alloc<Z>(&self) -> UninitAllocation<Z> {
UninitAllocation {
uninit: Uninit::invent_for_zst(),
level: self.level(),
}
}
fn bump(&self, expect_consumed: usize, new_consumed: usize) -> Result<(), usize> {
assert!(expect_consumed <= new_consumed);
assert!(new_consumed <= mem::size_of::<T>());
let observed = self.consumed.compare_and_swap(
expect_consumed,
new_consumed,
Ordering::SeqCst);
if expect_consumed == observed {
Ok(())
} else {
Err(observed)
}
}
}
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 Slab<T> { }
unsafe impl<T> GlobalAlloc for Slab<T> {
unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
Slab::alloc(self, layout)
.map(NonNull::as_ptr)
.unwrap_or_else(null_mut)
}
unsafe fn dealloc(&self, _ptr: *mut u8, _layout: Layout) {
}
}
#[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 = Slab::<()>::uninit();
let _ = alloc.leak(PanicOnDrop).unwrap();
}
}