Compiler Explorer

Source code

mod slab {
    use std::alloc::Layout;
    use std::cell::Cell;
    use std::marker::PhantomData;
    use std::mem::ManuallyDrop;
    use std::ops::{Deref, DerefMut};
    use std::ptr::NonNull;

type SlotPtr<T> = Option<NonNull<Slot<T>>>;

union Slot<T> {
        next: SlotPtr<T>,
        value: ManuallyDrop<T>,
    }

struct FreeList<T> {
        head: SlotPtr<T>,
        usage: usize,
    }
    impl<T> Clone for FreeList<T> {
        fn clone(&self) -> Self {
            *self
        }
    }
    impl<T> Copy for FreeList<T> {}

pub struct Slab<T> {
        object_store: NonNull<Slot<T>>,
        store_layout: Layout,
        free_list: Cell<FreeList<T>>,
    }

impl<T> Slab<T> {
        pub fn new(capacity: usize) -> Self {
            assert!(capacity > 0);

unsafe {
                let store_layout = Layout::array::<Slot<T>>(capacity).unwrap();
                let ptr = std::alloc::alloc(store_layout);
                if ptr.is_null() {
                    std::alloc::handle_alloc_error(store_layout);
                }

let object_store = NonNull::new_unchecked(object_store);
                Slab {
                    object_store,
                    store_layout,
                    free_list: Cell::new(FreeList {
                        head: Some(object_store),
                        usage: 0,
                    }),
                }
            }
        }

pub fn allocate(&self, value: T) -> Option<Handle<'_, T>> {
            unsafe {
                let free_list = self.free_list.get();
                let head = free_list.head?;

let empty_slot = head.as_ptr().read();
                self.free_list.set(FreeList {
                    head: empty_slot.next,
                    usage: free_list.usage + 1,
                });

head.as_ptr().write(Slot {
                    value: ManuallyDrop::new(value),
                });

Some(Handle {
                    item: head,
                    free_list: &self.free_list,
                    _marker: PhantomData,
                })
            }
        }
    }

impl<T> Drop for Slab<T> {
        fn drop(&mut self) {
            unsafe { std::alloc::dealloc(self.object_store.as_ptr().cast(), self.store_layout) }
        }
    }

pub struct Handle<'slab, T> {
        item: NonNull<Slot<T>>,
        free_list: &'slab Cell<FreeList<T>>,
        _marker: PhantomData<T>,
    }

impl<T> Drop for Handle<'_, T> {
        fn drop(&mut self) {
            unsafe {
                let free_list = self.free_list.get();
                ManuallyDrop::drop(&mut self.item.as_mut().value);
                self.item.as_ptr().write(Slot {
                    next: free_list.head,
                });

self.free_list.set(FreeList {
                    head: Some(self.item),
                    usage: free_list.usage - 1,
                })
            }
        }
    }

impl<T> Deref for Handle<'_, T> {
        type Target = T;
        fn deref(&self) -> &T {
            unsafe { &self.item.as_ref().value }
        }
    }

impl<T> DerefMut for Handle<'_, T> {
        fn deref_mut(&mut self) -> &mut T {
            unsafe { &mut self.item.as_mut().value }
        }
    }
}

pub fn make_slab(capacity: usize) -> slab::Slab<i32> {
    slab::Slab::new(capacity)
}

pub fn alloc_slot(s: &slab::Slab<i32>, v: i32) -> Option<slab::Handle<'_, i32>> {
    s.allocate(v)
}

pub fn free_slot(_: slab::Handle<'_, i32>) {}

fn main() {
    let slab = slab::Slab::new(3);

let mut a = slab.allocate(32).unwrap();
    let b = slab.allocate(56).unwrap();
    let c = slab.allocate(78).unwrap();
    let d = slab.allocate(45);

dbg!(*a, *b, *c, d.is_none());
    drop(b);

*a = 65;
    let e = slab.allocate(91).unwrap();
    dbg!(*a, *c, *e);
}