nexus_slab/

bounded.rs

//! Fixed-capacity slab allocator.
//!
//! This module provides a bounded (fixed-capacity) slab allocator.
//!
//! # Example
//!
//! ```
//! use nexus_slab::bounded::Slab;
//!
//! // SAFETY: caller guarantees slab contract (see struct docs)
//! let slab = unsafe { Slab::with_capacity(1024) };
//! let slot = slab.alloc(42u64);
//! assert_eq!(*slot, 42);
//! slab.free(slot);
//! ```

use core::cell::Cell;
use core::fmt;
use core::mem;
use core::ptr;

use alloc::vec::Vec;

use crate::shared::{Full, Slot, SlotCell};

// =============================================================================
// Claim
// =============================================================================

/// A claimed slot that has not yet been written to.
///
/// Created by [`Slab::claim()`]. Must be consumed via [`write()`](Self::write)
/// to complete the allocation. If dropped without calling `write()`, the slot
/// is returned to the freelist.
///
/// The `write()` method is `#[inline]`, enabling the compiler to potentially
/// optimize the value write as a placement new (constructing directly into
/// the slot memory).
pub struct Claim<'a, T> {
    slot_ptr: *mut SlotCell<T>,
    slab: &'a Slab<T>,
}

impl<T> Claim<'_, T> {
    /// Writes the value to the claimed slot and returns the [`Slot`] handle.
    ///
    /// This consumes the claim. The value is written directly to the slot's
    /// memory, which may enable placement new optimization.
    #[inline]
    pub fn write(self, value: T) -> Slot<T> {
        let slot_ptr = self.slot_ptr;
        // SAFETY: We own this slot from claim(), it's valid and vacant
        unsafe {
            (*slot_ptr).write_value(value);
        }
        // Don't run Drop - we're completing the allocation
        mem::forget(self);
        // SAFETY: slot_ptr is valid and now occupied
        unsafe { Slot::from_ptr(slot_ptr) }
    }
}

impl<T> fmt::Debug for Claim<'_, T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("Claim")
            .field("slot_ptr", &self.slot_ptr)
            .finish()
    }
}

impl<T> Drop for Claim<'_, T> {
    fn drop(&mut self) {
        // Abandoned claim - return slot to freelist
        // SAFETY: slot_ptr is valid and still vacant (never written to)
        let free_head = self.slab.free_head.get();
        unsafe {
            (*self.slot_ptr).set_next_free(free_head);
        }
        self.slab.free_head.set(self.slot_ptr);
    }
}

// =============================================================================
// Slab
// =============================================================================

/// Fixed-capacity slab allocator for manual memory management.
///
/// Construction is `unsafe` — by creating a slab, you accept the contract:
///
/// - **Free everything you allocate.** Dropping the slab does NOT drop
///   values in occupied slots. Unfree'd slots leak silently.
/// - **Free from the same slab.** Passing a [`Slot`] to a different
///   slab's `free()` corrupts the freelist.
/// - **Don't share across threads.** The slab is `!Send` and `!Sync`.
///
/// In practice, most systems have one slab per type in `thread_local!`
/// storage. Cross-slab free is a programming error caught by `debug_assert`.
///
/// Uses pointer-based freelist for O(1) allocation. ~20-24 cycle operations.
///
/// # Const Construction
///
/// Supports const construction via [`new()`](Self::new) followed by
/// runtime initialization via [`init()`](Self::init). This enables use with
/// `thread_local!` using the `const { }` block syntax for zero-overhead TLS access.
///
/// ```no_run
/// use nexus_slab::bounded::Slab;
///
/// struct MyType(u64);
///
/// // SAFETY: single slab per type, freed before thread exit
/// thread_local! {
///     static SLAB: Slab<MyType> = const { unsafe { Slab::new() } };
/// }
///
/// // Later, at runtime:
/// SLAB.with(|s| unsafe { s.init(1024) });
/// ```
///
/// For direct usage, prefer [`with_capacity()`](Self::with_capacity).
pub struct Slab<T> {
    /// Slot storage. Wrapped in UnsafeCell for interior mutability.
    slots: core::cell::UnsafeCell<Vec<SlotCell<T>>>,
    /// Capacity. Wrapped in Cell so it can be set during init.
    capacity: Cell<usize>,
    /// Head of freelist - raw pointer for fast allocation.
    /// NULL when slab is full or uninitialized.
    pub(crate) free_head: Cell<*mut SlotCell<T>>,
}

impl<T> Slab<T> {
    /// Creates an empty, uninitialized slab.
    ///
    /// This is a const function that performs no allocation. Call [`init()`](Self::init)
    /// to allocate storage before use.
    ///
    /// # Safety
    ///
    /// See [struct-level safety contract](Self).
    ///
    /// # Example
    ///
    /// ```no_run
    /// use nexus_slab::bounded::Slab;
    ///
    /// // SAFETY: single slab per type, freed before thread exit
    /// thread_local! {
    ///     static SLAB: Slab<u64> = const { unsafe { Slab::new() } };
    /// }
    /// ```
    #[inline]
    pub const unsafe fn new() -> Self {
        Self {
            slots: core::cell::UnsafeCell::new(Vec::new()),
            capacity: Cell::new(0),
            free_head: Cell::new(ptr::null_mut()),
        }
    }

    /// Creates a new slab with the given capacity.
    ///
    /// # Safety
    ///
    /// See [struct-level safety contract](Self).
    ///
    /// # Panics
    ///
    /// Panics if capacity is zero.
    #[inline]
    pub unsafe fn with_capacity(capacity: usize) -> Self {
        // SAFETY: caller upholds the slab contract
        let slab = unsafe { Self::new() };
        // SAFETY: caller upholds the slab contract
        unsafe { slab.init(capacity) };
        slab
    }

    /// Initializes the slab with the given capacity.
    ///
    /// This allocates slot storage and builds the freelist. Must be called
    /// exactly once before any allocations.
    ///
    /// # Safety
    ///
    /// See [struct-level safety contract](Self).
    ///
    /// # Panics
    ///
    /// - Panics if the slab is already initialized (capacity > 0)
    /// - Panics if capacity is zero
    pub unsafe fn init(&self, capacity: usize) {
        assert!(self.capacity.get() == 0, "Slab already initialized");
        assert!(capacity > 0, "capacity must be non-zero");

        // SAFETY: We have &self and verified capacity == 0, so no other code
        // can be accessing slots. This is the only mutation point.
        let slots = unsafe { &mut *self.slots.get() };

        // Allocate slots — all initially vacant
        slots.reserve_exact(capacity);
        for _ in 0..capacity {
            slots.push(SlotCell::vacant(ptr::null_mut()));
        }

        // Wire up the freelist: each slot's next_free points to the next slot
        for i in 0..(capacity - 1) {
            let next_ptr = slots.as_mut_ptr().wrapping_add(i + 1);
            // SAFETY: Slot is vacant, wiring up the freelist during init
            unsafe { slots[i].set_next_free(next_ptr) };
        }
        // Last slot points to NULL (end of freelist) — already null from vacant()

        let free_head = slots.as_mut_ptr();
        self.capacity.set(capacity);
        self.free_head.set(free_head);
    }

    /// Returns true if the slab has been initialized.
    #[inline]
    pub fn is_initialized(&self) -> bool {
        self.capacity.get() > 0
    }

    /// Returns the capacity.
    #[inline]
    pub fn capacity(&self) -> usize {
        self.capacity.get()
    }

    /// Returns the base pointer to the slots array.
    #[inline]
    pub(crate) fn slots_ptr(&self) -> *mut SlotCell<T> {
        // SAFETY: We're returning a pointer for use with raw pointer access
        let slots = unsafe { &*self.slots.get() };
        slots.as_ptr().cast_mut()
    }

    /// Returns `true` if `ptr` falls within this slab's slot array.
    ///
    /// O(1) range check. Used in `debug_assert!` to validate provenance.
    #[doc(hidden)]
    #[inline]
    pub fn contains_ptr(&self, ptr: *const ()) -> bool {
        let base = self.slots_ptr() as usize;
        let end = base + self.capacity.get() * core::mem::size_of::<SlotCell<T>>();
        let addr = ptr as usize;
        addr >= base && addr < end
    }

    // =========================================================================
    // Allocation API
    // =========================================================================

    /// Claims a slot from the freelist without writing a value.
    ///
    /// Returns `None` if the slab is full. The returned [`Claim`] must be
    /// consumed via [`Claim::write()`] to complete the allocation.
    ///
    /// This two-phase allocation enables placement new optimization: the
    /// value can be constructed directly into the slot memory.
    ///
    /// # Example
    ///
    /// ```
    /// use nexus_slab::bounded::Slab;
    ///
    /// // SAFETY: caller guarantees slab contract (see struct docs)
    /// let slab = unsafe { Slab::with_capacity(10) };
    /// if let Some(claim) = slab.claim() {
    ///     let slot = claim.write(42u64);
    ///     assert_eq!(*slot, 42);
    ///     slab.free(slot);
    /// }
    /// ```
    #[inline]
    pub fn claim(&self) -> Option<Claim<'_, T>> {
        self.claim_ptr().map(|slot_ptr| Claim {
            slot_ptr,
            slab: self,
        })
    }

    /// Claims a slot from the freelist, returning the raw pointer.
    ///
    /// Returns `None` if the slab is full. This is a low-level API for
    /// macro-generated code that needs to escape TLS closures.
    ///
    /// # Safety Contract
    ///
    /// The caller MUST either:
    /// - Write a value to the slot and use it as an allocated slot, OR
    /// - Return the pointer to the freelist via `free_ptr()` if abandoning
    #[doc(hidden)]
    #[inline]
    pub(crate) fn claim_ptr(&self) -> Option<*mut SlotCell<T>> {
        let slot_ptr = self.free_head.get();

        if slot_ptr.is_null() {
            return None;
        }

        // SAFETY: slot_ptr came from the freelist within this slab.
        // The slot is vacant, so next_free is the active union field.
        let next_free = unsafe { (*slot_ptr).get_next_free() };

        // Update freelist head
        self.free_head.set(next_free);

        Some(slot_ptr)
    }

    /// Allocates a slot and writes the value.
    ///
    /// # Panics
    ///
    /// Panics if the slab is full.
    #[inline]
    pub fn alloc(&self, value: T) -> Slot<T> {
        self.claim().expect("slab full").write(value)
    }

    /// Tries to allocate a slot and write the value.
    ///
    /// Returns `Err(Full(value))` if the slab is at capacity.
    #[inline]
    pub fn try_alloc(&self, value: T) -> Result<Slot<T>, Full<T>> {
        match self.claim() {
            Some(claim) => Ok(claim.write(value)),
            None => Err(Full(value)),
        }
    }

    /// Frees a slot, dropping the value and returning storage to the freelist.
    ///
    /// Consumes the handle — the slot cannot be used after this call.
    /// The caller's safety obligation (free from the correct slab) was
    /// accepted at construction time.
    #[inline]
    #[allow(clippy::needless_pass_by_value)]
    pub fn free(&self, slot: Slot<T>) {
        let slot_ptr = slot.into_raw();
        debug_assert!(
            self.contains_ptr(slot_ptr as *const ()),
            "slot was not allocated from this slab"
        );
        // SAFETY: Caller guarantees slot is valid and occupied
        unsafe {
            (*slot_ptr).drop_value_in_place();
            self.free_ptr(slot_ptr);
        }
    }

    /// Frees a slot and returns the value without dropping it.
    ///
    /// Consumes the handle — the slot cannot be used after this call.
    #[inline]
    #[allow(clippy::needless_pass_by_value)]
    pub fn take(&self, slot: Slot<T>) -> T {
        let slot_ptr = slot.into_raw();
        debug_assert!(
            self.contains_ptr(slot_ptr as *const ()),
            "slot was not allocated from this slab"
        );
        // SAFETY: Caller guarantees slot is valid and occupied
        unsafe {
            let value = (*slot_ptr).read_value();
            self.free_ptr(slot_ptr);
            value
        }
    }

    /// Returns a slot to the freelist by pointer.
    ///
    /// Does NOT drop the value — caller must drop before calling.
    ///
    /// # Safety
    ///
    /// - `slot_ptr` must point to a slot within this slab
    /// - Value must already be dropped or moved out
    #[doc(hidden)]
    #[inline]
    pub(crate) unsafe fn free_ptr(&self, slot_ptr: *mut SlotCell<T>) {
        debug_assert!(
            self.contains_ptr(slot_ptr as *const ()),
            "slot was not allocated from this slab"
        );
        let free_head = self.free_head.get();
        // SAFETY: Caller guarantees slot_ptr is valid, transitioning to vacant
        unsafe {
            (*slot_ptr).set_next_free(free_head);
        }
        self.free_head.set(slot_ptr);
    }
}

impl<T> fmt::Debug for Slab<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("Slab")
            .field("capacity", &self.capacity.get())
            .finish()
    }
}

// =============================================================================
// Tests
// =============================================================================

#[cfg(test)]
mod tests {
    use super::*;
    use std::borrow::{Borrow, BorrowMut};

    #[test]
    fn slab_basic() {
        let slab = unsafe { Slab::<u64>::with_capacity(100) };
        assert_eq!(slab.capacity(), 100);

        let slot = slab.alloc(42);
        assert_eq!(*slot, 42);
        slab.free(slot);
    }

    #[test]
    fn slab_full() {
        let slab = unsafe { Slab::<u64>::with_capacity(2) };
        let s1 = slab.alloc(1);
        let s2 = slab.alloc(2);

        let result = slab.try_alloc(3);
        assert!(result.is_err());
        let recovered = result.unwrap_err().into_inner();
        assert_eq!(recovered, 3);

        slab.free(s1);
        slab.free(s2);
    }

    #[test]
    fn slot_deref_mut() {
        let slab = unsafe { Slab::<String>::with_capacity(10) };
        let mut slot = slab.alloc("hello".to_string());
        slot.push_str(" world");
        assert_eq!(&*slot, "hello world");
        slab.free(slot);
    }

    #[test]
    fn slot_dealloc_take() {
        let slab = unsafe { Slab::<String>::with_capacity(10) };
        let slot = slab.alloc("hello".to_string());

        let value = slab.take(slot);
        assert_eq!(value, "hello");
    }

    #[test]
    fn slot_size() {
        assert_eq!(std::mem::size_of::<Slot<u64>>(), 8);
    }

    #[test]
    fn slab_debug() {
        let slab = unsafe { Slab::<u64>::with_capacity(10) };
        let s = slab.alloc(42);
        let debug = format!("{:?}", slab);
        assert!(debug.contains("Slab"));
        assert!(debug.contains("capacity"));
        slab.free(s);
    }

    #[test]
    fn borrow_traits() {
        let slab = unsafe { Slab::<u64>::with_capacity(10) };
        let mut slot = slab.alloc(42);

        let borrowed: &u64 = slot.borrow();
        assert_eq!(*borrowed, 42);

        let borrowed_mut: &mut u64 = slot.borrow_mut();
        *borrowed_mut = 100;
        assert_eq!(*slot, 100);

        slab.free(slot);
    }

    #[cfg(debug_assertions)]
    #[test]
    fn slot_debug_drop_panics() {
        let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
            let slab = unsafe { Slab::<u64>::with_capacity(10) };
            let _slot = slab.alloc(42u64);
            // slot drops here without being freed
        }));
        assert!(result.is_err(), "Slot should panic on drop in debug mode");
    }

    #[test]
    fn capacity_one() {
        let slab = unsafe { Slab::<u64>::with_capacity(1) };

        assert_eq!(slab.capacity(), 1);

        let slot = slab.alloc(42);
        assert!(slab.try_alloc(100).is_err());

        slab.free(slot);

        let slot2 = slab.alloc(100);
        assert_eq!(*slot2, 100);
        slab.free(slot2);
    }
}