1use std::alloc::{Layout, alloc, dealloc};
43use std::marker::PhantomData;
44use std::mem::MaybeUninit;
45use std::ptr::NonNull;
46use std::sync::atomic::{AtomicU32, AtomicU64, AtomicUsize, Ordering};
47
48#[derive(Debug, Clone, Copy, PartialEq, Eq)]
50#[repr(u8)]
51pub enum SizeClass {
52 Size64 = 0,
54 Size128 = 1,
56 Size256 = 2,
58 Size512 = 3,
60 Size1024 = 4,
62}
63
64impl SizeClass {
65 #[inline]
67 pub const fn size_bytes(&self) -> usize {
68 match self {
69 Self::Size64 => 64,
70 Self::Size128 => 128,
71 Self::Size256 => 256,
72 Self::Size512 => 512,
73 Self::Size1024 => 1024,
74 }
75 }
76
77 pub fn for_size(size: usize) -> Option<Self> {
79 match size {
80 0..=64 => Some(Self::Size64),
81 65..=128 => Some(Self::Size128),
82 129..=256 => Some(Self::Size256),
83 257..=512 => Some(Self::Size512),
84 513..=1024 => Some(Self::Size1024),
85 _ => None,
86 }
87 }
88
89 pub const fn all() -> [Self; 5] {
91 [
92 Self::Size64,
93 Self::Size128,
94 Self::Size256,
95 Self::Size512,
96 Self::Size1024,
97 ]
98 }
99}
100
101#[derive(Clone, Copy, Debug)]
109pub struct GenerationalHandle {
110 packed: u64,
112}
113
114impl GenerationalHandle {
115 #[inline]
117 pub const fn new(slot: u32, generation: u32) -> Self {
118 Self {
119 packed: ((generation as u64) << 32) | (slot as u64),
120 }
121 }
122
123 #[inline]
125 pub const fn slot(&self) -> u32 {
126 self.packed as u32
127 }
128
129 #[inline]
131 pub const fn generation(&self) -> u32 {
132 (self.packed >> 32) as u32
133 }
134
135 #[inline]
137 pub const fn from_raw(packed: u64) -> Self {
138 Self { packed }
139 }
140
141 #[inline]
143 pub const fn to_raw(&self) -> u64 {
144 self.packed
145 }
146
147 #[inline]
149 pub const fn invalid() -> Self {
150 Self { packed: u64::MAX }
151 }
152
153 #[inline]
155 pub const fn is_valid(&self) -> bool {
156 self.packed != u64::MAX
157 }
158}
159
160#[repr(C)]
166struct SlotHeader {
167 generation: AtomicU32,
169 next_free: AtomicU32,
171 flags: AtomicU32,
173 _reserved: u32,
175}
176
177impl SlotHeader {
178 const FLAG_ALLOCATED: u32 = 1 << 0;
179
180 fn new() -> Self {
181 Self {
182 generation: AtomicU32::new(0),
183 next_free: AtomicU32::new(u32::MAX),
184 flags: AtomicU32::new(0),
185 _reserved: 0,
186 }
187 }
188
189 #[inline]
190 fn is_allocated(&self) -> bool {
191 self.flags.load(Ordering::Acquire) & Self::FLAG_ALLOCATED != 0
192 }
193
194 #[inline]
195 fn set_allocated(&self, allocated: bool) {
196 if allocated {
197 self.flags.fetch_or(Self::FLAG_ALLOCATED, Ordering::Release);
198 } else {
199 self.flags
200 .fetch_and(!Self::FLAG_ALLOCATED, Ordering::Release);
201 }
202 }
203
204 #[inline]
205 fn increment_generation(&self) -> u32 {
206 self.generation.fetch_add(1, Ordering::Release) + 1
207 }
208}
209
210struct Slab {
216 data: NonNull<u8>,
218 slot_size: usize,
220 slot_count: usize,
222 layout: Layout,
224 free_head: AtomicU32,
226 allocated_count: AtomicUsize,
228}
229
230impl Slab {
231 fn new(size_class: SizeClass, slot_count: usize) -> Option<Self> {
233 let user_size = size_class.size_bytes();
234 let header_size = std::mem::size_of::<SlotHeader>();
235 let slot_size = header_size + user_size;
236
237 let slot_size = (slot_size + 15) & !15; let total_size = slot_size * slot_count;
241 let layout = Layout::from_size_align(total_size, 64).ok()?; let ptr = unsafe { alloc(layout) };
244 let data = NonNull::new(ptr)?;
245
246 unsafe {
248 for i in 0..slot_count {
249 let slot_ptr = data.as_ptr().add(i * slot_size);
250 let header = &mut *(slot_ptr as *mut SlotHeader);
251 *header = SlotHeader::new();
252
253 if i < slot_count - 1 {
255 header.next_free.store((i + 1) as u32, Ordering::Relaxed);
256 } else {
257 header.next_free.store(u32::MAX, Ordering::Relaxed);
258 }
259 }
260 }
261
262 Some(Self {
263 data,
264 slot_size,
265 slot_count,
266 layout,
267 free_head: AtomicU32::new(0),
268 allocated_count: AtomicUsize::new(0),
269 })
270 }
271
272 fn allocate(&self) -> Option<GenerationalHandle> {
274 loop {
275 let head = self.free_head.load(Ordering::Acquire);
276
277 if head == u32::MAX {
278 return None; }
280
281 let header = self.get_header(head as usize)?;
282 let next = header.next_free.load(Ordering::Acquire);
283
284 match self.free_head.compare_exchange_weak(
285 head,
286 next,
287 Ordering::AcqRel,
288 Ordering::Acquire,
289 ) {
290 Ok(_) => {
291 let generation = header.increment_generation();
293 header.set_allocated(true);
294 self.allocated_count.fetch_add(1, Ordering::Relaxed);
295 return Some(GenerationalHandle::new(head, generation));
296 }
297 Err(_) => continue, }
299 }
300 }
301
302 fn free(&self, handle: GenerationalHandle) -> bool {
304 let slot = handle.slot() as usize;
305
306 if slot >= self.slot_count {
307 return false;
308 }
309
310 let header = match self.get_header(slot) {
311 Some(h) => h,
312 None => return false,
313 };
314
315 if header.generation.load(Ordering::Acquire) != handle.generation() {
317 return false; }
319
320 if !header.is_allocated() {
322 return false; }
324
325 header.set_allocated(false);
326
327 loop {
329 let head = self.free_head.load(Ordering::Acquire);
330 header.next_free.store(head, Ordering::Release);
331
332 match self.free_head.compare_exchange_weak(
333 head,
334 slot as u32,
335 Ordering::AcqRel,
336 Ordering::Acquire,
337 ) {
338 Ok(_) => {
339 self.allocated_count.fetch_sub(1, Ordering::Relaxed);
340 return true;
341 }
342 Err(_) => continue,
343 }
344 }
345 }
346
347 fn get_ptr(&self, handle: GenerationalHandle) -> Option<NonNull<u8>> {
349 let slot = handle.slot() as usize;
350
351 if slot >= self.slot_count {
352 return None;
353 }
354
355 let header = self.get_header(slot)?;
356
357 if header.generation.load(Ordering::Acquire) != handle.generation() {
359 return None;
360 }
361
362 if !header.is_allocated() {
364 return None;
365 }
366
367 let header_size = std::mem::size_of::<SlotHeader>();
368 let slot_ptr = unsafe { self.data.as_ptr().add(slot * self.slot_size) };
369 let user_ptr = unsafe { slot_ptr.add(header_size) };
370
371 NonNull::new(user_ptr)
372 }
373
374 fn get_header(&self, slot: usize) -> Option<&SlotHeader> {
376 if slot >= self.slot_count {
377 return None;
378 }
379
380 let slot_ptr = unsafe { self.data.as_ptr().add(slot * self.slot_size) };
381 Some(unsafe { &*(slot_ptr as *const SlotHeader) })
382 }
383
384 fn stats(&self) -> SlabStats {
386 SlabStats {
387 slot_count: self.slot_count,
388 allocated_count: self.allocated_count.load(Ordering::Relaxed),
389 slot_size: self.slot_size,
390 }
391 }
392}
393
394impl Drop for Slab {
395 fn drop(&mut self) {
396 unsafe {
397 dealloc(self.data.as_ptr(), self.layout);
398 }
399 }
400}
401
402unsafe impl Send for Slab {}
404unsafe impl Sync for Slab {}
405
406#[derive(Debug, Clone)]
408pub struct SlabStats {
409 pub slot_count: usize,
411 pub allocated_count: usize,
413 pub slot_size: usize,
415}
416
417#[derive(Clone)]
423pub struct SlabAllocatorConfig {
424 pub initial_slots: [usize; 5],
426 pub max_slabs: usize,
428}
429
430impl Default for SlabAllocatorConfig {
431 fn default() -> Self {
432 Self {
433 initial_slots: [1024, 512, 256, 128, 64], max_slabs: 64,
435 }
436 }
437}
438
439pub struct SlabAllocator {
441 slabs: [parking_lot::RwLock<Vec<Slab>>; 5],
443 config: SlabAllocatorConfig,
445 total_allocations: AtomicU64,
447 total_frees: AtomicU64,
449}
450
451impl SlabAllocator {
452 pub fn new() -> Self {
454 Self::with_config(SlabAllocatorConfig::default())
455 }
456
457 pub fn with_config(config: SlabAllocatorConfig) -> Self {
459 let slabs = std::array::from_fn(|i| {
460 let size_class = SizeClass::all()[i];
461 let initial_slots = config.initial_slots[i];
462 let slab = Slab::new(size_class, initial_slots).expect("Failed to create initial slab");
463 parking_lot::RwLock::new(vec![slab])
464 });
465
466 Self {
467 slabs,
468 config,
469 total_allocations: AtomicU64::new(0),
470 total_frees: AtomicU64::new(0),
471 }
472 }
473
474 pub fn allocate(&self, size: usize) -> Option<(GenerationalHandle, NonNull<u8>)> {
476 let size_class = SizeClass::for_size(size)?;
477 self.allocate_from_class(size_class)
478 }
479
480 pub fn allocate_from_class(
482 &self,
483 size_class: SizeClass,
484 ) -> Option<(GenerationalHandle, NonNull<u8>)> {
485 let class_idx = size_class as usize;
486
487 {
489 let slabs = self.slabs[class_idx].read();
490 for (slab_idx, slab) in slabs.iter().enumerate() {
491 if let Some(handle) = slab.allocate() {
492 let ptr = slab.get_ptr(handle)?;
493 let full_handle = GenerationalHandle::new(
495 ((slab_idx as u32) << 24) | handle.slot(),
496 handle.generation(),
497 );
498 self.total_allocations.fetch_add(1, Ordering::Relaxed);
499 return Some((full_handle, ptr));
500 }
501 }
502 }
503
504 let mut slabs = self.slabs[class_idx].write();
506
507 for (slab_idx, slab) in slabs.iter().enumerate() {
509 if let Some(handle) = slab.allocate() {
510 let ptr = slab.get_ptr(handle)?;
511 let full_handle = GenerationalHandle::new(
512 ((slab_idx as u32) << 24) | handle.slot(),
513 handle.generation(),
514 );
515 self.total_allocations.fetch_add(1, Ordering::Relaxed);
516 return Some((full_handle, ptr));
517 }
518 }
519
520 if slabs.len() >= self.config.max_slabs {
522 return None; }
524
525 let new_slab = Slab::new(size_class, self.config.initial_slots[class_idx])?;
526 let handle = new_slab.allocate()?;
527 let ptr = new_slab.get_ptr(handle)?;
528
529 let slab_idx = slabs.len();
530 slabs.push(new_slab);
531
532 let full_handle = GenerationalHandle::new(
533 ((slab_idx as u32) << 24) | handle.slot(),
534 handle.generation(),
535 );
536 self.total_allocations.fetch_add(1, Ordering::Relaxed);
537 Some((full_handle, ptr))
538 }
539
540 pub fn free(&self, size_class: SizeClass, handle: GenerationalHandle) -> bool {
542 let class_idx = size_class as usize;
543 let slab_idx = (handle.slot() >> 24) as usize;
544 let slot = handle.slot() & 0x00FFFFFF;
545 let local_handle = GenerationalHandle::new(slot, handle.generation());
546
547 let slabs = self.slabs[class_idx].read();
548
549 if slab_idx >= slabs.len() {
550 return false;
551 }
552
553 let result = slabs[slab_idx].free(local_handle);
554 if result {
555 self.total_frees.fetch_add(1, Ordering::Relaxed);
556 }
557 result
558 }
559
560 pub fn get_ptr(
562 &self,
563 size_class: SizeClass,
564 handle: GenerationalHandle,
565 ) -> Option<NonNull<u8>> {
566 let class_idx = size_class as usize;
567 let slab_idx = (handle.slot() >> 24) as usize;
568 let slot = handle.slot() & 0x00FFFFFF;
569 let local_handle = GenerationalHandle::new(slot, handle.generation());
570
571 let slabs = self.slabs[class_idx].read();
572
573 if slab_idx >= slabs.len() {
574 return None;
575 }
576
577 slabs[slab_idx].get_ptr(local_handle)
578 }
579
580 pub fn stats(&self) -> AllocatorStats {
582 let mut class_stats = Vec::new();
583
584 for (i, size_class) in SizeClass::all().iter().enumerate() {
585 let slabs = self.slabs[i].read();
586 let slab_stats: Vec<_> = slabs.iter().map(|s| s.stats()).collect();
587 class_stats.push(SizeClassStats {
588 size_class: *size_class,
589 slab_count: slab_stats.len(),
590 total_slots: slab_stats.iter().map(|s| s.slot_count).sum(),
591 allocated_slots: slab_stats.iter().map(|s| s.allocated_count).sum(),
592 });
593 }
594
595 AllocatorStats {
596 total_allocations: self.total_allocations.load(Ordering::Relaxed),
597 total_frees: self.total_frees.load(Ordering::Relaxed),
598 class_stats,
599 }
600 }
601}
602
603impl Default for SlabAllocator {
604 fn default() -> Self {
605 Self::new()
606 }
607}
608
609#[derive(Debug, Clone)]
611pub struct SizeClassStats {
612 pub size_class: SizeClass,
614 pub slab_count: usize,
616 pub total_slots: usize,
618 pub allocated_slots: usize,
620}
621
622#[derive(Debug, Clone)]
624pub struct AllocatorStats {
625 pub total_allocations: u64,
627 pub total_frees: u64,
629 pub class_stats: Vec<SizeClassStats>,
631}
632
633pub struct TypedSlabAllocator<T> {
639 allocator: SlabAllocator,
641 size_class: SizeClass,
643 _marker: PhantomData<T>,
645}
646
647impl<T: Sized> TypedSlabAllocator<T> {
648 pub fn new() -> Option<Self> {
650 let size = std::mem::size_of::<T>();
651 let size_class = SizeClass::for_size(size)?;
652
653 Some(Self {
654 allocator: SlabAllocator::new(),
655 size_class,
656 _marker: PhantomData,
657 })
658 }
659
660 pub fn allocate(&self, value: T) -> Option<(GenerationalHandle, NonNull<T>)> {
662 let (handle, ptr) = self.allocator.allocate_from_class(self.size_class)?;
663
664 unsafe {
666 std::ptr::write(ptr.as_ptr() as *mut T, value);
667 }
668
669 Some((handle, ptr.cast()))
670 }
671
672 pub fn allocate_uninit(&self) -> Option<(GenerationalHandle, NonNull<MaybeUninit<T>>)> {
674 let (handle, ptr) = self.allocator.allocate_from_class(self.size_class)?;
675 Some((handle, ptr.cast()))
676 }
677
678 pub fn free(&self, handle: GenerationalHandle) -> bool {
680 if let Some(ptr) = self.allocator.get_ptr(self.size_class, handle) {
682 unsafe {
683 std::ptr::drop_in_place(ptr.as_ptr() as *mut T);
684 }
685 }
686
687 self.allocator.free(self.size_class, handle)
688 }
689
690 pub fn get(&self, handle: GenerationalHandle) -> Option<&T> {
692 let ptr = self.allocator.get_ptr(self.size_class, handle)?;
693 Some(unsafe { &*(ptr.as_ptr() as *const T) })
694 }
695
696 #[allow(clippy::mut_from_ref)]
706 pub unsafe fn get_mut(&self, handle: GenerationalHandle) -> Option<&mut T> {
707 let ptr = self.allocator.get_ptr(self.size_class, handle)?;
708 Some(unsafe { &mut *(ptr.as_ptr() as *mut T) })
709 }
710}
711
712impl<T: Sized> Default for TypedSlabAllocator<T> {
713 fn default() -> Self {
714 Self::new().expect("Type too large for slab allocation")
715 }
716}
717
718#[cfg(test)]
719mod tests {
720 use super::*;
721 use std::sync::Arc;
722 use std::thread;
723
724 #[test]
725 fn test_generational_handle() {
726 let handle = GenerationalHandle::new(42, 7);
727 assert_eq!(handle.slot(), 42);
728 assert_eq!(handle.generation(), 7);
729
730 let raw = handle.to_raw();
731 let handle2 = GenerationalHandle::from_raw(raw);
732 assert_eq!(handle2.slot(), 42);
733 assert_eq!(handle2.generation(), 7);
734 }
735
736 #[test]
737 fn test_size_class() {
738 assert_eq!(SizeClass::for_size(32), Some(SizeClass::Size64));
739 assert_eq!(SizeClass::for_size(64), Some(SizeClass::Size64));
740 assert_eq!(SizeClass::for_size(65), Some(SizeClass::Size128));
741 assert_eq!(SizeClass::for_size(512), Some(SizeClass::Size512));
742 assert_eq!(SizeClass::for_size(1025), None);
743 }
744
745 #[test]
746 fn test_slab_basic() {
747 let slab = Slab::new(SizeClass::Size64, 10).unwrap();
748
749 let h1 = slab.allocate().unwrap();
750 let h2 = slab.allocate().unwrap();
751 let h3 = slab.allocate().unwrap();
752
753 assert_ne!(h1.slot(), h2.slot());
754 assert_ne!(h2.slot(), h3.slot());
755
756 assert!(slab.free(h2));
757
758 let h4 = slab.allocate().unwrap();
759 assert_eq!(h4.slot(), h2.slot()); assert_ne!(h4.generation(), h2.generation()); }
762
763 #[test]
764 fn test_slab_generation() {
765 let slab = Slab::new(SizeClass::Size64, 10).unwrap();
766
767 let h1 = slab.allocate().unwrap();
768 assert!(slab.free(h1));
769
770 let ptr = slab.get_ptr(h1);
772 assert!(ptr.is_none());
773
774 assert!(!slab.free(h1));
776 }
777
778 #[test]
779 fn test_allocator_basic() {
780 let allocator = SlabAllocator::new();
781
782 let (h1, p1) = allocator.allocate(32).unwrap();
783 let (h2, p2) = allocator.allocate(100).unwrap();
784 let (h3, p3) = allocator.allocate(300).unwrap();
785
786 assert!(p1.as_ptr() != p2.as_ptr());
787 assert!(p2.as_ptr() != p3.as_ptr());
788
789 assert!(allocator.free(SizeClass::Size64, h1));
790 assert!(allocator.free(SizeClass::Size128, h2));
791 assert!(allocator.free(SizeClass::Size512, h3));
792 }
793
794 #[test]
795 fn test_allocator_concurrent() {
796 let allocator = Arc::new(SlabAllocator::new());
797 let mut handles = vec![];
798
799 for _ in 0..4 {
800 let allocator = allocator.clone();
801 handles.push(thread::spawn(move || {
802 let mut local_handles = Vec::new();
803 for _ in 0..1000 {
804 let (handle, _ptr) = allocator.allocate(64).unwrap();
805 local_handles.push(handle);
806 }
807
808 for handle in local_handles.drain(..500) {
810 allocator.free(SizeClass::Size64, handle);
811 }
812
813 local_handles
814 }));
815 }
816
817 let mut all_handles = Vec::new();
818 for handle in handles {
819 all_handles.extend(handle.join().unwrap());
820 }
821
822 let stats = allocator.stats();
823 assert_eq!(stats.total_allocations, 4000);
824 assert_eq!(stats.total_frees, 2000);
825 }
826
827 #[test]
828 fn test_typed_allocator() {
829 #[derive(Debug, PartialEq)]
830 struct TestNode {
831 value: i32,
832 next: Option<u64>,
833 }
834
835 let allocator = TypedSlabAllocator::<TestNode>::new().unwrap();
836
837 let (h1, _) = allocator
838 .allocate(TestNode {
839 value: 42,
840 next: None,
841 })
842 .unwrap();
843 let (h2, _) = allocator
844 .allocate(TestNode {
845 value: 100,
846 next: Some(h1.to_raw()),
847 })
848 .unwrap();
849
850 let node1 = allocator.get(h1).unwrap();
851 assert_eq!(node1.value, 42);
852
853 let node2 = allocator.get(h2).unwrap();
854 assert_eq!(node2.value, 100);
855 assert_eq!(node2.next, Some(h1.to_raw()));
856
857 assert!(allocator.free(h1));
858 assert!(allocator.free(h2));
859 }
860
861 #[test]
862 fn test_stats() {
863 let allocator = SlabAllocator::new();
864
865 for _ in 0..50 {
866 allocator.allocate(32).unwrap();
867 }
868 for _ in 0..30 {
869 allocator.allocate(200).unwrap();
870 }
871
872 let stats = allocator.stats();
873 assert_eq!(stats.total_allocations, 80);
874
875 let size64_stats = stats
877 .class_stats
878 .iter()
879 .find(|s| s.size_class == SizeClass::Size64)
880 .unwrap();
881 assert_eq!(size64_stats.allocated_slots, 50);
882 }
883}