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 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 try_claim(&self) -> Option<super::ByteClaim<'_>> {
let claim = self.inner.claim()?;
let ptr = claim.into_ptr().cast::<u8>();
let slab_ptr = core::ptr::from_ref(&self.inner).cast::<u8>();
Some(unsafe { super::ByteClaim::from_raw_parts(ptr, slab_ptr, free_raw_impl::<N>, 0, N) })
}
#[inline]
pub fn claim(&self) -> super::ByteClaim<'_> {
self.try_claim().expect("byte slab full")
}
#[inline]
pub unsafe fn free_raw(&self, ptr: *mut u8) {
unsafe {
self.inner.free_ptr(ptr.cast());
}
}
#[inline]
pub unsafe fn alloc_raw(&self, src: *const u8, size: usize) -> *mut u8 {
assert!(size <= N, "raw alloc size {size} exceeds slot size {N}");
let slot_ptr = self
.inner
.claim_ptr()
.unwrap_or_else(|| panic!("byte slab full"));
let dst = slot_ptr.cast::<u8>();
unsafe { core::ptr::copy_nonoverlapping(src, dst, size) };
dst
}
#[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()
}
}
unsafe fn free_raw_impl<const N: usize>(slab_ptr: *const u8, slot_ptr: *mut u8, _chunk_idx: usize) {
let slab = unsafe { &*(slab_ptr as *const crate::bounded::Slab<super::AlignedBytes<N>>) };
unsafe {
slab.free_ptr(slot_ptr.cast());
}
}
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());
}
#[test]
fn claim_write_typed() {
let slab: Slab<64> = unsafe { Slab::with_capacity(4) };
let claim = slab.claim();
let slot = claim.write(42u64);
assert_eq!(*slot, 42);
slab.free(slot);
}
#[test]
fn claim_write_raw() {
let slab: Slab<64> = unsafe { Slab::with_capacity(4) };
let claim = slab.claim();
let val: u64 = 99;
let ptr = unsafe {
claim.write_raw(&val as *const u64 as *const u8, core::mem::size_of::<u64>())
};
assert_eq!(unsafe { *(ptr as *const u64) }, 99);
let slot = unsafe { super::Slot::<u64>::from_raw(ptr) };
slab.free(slot);
}
#[test]
fn claim_drop_returns_to_freelist() {
let slab: Slab<64> = unsafe { Slab::with_capacity(1) };
let claim = slab.claim();
drop(claim);
let claim = slab.claim();
let slot = claim.write(7u64);
assert_eq!(*slot, 7);
slab.free(slot);
}
#[test]
fn try_claim_returns_none_when_full() {
let slab: Slab<64> = unsafe { Slab::with_capacity(1) };
let _held = slab.claim();
assert!(slab.try_claim().is_none());
}
#[test]
fn try_claim_succeeds_after_abandon() {
let slab: Slab<64> = unsafe { Slab::with_capacity(1) };
let claim = slab.claim();
drop(claim);
let claim2 = slab.try_claim();
assert!(claim2.is_some());
let slot = claim2.unwrap().write(42u64);
slab.free(slot);
}
}