use core::marker::PhantomData;
use core::mem;
use crate::shared::{Full, SlotCell};
use super::{AlignedBytes, Slot, validate_type};
pub struct Slab<const N: usize> {
inner: crate::bounded::Slab<AlignedBytes<N>>,
}
impl<const N: usize> Slab<N> {
#[inline]
pub const unsafe fn new() -> Self {
Self {
inner: unsafe { crate::bounded::Slab::new() },
}
}
#[inline]
pub unsafe fn init(&self, capacity: usize) {
unsafe { self.inner.init(capacity) };
}
#[inline]
pub unsafe fn with_capacity(capacity: usize) -> Self {
Self {
inner: unsafe { crate::bounded::Slab::with_capacity(capacity) },
}
}
#[inline]
pub fn alloc<T>(&self, value: T) -> Slot<T> {
self.try_alloc(value)
.unwrap_or_else(|_| panic!("byte slab full"))
}
pub fn try_alloc<T>(&self, value: T) -> Result<Slot<T>, Full<T>> {
validate_type::<T, N>();
let Some(slot_ptr) = self.inner.claim_ptr() else {
return Err(Full(value));
};
unsafe {
let data_ptr = slot_ptr.cast::<T>();
core::ptr::write(data_ptr, value);
}
Ok(Slot {
ptr: slot_ptr.cast::<u8>(),
_marker: PhantomData,
})
}
#[inline]
pub fn free<T>(&self, ptr: Slot<T>) {
let data_ptr = ptr.ptr;
debug_assert!(
self.inner.contains_ptr(data_ptr as *const ()),
"slot was not allocated from this slab"
);
mem::forget(ptr);
unsafe {
core::ptr::drop_in_place(data_ptr.cast::<T>());
self.inner
.free_ptr(data_ptr.cast::<SlotCell<AlignedBytes<N>>>());
}
}
#[inline]
pub fn take<T>(&self, ptr: Slot<T>) -> T {
let data_ptr = ptr.ptr;
debug_assert!(
self.inner.contains_ptr(data_ptr as *const ()),
"slot was not allocated from this slab"
);
mem::forget(ptr);
unsafe {
let value = core::ptr::read(data_ptr.cast::<T>());
self.inner
.free_ptr(data_ptr.cast::<SlotCell<AlignedBytes<N>>>());
value
}
}
#[inline]
pub fn capacity(&self) -> usize {
self.inner.capacity()
}
}
impl<const N: usize> core::fmt::Debug for Slab<N> {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
f.debug_struct("byte::bounded::Slab")
.field("slot_size", &N)
.field("capacity", &self.capacity())
.finish()
}
}
#[cfg(test)]
#[allow(clippy::float_cmp)]
mod tests {
use super::*;
#[test]
fn basic_alloc_free() {
let slab: Slab<128> = unsafe { Slab::with_capacity(10) };
let ptr = slab.alloc(42u64);
assert_eq!(*ptr, 42);
slab.free(ptr);
}
#[test]
fn heterogeneous_types() {
let slab: Slab<128> = unsafe { Slab::with_capacity(10) };
let p1 = slab.alloc(42u64);
let p2 = slab.alloc([1.0f64; 4]);
let p3 = slab.alloc(String::from("hello"));
assert_eq!(*p1, 42);
assert_eq!(p2[0], 1.0);
assert_eq!(&*p3, "hello");
slab.free(p3);
slab.free(p2);
slab.free(p1);
}
#[test]
fn take_returns_value() {
let slab: Slab<64> = unsafe { Slab::with_capacity(10) };
let ptr = slab.alloc(String::from("owned"));
let val = slab.take(ptr);
assert_eq!(val, "owned");
}
#[test]
fn full_returns_error() {
let slab: Slab<64> = unsafe { Slab::with_capacity(1) };
let p1 = slab.alloc(1u64);
let result = slab.try_alloc(2u64);
assert!(result.is_err());
assert_eq!(result.unwrap_err().0, 2);
slab.free(p1);
}
#[test]
#[should_panic(expected = "exceeds byte slab slot size")]
fn rejects_oversized_type() {
let slab: Slab<8> = unsafe { Slab::with_capacity(1) };
let _p = slab.alloc([0u64; 2]);
}
#[test]
fn deref_mut() {
let slab: Slab<64> = unsafe { Slab::with_capacity(10) };
let mut ptr = slab.alloc(String::from("hello"));
ptr.push_str(" world");
assert_eq!(&*ptr, "hello world");
slab.free(ptr);
}
#[test]
fn reuse_after_free() {
let slab: Slab<64> = unsafe { Slab::with_capacity(1) };
let ptr = slab.alloc(1u64);
slab.free(ptr);
let ptr = slab.alloc(2u64);
assert_eq!(*ptr, 2);
slab.free(ptr);
}
#[cfg(debug_assertions)]
#[test]
fn debug_drop_panics() {
let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
let slab: Slab<64> = unsafe { Slab::with_capacity(10) };
let _ptr = slab.alloc(42u64);
}));
assert!(result.is_err());
}
}