#![allow(clippy::inline_always)]
#![forbid(unsafe_op_in_unsafe_fn)]
#![no_std]
extern crate alloc;
use alloc::alloc::{alloc, dealloc, handle_alloc_error, Layout};
use core::{
cell::Cell,
fmt,
mem::{self, ManuallyDrop, MaybeUninit},
num::NonZeroUsize,
ptr::{self, NonNull},
};
pub struct SlabAllocator<T> {
free_list_head: Cell<Option<NonNull<Slot<T>>>>,
slab_list_head: Cell<Option<NonNull<Slab<T>>>>,
slab_list_next: Cell<Option<NonNull<Slab<T>>>>,
free_start: Cell<NonNull<Slot<T>>>,
free_end: Cell<NonNull<Slot<T>>>,
slab_capacity: NonZeroUsize,
}
unsafe impl<T> Send for SlabAllocator<T> {}
impl<T> SlabAllocator<T> {
#[inline]
#[must_use]
pub const fn new(slab_capacity: usize) -> Self {
if let Some(slab_capacity) = NonZeroUsize::new(slab_capacity) {
let dangling = NonNull::dangling();
SlabAllocator {
free_list_head: Cell::new(None),
slab_list_head: Cell::new(None),
slab_list_next: Cell::new(None),
free_start: Cell::new(dangling),
free_end: Cell::new(dangling),
slab_capacity,
}
} else {
panic!("`slab_capacity` must be non-zero");
}
}
#[inline(always)]
#[must_use]
pub fn allocate(&self) -> NonNull<T> {
let ptr = if let Some(ptr) = self.allocate_fast() {
ptr
} else if let Some(ptr) = self.allocate_fast2() {
ptr
} else {
self.allocate_slow()
.unwrap_or_else(|_| handle_alloc_error(self.slab_layout()))
};
ptr.cast::<T>()
}
#[inline(always)]
pub fn try_allocate(&self) -> Result<NonNull<T>, AllocError> {
let ptr = if let Some(ptr) = self.allocate_fast() {
ptr
} else if let Some(ptr) = self.allocate_fast2() {
ptr
} else {
self.allocate_slow()?
};
Ok(ptr.cast::<T>())
}
#[inline(always)]
fn allocate_fast(&self) -> Option<NonNull<Slot<T>>> {
let head = self.free_list_head.get()?;
let next = unsafe { (*head.as_ptr()).next_free };
self.free_list_head.set(next);
if cfg!(miri) {
let ptr = head.as_ptr().cast::<MaybeUninit<Slot<T>>>();
unsafe { ptr.write(MaybeUninit::uninit()) };
}
Some(head)
}
#[inline(always)]
fn allocate_fast2(&self) -> Option<NonNull<Slot<T>>> {
let ptr = self.free_start.get();
if ptr < self.free_end.get() {
let free_start = unsafe { NonNull::new_unchecked(ptr.as_ptr().add(1)) };
self.free_start.set(free_start);
Some(ptr)
} else {
None
}
}
#[cold]
fn allocate_slow(&self) -> Result<NonNull<Slot<T>>, AllocError> {
let slab = if let Some(slab) = self.slab_list_next.get() {
let next = unsafe { (*slab.as_ptr()).next };
self.slab_list_next.set(next);
slab
} else {
self.add_slab()?
};
let slots = unsafe { NonNull::new_unchecked(ptr::addr_of_mut!((*slab.as_ptr()).slots)) };
let free_start = unsafe { NonNull::new_unchecked(slots.as_ptr().add(1)) };
let free_end =
unsafe { NonNull::new_unchecked(slots.as_ptr().add(self.slab_capacity.get())) };
self.free_start.set(free_start);
self.free_end.set(free_end);
Ok(slots)
}
fn add_slab(&self) -> Result<NonNull<Slab<T>>, AllocError> {
let bytes = unsafe { alloc(self.slab_layout()) };
let slab = NonNull::new(bytes.cast::<Slab<T>>()).ok_or(AllocError)?;
unsafe { (*slab.as_ptr()).next = self.slab_list_head.get() };
self.slab_list_head.set(Some(slab));
Ok(slab)
}
fn slab_layout(&self) -> Layout {
Layout::new::<Option<NonNull<Slab<T>>>>()
.extend(Layout::array::<Slot<T>>(self.slab_capacity.get()).unwrap())
.unwrap()
.0
.pad_to_align()
}
#[inline(always)]
pub unsafe fn deallocate(&self, ptr: NonNull<T>) {
debug_assert!(
self.contains(ptr),
"attempted to deallocate a slot that does not belong to this allocator",
);
let ptr = ptr.cast::<Slot<T>>();
unsafe { (*ptr.as_ptr()).next_free = self.free_list_head.get() };
self.free_list_head.set(Some(ptr));
}
pub unsafe fn reset(&self) {
self.free_list_head.set(None);
self.slab_list_next.set(self.slab_list_head.get());
let dangling = NonNull::dangling();
self.free_start.set(dangling);
self.free_end.set(dangling);
}
pub fn contains(&self, ptr: NonNull<T>) -> bool {
let ptr = ptr.cast::<Slot<T>>();
#[allow(clippy::transmutes_expressible_as_ptr_casts)]
let addr = unsafe { mem::transmute::<*mut Slot<T>, usize>(ptr.as_ptr()) };
if addr & (mem::align_of::<Slot<T>>() - 1) != 0 {
return false;
}
let mut head = self.slab_list_head.get();
loop {
let slab = if let Some(slab) = head {
slab
} else {
return false;
};
if Some(slab) == self.slab_list_next.get() {
return false;
}
let slots_start = unsafe { ptr::addr_of_mut!((*slab.as_ptr()).slots) };
let slots_end = unsafe { slots_start.add(self.slab_capacity.get()) };
if (slots_start..slots_end).contains(&ptr.as_ptr()) {
break;
}
head = unsafe { (*slab.as_ptr()).next };
}
if (self.free_start.get()..self.free_end.get()).contains(&ptr) {
return false;
}
let mut head = self.free_list_head.get();
while let Some(slot) = head {
if slot == ptr {
return false;
}
head = unsafe { (*slot.as_ptr()).next_free };
}
true
}
fn slab_count(&self) -> usize {
let mut head = self.slab_list_head.get();
let mut count = 0;
while let Some(slab) = head {
head = unsafe { (*slab.as_ptr()).next };
count += 1;
}
count
}
}
#[allow(clippy::missing_fields_in_debug)]
impl<T> fmt::Debug for SlabAllocator<T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("SlabAllocator")
.field("slab_count", &self.slab_count())
.field("slab_capacity", &self.slab_capacity)
.finish()
}
}
impl<T> Drop for SlabAllocator<T> {
fn drop(&mut self) {
let slab_layout = self.slab_layout();
while let Some(slab) = self.slab_list_head.get() {
*self.slab_list_head.get_mut() = unsafe { (*slab.as_ptr()).next };
unsafe { dealloc(slab.as_ptr().cast(), slab_layout) };
}
}
}
#[repr(C)]
struct Slab<T> {
next: Option<NonNull<Self>>,
slots: Slot<T>,
}
#[repr(C)]
union Slot<T> {
next_free: Option<NonNull<Self>>,
value: ManuallyDrop<T>,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct AllocError;
impl fmt::Display for AllocError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str("memory allocation failed")
}
}
#[cfg(test)]
mod tests {
use super::*;
use core::mem;
#[test]
fn basic_usage1() {
let allocator = SlabAllocator::<i32>::new(2);
let mut x = allocator.allocate();
unsafe { x.as_ptr().write(69) };
let mut y = allocator.allocate();
unsafe { y.as_ptr().write(42) };
assert_eq!(allocator.slab_count(), 1);
mem::swap(unsafe { x.as_mut() }, unsafe { y.as_mut() });
unsafe { allocator.deallocate(x) };
let mut x2 = allocator.allocate();
unsafe { x2.as_ptr().write(12) };
assert_eq!(allocator.slab_count(), 1);
mem::swap(unsafe { y.as_mut() }, unsafe { x2.as_mut() });
unsafe { allocator.deallocate(y) };
unsafe { allocator.deallocate(x2) };
}
#[test]
fn basic_usage2() {
let allocator = SlabAllocator::<i32>::new(1);
let mut x = allocator.allocate();
unsafe { x.as_ptr().write(1) };
let mut y = allocator.allocate();
unsafe { y.as_ptr().write(2) };
let mut z = allocator.allocate();
unsafe { z.as_ptr().write(3) };
assert_eq!(allocator.slab_count(), 3);
mem::swap(unsafe { x.as_mut() }, unsafe { y.as_mut() });
mem::swap(unsafe { y.as_mut() }, unsafe { z.as_mut() });
mem::swap(unsafe { z.as_mut() }, unsafe { x.as_mut() });
unsafe { allocator.deallocate(y) };
let mut y2 = allocator.allocate();
unsafe { y2.as_ptr().write(20) };
assert_eq!(allocator.slab_count(), 3);
mem::swap(unsafe { x.as_mut() }, unsafe { y2.as_mut() });
unsafe { allocator.deallocate(x) };
unsafe { allocator.deallocate(z) };
let mut x2 = allocator.allocate();
unsafe { x2.as_ptr().write(10) };
mem::swap(unsafe { y2.as_mut() }, unsafe { x2.as_mut() });
let mut z2 = allocator.allocate();
unsafe { z2.as_ptr().write(30) };
assert_eq!(allocator.slab_count(), 3);
mem::swap(unsafe { x2.as_mut() }, unsafe { z2.as_mut() });
unsafe { allocator.deallocate(x2) };
mem::swap(unsafe { z2.as_mut() }, unsafe { y2.as_mut() });
unsafe { allocator.deallocate(y2) };
unsafe { allocator.deallocate(z2) };
}
#[test]
fn basic_usage3() {
let allocator = SlabAllocator::<i32>::new(2);
let mut x = allocator.allocate();
unsafe { x.as_ptr().write(1) };
let mut y = allocator.allocate();
unsafe { y.as_ptr().write(2) };
assert_eq!(allocator.slab_count(), 1);
mem::swap(unsafe { x.as_mut() }, unsafe { y.as_mut() });
let z = allocator.allocate();
unsafe { z.as_ptr().write(3) };
assert_eq!(allocator.slab_count(), 2);
unsafe { allocator.deallocate(x) };
unsafe { allocator.deallocate(z) };
let mut z2 = allocator.allocate();
unsafe { z2.as_ptr().write(30) };
let mut x2 = allocator.allocate();
unsafe { x2.as_ptr().write(10) };
assert_eq!(allocator.slab_count(), 2);
mem::swap(unsafe { x2.as_mut() }, unsafe { z2.as_mut() });
unsafe { allocator.deallocate(x2) };
unsafe { allocator.deallocate(y) };
unsafe { allocator.deallocate(z2) };
}
#[test]
fn reusing_slots1() {
let allocator = SlabAllocator::<i32>::new(2);
let x = allocator.allocate();
let y = allocator.allocate();
unsafe { allocator.deallocate(y) };
let y2 = allocator.allocate();
assert_eq!(y2, y);
unsafe { allocator.deallocate(x) };
let x2 = allocator.allocate();
assert_eq!(x2, x);
unsafe { allocator.deallocate(y2) };
unsafe { allocator.deallocate(x2) };
}
#[test]
fn reusing_slots2() {
let allocator = SlabAllocator::<i32>::new(1);
let x = allocator.allocate();
unsafe { allocator.deallocate(x) };
let x2 = allocator.allocate();
assert_eq!(x, x2);
let y = allocator.allocate();
let z = allocator.allocate();
unsafe { allocator.deallocate(y) };
unsafe { allocator.deallocate(x2) };
let x3 = allocator.allocate();
let y2 = allocator.allocate();
assert_eq!(x3, x2);
assert_eq!(y2, y);
unsafe { allocator.deallocate(x3) };
unsafe { allocator.deallocate(y2) };
unsafe { allocator.deallocate(z) };
}
#[test]
fn reusing_slots3() {
let allocator = SlabAllocator::<i32>::new(2);
let x = allocator.allocate();
let y = allocator.allocate();
unsafe { allocator.deallocate(x) };
unsafe { allocator.deallocate(y) };
let y2 = allocator.allocate();
let x2 = allocator.allocate();
let z = allocator.allocate();
assert_eq!(x2, x);
assert_eq!(y2, y);
unsafe { allocator.deallocate(x2) };
unsafe { allocator.deallocate(z) };
unsafe { allocator.deallocate(y2) };
let y3 = allocator.allocate();
let z2 = allocator.allocate();
let x3 = allocator.allocate();
assert_eq!(y3, y2);
assert_eq!(z2, z);
assert_eq!(x3, x2);
unsafe { allocator.deallocate(x3) };
unsafe { allocator.deallocate(y3) };
unsafe { allocator.deallocate(z2) };
}
#[test]
fn reusing_slots4() {
let allocator = SlabAllocator::<i32>::new(2);
let x = allocator.allocate();
let y = allocator.allocate();
unsafe { allocator.deallocate(y) };
let y2 = allocator.allocate();
assert_eq!(y2, y);
unsafe { allocator.reset() };
let x2 = allocator.allocate();
let y3 = allocator.allocate();
assert_eq!(x2, x);
assert_eq!(y3, y2);
unsafe { allocator.deallocate(y3) };
unsafe { allocator.deallocate(x2) };
}
#[test]
fn reusing_slots5() {
let allocator = SlabAllocator::<i32>::new(1);
let x = allocator.allocate();
unsafe { allocator.deallocate(x) };
let x2 = allocator.allocate();
assert_eq!(x, x2);
let y = allocator.allocate();
let z = allocator.allocate();
unsafe { allocator.reset() };
let z2 = allocator.allocate();
let y2 = allocator.allocate();
assert_eq!(z2, z);
assert_eq!(y2, y);
unsafe { allocator.deallocate(z2) };
unsafe { allocator.deallocate(y2) };
}
#[test]
fn reusing_slots6() {
let allocator = SlabAllocator::<i32>::new(2);
let x = allocator.allocate();
let y = allocator.allocate();
unsafe { allocator.reset() };
let x2 = allocator.allocate();
let y2 = allocator.allocate();
let z = allocator.allocate();
assert_eq!(x2, x);
assert_eq!(y2, y);
unsafe { allocator.deallocate(x2) };
unsafe { allocator.deallocate(z) };
unsafe { allocator.deallocate(y2) };
unsafe { allocator.reset() };
let z2 = allocator.allocate();
let _ = allocator.allocate();
let x3 = allocator.allocate();
let y3 = allocator.allocate();
assert_eq!(z2, z);
assert_eq!(x3, x2);
assert_eq!(y3, y2);
unsafe { allocator.reset() };
}
#[test]
fn same_slab() {
const MAX_DIFF: usize = 2 * mem::size_of::<Slot<i32>>();
let allocator = SlabAllocator::<i32>::new(3);
let x = allocator.allocate();
let y = allocator.allocate();
let z = allocator.allocate();
assert!((x.as_ptr() as usize).abs_diff(y.as_ptr() as usize) <= MAX_DIFF);
assert!((y.as_ptr() as usize).abs_diff(z.as_ptr() as usize) <= MAX_DIFF);
assert!((z.as_ptr() as usize).abs_diff(x.as_ptr() as usize) <= MAX_DIFF);
}
#[test]
fn different_slabs() {
const MIN_DIFF: usize = mem::size_of::<Slab<i32>>();
let allocator = SlabAllocator::<i32>::new(1);
let x = allocator.allocate();
let y = allocator.allocate();
let z = allocator.allocate();
assert!((x.as_ptr() as usize).abs_diff(y.as_ptr() as usize) >= MIN_DIFF);
assert!((y.as_ptr() as usize).abs_diff(z.as_ptr() as usize) >= MIN_DIFF);
assert!((z.as_ptr() as usize).abs_diff(x.as_ptr() as usize) >= MIN_DIFF);
}
#[test]
#[should_panic]
fn zero_slab_capacity() {
let _ = SlabAllocator::<()>::new(0);
}
#[cfg(debug_assertions)]
#[test]
#[should_panic]
fn unaligned_ptr() {
let allocator = SlabAllocator::<i32>::new(1);
let x = allocator.allocate();
unsafe { allocator.deallocate(x.byte_add(1)) };
}
#[cfg(debug_assertions)]
#[test]
#[should_panic]
fn foreign_ptr() {
let allocator = SlabAllocator::<i32>::new(1);
let mut x = 69;
unsafe { allocator.deallocate((&mut x).into()) };
}
#[cfg(debug_assertions)]
#[test]
#[should_panic]
fn not_yet_allocated() {
let allocator = SlabAllocator::<i32>::new(2);
let x = allocator.allocate();
unsafe { allocator.deallocate(x.byte_add(size_of::<Slot<i32>>())) };
}
#[cfg(debug_assertions)]
#[test]
#[should_panic]
fn already_deallocated() {
let allocator = SlabAllocator::<i32>::new(2);
let x = allocator.allocate();
unsafe { allocator.deallocate(x) };
unsafe { allocator.deallocate(x) };
}
}