use std::alloc::{alloc, dealloc, Layout};
use std::marker::PhantomData;
use std::mem::MaybeUninit;
use std::ptr::NonNull;
use std::sync::atomic::{AtomicU32, AtomicU64, AtomicUsize, Ordering};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[repr(u8)]
pub enum SizeClass {
Size64 = 0,
Size128 = 1,
Size256 = 2,
Size512 = 3,
Size1024 = 4,
}
impl SizeClass {
#[inline]
pub const fn size_bytes(&self) -> usize {
match self {
Self::Size64 => 64,
Self::Size128 => 128,
Self::Size256 => 256,
Self::Size512 => 512,
Self::Size1024 => 1024,
}
}
pub fn for_size(size: usize) -> Option<Self> {
match size {
0..=64 => Some(Self::Size64),
65..=128 => Some(Self::Size128),
129..=256 => Some(Self::Size256),
257..=512 => Some(Self::Size512),
513..=1024 => Some(Self::Size1024),
_ => None,
}
}
pub const fn all() -> [Self; 5] {
[Self::Size64, Self::Size128, Self::Size256, Self::Size512, Self::Size1024]
}
}
#[derive(Clone, Copy, Debug)]
pub struct GenerationalHandle {
packed: u64,
}
impl GenerationalHandle {
#[inline]
pub const fn new(slot: u32, generation: u32) -> Self {
Self {
packed: ((generation as u64) << 32) | (slot as u64),
}
}
#[inline]
pub const fn slot(&self) -> u32 {
self.packed as u32
}
#[inline]
pub const fn generation(&self) -> u32 {
(self.packed >> 32) as u32
}
#[inline]
pub const fn from_raw(packed: u64) -> Self {
Self { packed }
}
#[inline]
pub const fn to_raw(&self) -> u64 {
self.packed
}
#[inline]
pub const fn invalid() -> Self {
Self { packed: u64::MAX }
}
#[inline]
pub const fn is_valid(&self) -> bool {
self.packed != u64::MAX
}
}
#[repr(C)]
struct SlotHeader {
generation: AtomicU32,
next_free: AtomicU32,
flags: AtomicU32,
_reserved: u32,
}
impl SlotHeader {
const FLAG_ALLOCATED: u32 = 1 << 0;
fn new() -> Self {
Self {
generation: AtomicU32::new(0),
next_free: AtomicU32::new(u32::MAX),
flags: AtomicU32::new(0),
_reserved: 0,
}
}
#[inline]
fn is_allocated(&self) -> bool {
self.flags.load(Ordering::Acquire) & Self::FLAG_ALLOCATED != 0
}
#[inline]
fn set_allocated(&self, allocated: bool) {
if allocated {
self.flags.fetch_or(Self::FLAG_ALLOCATED, Ordering::Release);
} else {
self.flags.fetch_and(!Self::FLAG_ALLOCATED, Ordering::Release);
}
}
#[inline]
fn increment_generation(&self) -> u32 {
self.generation.fetch_add(1, Ordering::Release) + 1
}
}
struct Slab {
data: NonNull<u8>,
slot_size: usize,
slot_count: usize,
layout: Layout,
free_head: AtomicU32,
allocated_count: AtomicUsize,
}
impl Slab {
fn new(size_class: SizeClass, slot_count: usize) -> Option<Self> {
let user_size = size_class.size_bytes();
let header_size = std::mem::size_of::<SlotHeader>();
let slot_size = header_size + user_size;
let slot_size = (slot_size + 15) & !15;
let total_size = slot_size * slot_count;
let layout = Layout::from_size_align(total_size, 64).ok()?;
let ptr = unsafe { alloc(layout) };
let data = NonNull::new(ptr)?;
unsafe {
for i in 0..slot_count {
let slot_ptr = data.as_ptr().add(i * slot_size);
let header = &mut *(slot_ptr as *mut SlotHeader);
*header = SlotHeader::new();
if i < slot_count - 1 {
header.next_free.store((i + 1) as u32, Ordering::Relaxed);
} else {
header.next_free.store(u32::MAX, Ordering::Relaxed);
}
}
}
Some(Self {
data,
slot_size,
slot_count,
layout,
free_head: AtomicU32::new(0),
allocated_count: AtomicUsize::new(0),
})
}
fn allocate(&self) -> Option<GenerationalHandle> {
loop {
let head = self.free_head.load(Ordering::Acquire);
if head == u32::MAX {
return None; }
let header = self.get_header(head as usize)?;
let next = header.next_free.load(Ordering::Acquire);
match self.free_head.compare_exchange_weak(
head,
next,
Ordering::AcqRel,
Ordering::Acquire,
) {
Ok(_) => {
let generation = header.increment_generation();
header.set_allocated(true);
self.allocated_count.fetch_add(1, Ordering::Relaxed);
return Some(GenerationalHandle::new(head, generation));
}
Err(_) => continue, }
}
}
fn free(&self, handle: GenerationalHandle) -> bool {
let slot = handle.slot() as usize;
if slot >= self.slot_count {
return false;
}
let header = match self.get_header(slot) {
Some(h) => h,
None => return false,
};
if header.generation.load(Ordering::Acquire) != handle.generation() {
return false; }
if !header.is_allocated() {
return false; }
header.set_allocated(false);
loop {
let head = self.free_head.load(Ordering::Acquire);
header.next_free.store(head, Ordering::Release);
match self.free_head.compare_exchange_weak(
head,
slot as u32,
Ordering::AcqRel,
Ordering::Acquire,
) {
Ok(_) => {
self.allocated_count.fetch_sub(1, Ordering::Relaxed);
return true;
}
Err(_) => continue,
}
}
}
fn get_ptr(&self, handle: GenerationalHandle) -> Option<NonNull<u8>> {
let slot = handle.slot() as usize;
if slot >= self.slot_count {
return None;
}
let header = self.get_header(slot)?;
if header.generation.load(Ordering::Acquire) != handle.generation() {
return None;
}
if !header.is_allocated() {
return None;
}
let header_size = std::mem::size_of::<SlotHeader>();
let slot_ptr = unsafe { self.data.as_ptr().add(slot * self.slot_size) };
let user_ptr = unsafe { slot_ptr.add(header_size) };
NonNull::new(user_ptr)
}
fn get_header(&self, slot: usize) -> Option<&SlotHeader> {
if slot >= self.slot_count {
return None;
}
let slot_ptr = unsafe { self.data.as_ptr().add(slot * self.slot_size) };
Some(unsafe { &*(slot_ptr as *const SlotHeader) })
}
fn stats(&self) -> SlabStats {
SlabStats {
slot_count: self.slot_count,
allocated_count: self.allocated_count.load(Ordering::Relaxed),
slot_size: self.slot_size,
}
}
}
impl Drop for Slab {
fn drop(&mut self) {
unsafe {
dealloc(self.data.as_ptr(), self.layout);
}
}
}
unsafe impl Send for Slab {}
unsafe impl Sync for Slab {}
#[derive(Debug, Clone)]
pub struct SlabStats {
pub slot_count: usize,
pub allocated_count: usize,
pub slot_size: usize,
}
#[derive(Clone)]
pub struct SlabAllocatorConfig {
pub initial_slots: [usize; 5],
pub max_slabs: usize,
}
impl Default for SlabAllocatorConfig {
fn default() -> Self {
Self {
initial_slots: [1024, 512, 256, 128, 64], max_slabs: 64,
}
}
}
pub struct SlabAllocator {
slabs: [parking_lot::RwLock<Vec<Slab>>; 5],
config: SlabAllocatorConfig,
total_allocations: AtomicU64,
total_frees: AtomicU64,
}
impl SlabAllocator {
pub fn new() -> Self {
Self::with_config(SlabAllocatorConfig::default())
}
pub fn with_config(config: SlabAllocatorConfig) -> Self {
let slabs = std::array::from_fn(|i| {
let size_class = SizeClass::all()[i];
let initial_slots = config.initial_slots[i];
let slab = Slab::new(size_class, initial_slots)
.expect("Failed to create initial slab");
parking_lot::RwLock::new(vec![slab])
});
Self {
slabs,
config,
total_allocations: AtomicU64::new(0),
total_frees: AtomicU64::new(0),
}
}
pub fn allocate(&self, size: usize) -> Option<(GenerationalHandle, NonNull<u8>)> {
let size_class = SizeClass::for_size(size)?;
self.allocate_from_class(size_class)
}
pub fn allocate_from_class(&self, size_class: SizeClass) -> Option<(GenerationalHandle, NonNull<u8>)> {
let class_idx = size_class as usize;
{
let slabs = self.slabs[class_idx].read();
for (slab_idx, slab) in slabs.iter().enumerate() {
if let Some(handle) = slab.allocate() {
let ptr = slab.get_ptr(handle)?;
let full_handle = GenerationalHandle::new(
((slab_idx as u32) << 24) | handle.slot(),
handle.generation(),
);
self.total_allocations.fetch_add(1, Ordering::Relaxed);
return Some((full_handle, ptr));
}
}
}
let mut slabs = self.slabs[class_idx].write();
for (slab_idx, slab) in slabs.iter().enumerate() {
if let Some(handle) = slab.allocate() {
let ptr = slab.get_ptr(handle)?;
let full_handle = GenerationalHandle::new(
((slab_idx as u32) << 24) | handle.slot(),
handle.generation(),
);
self.total_allocations.fetch_add(1, Ordering::Relaxed);
return Some((full_handle, ptr));
}
}
if slabs.len() >= self.config.max_slabs {
return None; }
let new_slab = Slab::new(size_class, self.config.initial_slots[class_idx])?;
let handle = new_slab.allocate()?;
let ptr = new_slab.get_ptr(handle)?;
let slab_idx = slabs.len();
slabs.push(new_slab);
let full_handle = GenerationalHandle::new(
((slab_idx as u32) << 24) | handle.slot(),
handle.generation(),
);
self.total_allocations.fetch_add(1, Ordering::Relaxed);
Some((full_handle, ptr))
}
pub fn free(&self, size_class: SizeClass, handle: GenerationalHandle) -> bool {
let class_idx = size_class as usize;
let slab_idx = (handle.slot() >> 24) as usize;
let slot = handle.slot() & 0x00FFFFFF;
let local_handle = GenerationalHandle::new(slot, handle.generation());
let slabs = self.slabs[class_idx].read();
if slab_idx >= slabs.len() {
return false;
}
let result = slabs[slab_idx].free(local_handle);
if result {
self.total_frees.fetch_add(1, Ordering::Relaxed);
}
result
}
pub fn get_ptr(&self, size_class: SizeClass, handle: GenerationalHandle) -> Option<NonNull<u8>> {
let class_idx = size_class as usize;
let slab_idx = (handle.slot() >> 24) as usize;
let slot = handle.slot() & 0x00FFFFFF;
let local_handle = GenerationalHandle::new(slot, handle.generation());
let slabs = self.slabs[class_idx].read();
if slab_idx >= slabs.len() {
return None;
}
slabs[slab_idx].get_ptr(local_handle)
}
pub fn stats(&self) -> AllocatorStats {
let mut class_stats = Vec::new();
for (i, size_class) in SizeClass::all().iter().enumerate() {
let slabs = self.slabs[i].read();
let slab_stats: Vec<_> = slabs.iter().map(|s| s.stats()).collect();
class_stats.push(SizeClassStats {
size_class: *size_class,
slab_count: slab_stats.len(),
total_slots: slab_stats.iter().map(|s| s.slot_count).sum(),
allocated_slots: slab_stats.iter().map(|s| s.allocated_count).sum(),
});
}
AllocatorStats {
total_allocations: self.total_allocations.load(Ordering::Relaxed),
total_frees: self.total_frees.load(Ordering::Relaxed),
class_stats,
}
}
}
impl Default for SlabAllocator {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug, Clone)]
pub struct SizeClassStats {
pub size_class: SizeClass,
pub slab_count: usize,
pub total_slots: usize,
pub allocated_slots: usize,
}
#[derive(Debug, Clone)]
pub struct AllocatorStats {
pub total_allocations: u64,
pub total_frees: u64,
pub class_stats: Vec<SizeClassStats>,
}
pub struct TypedSlabAllocator<T> {
allocator: SlabAllocator,
size_class: SizeClass,
_marker: PhantomData<T>,
}
impl<T: Sized> TypedSlabAllocator<T> {
pub fn new() -> Option<Self> {
let size = std::mem::size_of::<T>();
let size_class = SizeClass::for_size(size)?;
Some(Self {
allocator: SlabAllocator::new(),
size_class,
_marker: PhantomData,
})
}
pub fn allocate(&self, value: T) -> Option<(GenerationalHandle, NonNull<T>)> {
let (handle, ptr) = self.allocator.allocate_from_class(self.size_class)?;
unsafe {
std::ptr::write(ptr.as_ptr() as *mut T, value);
}
Some((handle, ptr.cast()))
}
pub fn allocate_uninit(&self) -> Option<(GenerationalHandle, NonNull<MaybeUninit<T>>)> {
let (handle, ptr) = self.allocator.allocate_from_class(self.size_class)?;
Some((handle, ptr.cast()))
}
pub fn free(&self, handle: GenerationalHandle) -> bool {
if let Some(ptr) = self.allocator.get_ptr(self.size_class, handle) {
unsafe {
std::ptr::drop_in_place(ptr.as_ptr() as *mut T);
}
}
self.allocator.free(self.size_class, handle)
}
pub fn get(&self, handle: GenerationalHandle) -> Option<&T> {
let ptr = self.allocator.get_ptr(self.size_class, handle)?;
Some(unsafe { &*(ptr.as_ptr() as *const T) })
}
pub fn get_mut(&self, handle: GenerationalHandle) -> Option<&mut T> {
let ptr = self.allocator.get_ptr(self.size_class, handle)?;
Some(unsafe { &mut *(ptr.as_ptr() as *mut T) })
}
}
impl<T: Sized> Default for TypedSlabAllocator<T> {
fn default() -> Self {
Self::new().expect("Type too large for slab allocation")
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::Arc;
use std::thread;
#[test]
fn test_generational_handle() {
let handle = GenerationalHandle::new(42, 7);
assert_eq!(handle.slot(), 42);
assert_eq!(handle.generation(), 7);
let raw = handle.to_raw();
let handle2 = GenerationalHandle::from_raw(raw);
assert_eq!(handle2.slot(), 42);
assert_eq!(handle2.generation(), 7);
}
#[test]
fn test_size_class() {
assert_eq!(SizeClass::for_size(32), Some(SizeClass::Size64));
assert_eq!(SizeClass::for_size(64), Some(SizeClass::Size64));
assert_eq!(SizeClass::for_size(65), Some(SizeClass::Size128));
assert_eq!(SizeClass::for_size(512), Some(SizeClass::Size512));
assert_eq!(SizeClass::for_size(1025), None);
}
#[test]
fn test_slab_basic() {
let slab = Slab::new(SizeClass::Size64, 10).unwrap();
let h1 = slab.allocate().unwrap();
let h2 = slab.allocate().unwrap();
let h3 = slab.allocate().unwrap();
assert_ne!(h1.slot(), h2.slot());
assert_ne!(h2.slot(), h3.slot());
assert!(slab.free(h2));
let h4 = slab.allocate().unwrap();
assert_eq!(h4.slot(), h2.slot()); assert_ne!(h4.generation(), h2.generation()); }
#[test]
fn test_slab_generation() {
let slab = Slab::new(SizeClass::Size64, 10).unwrap();
let h1 = slab.allocate().unwrap();
assert!(slab.free(h1));
let ptr = slab.get_ptr(h1);
assert!(ptr.is_none());
assert!(!slab.free(h1));
}
#[test]
fn test_allocator_basic() {
let allocator = SlabAllocator::new();
let (h1, p1) = allocator.allocate(32).unwrap();
let (h2, p2) = allocator.allocate(100).unwrap();
let (h3, p3) = allocator.allocate(300).unwrap();
assert!(p1.as_ptr() != p2.as_ptr());
assert!(p2.as_ptr() != p3.as_ptr());
assert!(allocator.free(SizeClass::Size64, h1));
assert!(allocator.free(SizeClass::Size128, h2));
assert!(allocator.free(SizeClass::Size512, h3));
}
#[test]
fn test_allocator_concurrent() {
let allocator = Arc::new(SlabAllocator::new());
let mut handles = vec![];
for _ in 0..4 {
let allocator = allocator.clone();
handles.push(thread::spawn(move || {
let mut local_handles = Vec::new();
for _ in 0..1000 {
let (handle, _ptr) = allocator.allocate(64).unwrap();
local_handles.push(handle);
}
for handle in local_handles.drain(..500) {
allocator.free(SizeClass::Size64, handle);
}
local_handles
}));
}
let mut all_handles = Vec::new();
for handle in handles {
all_handles.extend(handle.join().unwrap());
}
let stats = allocator.stats();
assert_eq!(stats.total_allocations, 4000);
assert_eq!(stats.total_frees, 2000);
}
#[test]
fn test_typed_allocator() {
#[derive(Debug, PartialEq)]
struct TestNode {
value: i32,
next: Option<u64>,
}
let allocator = TypedSlabAllocator::<TestNode>::new().unwrap();
let (h1, _) = allocator.allocate(TestNode { value: 42, next: None }).unwrap();
let (h2, _) = allocator.allocate(TestNode { value: 100, next: Some(h1.to_raw()) }).unwrap();
let node1 = allocator.get(h1).unwrap();
assert_eq!(node1.value, 42);
let node2 = allocator.get(h2).unwrap();
assert_eq!(node2.value, 100);
assert_eq!(node2.next, Some(h1.to_raw()));
assert!(allocator.free(h1));
assert!(allocator.free(h2));
}
#[test]
fn test_stats() {
let allocator = SlabAllocator::new();
for _ in 0..50 {
allocator.allocate(32).unwrap();
}
for _ in 0..30 {
allocator.allocate(200).unwrap();
}
let stats = allocator.stats();
assert_eq!(stats.total_allocations, 80);
let size64_stats = stats.class_stats.iter()
.find(|s| s.size_class == SizeClass::Size64)
.unwrap();
assert_eq!(size64_stats.allocated_slots, 50);
}
}