attachable-slab-allocator 0.1.0

//! # Heterogeneous Master-Slave Slab Allocator
//!
//! A high-performance, low-level memory management system designed for fixed-size allocations
//! using a **Master-Slave** architecture. This implementation provides $O(1)$ allocation
//! and deallocation while supporting "decentralized" deallocation (RAII) without requiring
//! the caller to maintain a handle to the original allocator.
//!
//! ## Architecture Overview
//!
//! The allocator organizes memory into "Slabs"—contiguous blocks situated at power-of-two
//! aligned boundaries. Slabs operate in two roles:
//!
//! 1.  **Master Slab**: The primary segment. It owns the synchronization primitive ([`LockTrait`]),
//!     tracks the `ref_count` of active slaves, and maintains an intrusive list of available
//!     segments (both itself and its slaves).
//! 2.  **Slave Slab**: Secondary segments that extend the capacity of a Master. Every slave
//!     holds a pointer back to its Master to facilitate global synchronization during
//!     deallocation and to access shared metadata.
//!
//! ## Memory Layout & Alignment
//!
//! Each slab header is positioned at the very start of a memory page (or a power-of-two aligned block).
//! Using `#[repr(C)]`, the header ensures a stable layout:
//! - **Magic Sentinel**: A 64-bit value ([`MAGIC_HEADER`]) used to validate pointers and prevent
//!   use-after-free via "tombstoning" (zeroing the magic on drop).
//! - **Intrusive Metadata**: Doubly-linked list pointers are stored in-place, removing external
//!   allocation overhead for the allocator's own management.
//! - **Polymorphic Dispatch**: Function pointers (`free_behaviour`) allow the system to switch
//!   between Master and Slave deallocation logic at runtime without expensive branching.
//!
//! ## Safety & Lifetime Invariants
//!
//! *   **Ref-Counting**: A Master slab remains alive as long as it has active Slaves or an external
//!     owner (like a `SlabCache`) holds a reference.
//! *   **Thread Safety**: A central lock in the Master protects the entire hierarchy.
//! *   **Tombstoning**: Upon destruction, the magic header is zeroed. Any subsequent attempt to
//!     access the slab will fail the `validate_slab` check.

use crate::get_field_ptr;
use crate::locks::LockTrait;
use crate::prelude::*;
use crate::slots::*;
use core::alloc::Layout;
use core::cell::UnsafeCell;
use core::ptr::NonNull;
use intrusive_doubly_list::DoublyLinkPointerRaw;
use intrusive_doubly_list::IntrusiveDLinkListRaw;

/// The 64-bit sentinel value used to verify the validity of a `Slab` header.
const MAGIC_HEADER: u64 = 0xFB4E_4C07_ACCE_45A2;

/// A function pointer type used for polymorphic deallocation.
///
/// This allows the `free_slot` logic to remain generic while the specific cleanup
/// behavior (Master vs Slave) is baked into the slab header at initialization.
type FreeFunc<T, LOCK> =
    fn(slab_ptr: NonNull<Slab<T, LOCK>>, slot_ptr: NonNull<T>, slab_lay: Layout) -> Result<()>;

/// A type alias for an `UnsafeCell` wrapped optional pointer, used for intrusive linking.
type SlabPtr<T, LOCK> = UnsafeCell<Option<NonNull<Slab<T, LOCK>>>>;

/// The core Slab structure.
///
/// A `Slab` manages a contiguous block of memory divided into "slots" of type `T`.
/// It acts as either a Master or a Slave depending on its initialization and
/// its `master_ptr` field.
#[repr(C)]
pub struct Slab<T, LOCK>
where
    LOCK: LockTrait,
{
    /// Sentinel value for runtime integrity checks.
    magic: UnsafeCell<u64>,

    /// The synchronization primitive. In Slave slabs, this is initialized but unused;
    /// Slaves use the Master's lock instead.
    lock: LOCK,

    /// An intrusive list header. Only meaningful in Master slabs, tracking its subordinate Slaves.
    slave_list: IntrusiveDLinkListRaw<Slab<T, LOCK>>,

    /// Pointer back to the Master slab. If this slab is the Master, it points to itself.
    master_ptr: Option<NonNull<Slab<T, LOCK>>>,

    /// Intrusive next pointer for the doubly linked list.
    next_slab: SlabPtr<T, LOCK>,
    /// Intrusive previous pointer for the doubly linked list.
    last_slab: SlabPtr<T, LOCK>,

    /// The specific deallocation logic to use for this slab instance.
    free_behaviour: FreeFunc<T, LOCK>,

    /// Manages the bitmask/stack of free slots within this slab's memory area.
    slots: UnsafeCell<Slots<T>>,

    /// Reference count of active Slave slabs. Only tracked in the Master slab.
    ref_count: UnsafeCell<u32>,

    /// Tracks whether this slab is currently present in an intrusive list.
    is_linked: UnsafeCell<bool>,
}

/// Deallocation logic for Master slabs.
///
/// Returns a slot to the Master's internal [`Slots`] manager. If the Master is empty
/// and has no active Slaves or external references, it triggers its own deallocation.
///
/// # Errors
/// - Returns `SlabError::InvalidPointer` if the slot is out of bounds for this slab.
///
/// # Safety
/// - The `master_ptr` must be a valid, initialized Master slab.
/// - The `slot_ptr` must have been originally allocated from this specific Master slab.
pub fn free_master<T, LOCK>(
    master_ptr: NonNull<Slab<T, LOCK>>,
    slot_ptr: NonNull<T>,
    slab_lay: Layout,
) -> Result<()>
where
    LOCK: LockTrait,
{
    let lock_ptr = get_field_ptr!(master_ptr, lock).as_ptr();
    let _guard = unsafe { (*lock_ptr).lock() };

    let slots_ptr = get_field_ptr!(master_ptr, slots);

    let slots_ = slots_ptr.unsafe_ref().get();
    unsafe {
        (*slots_).try_push_slot(slot_ptr)?;
    }
    let can_drop = Slab::<T, LOCK>::can_drop(master_ptr); // master_ptr.unsafe_ref().can_drop();

    if can_drop {
        Slab::<T, LOCK>::write_magic(master_ptr, 0);

        drop(_guard);
        free_slab(master_ptr.cast(), slab_lay)?;
        Ok(())
    } else {
        let slave_list_ptr = get_field_ptr!(master_ptr, slave_list);
        IntrusiveDLinkListRaw::<Slab<T, LOCK>>::push(slave_list_ptr, master_ptr);

        Ok(())
    }
}

/// Deallocation logic for Slave slabs.
///
/// Returns a slot to the Slave's internal [`Slots`] manager. If the Slave becomes entirely empty,
/// it removes itself from the Master's list and frees its memory. This may trigger
/// a recursive drop of the Master if the Master was waiting on this specific Slave.
///
/// # Errors
/// - Returns `SlabError::FatalError` if the master pointer cannot be resolved.
/// - Returns `SlabError::InvalidPointer` if validation fails.
///
/// # Safety
/// - The `slab_ptr` must be a valid, initialized Slave slab.
/// - The `slot_ptr` must have been originally allocated from this specific Slave slab.
pub fn free_slave<T, LOCK>(
    slab_ptr: NonNull<Slab<T, LOCK>>,
    slot_ptr: NonNull<T>,
    slab_lay: Layout,
) -> Result<()>
where
    LOCK: LockTrait,
{
    let master_ptr = Slab::<T, LOCK>::try_get_master_ptr(slab_ptr)?;

    Slab::<T, LOCK>::validate_slab(master_ptr)?;

    let lock_ptr = get_field_ptr!(master_ptr, lock).as_ptr();
    let _guard = unsafe { (*lock_ptr).lock() }; //lock_ptr.unsafe_ref().lock();

    let slave_list_ptr = get_field_ptr!(master_ptr, slave_list);

    let slots_ptr = get_field_ptr!(slab_ptr, slots);
    let slots_ = slots_ptr.unsafe_ref().get();
    unsafe {
        (*slots_).try_push_slot(slot_ptr)?;
    }

    let can_drop = Slab::<T, LOCK>::can_drop(slab_ptr);

    if can_drop {
        Slab::<T, LOCK>::ref_down(master_ptr);

        IntrusiveDLinkListRaw::<Slab<T, LOCK>>::remove(slave_list_ptr, slab_ptr);

        let can_drop_master = Slab::<T, LOCK>::can_drop(master_ptr); // master_ptr.unsafe_ref().can_drop();

        Slab::<T, LOCK>::write_magic(slab_ptr, 0);

        if can_drop_master {
            Slab::<T, LOCK>::write_magic(master_ptr, 0);

            drop(_guard);
            free_slab(slab_ptr.cast(), slab_lay)?;
            free_slab(master_ptr.cast(), slab_lay)?;
        } else {
            drop(_guard);
            free_slab(slab_ptr.cast(), slab_lay)?;
        }
        Ok(())
    } else {
        IntrusiveDLinkListRaw::<Slab<T, LOCK>>::push(slave_list_ptr, slab_ptr);
        Ok(())
    }
}

impl<T, LOCK> Slab<T, LOCK>
where
    LOCK: LockTrait,
{
    const SLAB_ALIGNMENT: usize = align_of::<Slab<T, LOCK>>();
    const SLAB_HEADE_SIZE: usize = size_of::<Slab<T, LOCK>>();

    /// Static assertion to ensure the lock implementation doesn't exceed 16 bytes,
    /// keeping the slab header compact.
    const _CHECK: u8 = const {
        let lock_size = size_of::<LOCK>();
        assert!(
            lock_size <= size_of::<u128>(),
            "The Lock Size Should Be Smaller Than 16 Bytes"
        );

        0
    };

    /// Checks if the slab is eligible for deallocation.
    ///
    /// A slab is droppable if:
    /// 1. Its `ref_count` (active slaves) is zero.
    /// 2. Its internal `slots` state is `Free` (all slots returned).
    fn can_drop(slab_ptr: NonNull<Slab<T, LOCK>>) -> bool {
        let slots_ptr = get_field_ptr!(slab_ptr, slots);

        let slots_ = slots_ptr.unsafe_ref().get();
        let state = unsafe { (*slots_).get_state() };

        let ref_count_ptr = get_field_ptr!(slab_ptr, ref_count);
        let ref_count = unsafe { *ref_count_ptr.unsafe_ref().get() };

        ref_count == 0 && state == SlotsState::Free
    }

    /// Allocates a new slot of type `T` from the Master-Slave hierarchy.
    ///
    /// This method iterates through the Master's `slave_list` to find a segment with capacity.
    /// If all existing slabs are full, it automatically allocates and links a new Slave slab.
    ///
    /// # Errors
    /// - `SlabError::OutOfMemory`: If the system cannot allocate a new Slave slab.
    /// - `SlabError::FatalError`: If the Master-Slave pointers are corrupted.
    ///
    /// # Safety
    /// - `slab_ptr_master` must point to a valid Master slab.
    /// - The `slab_memory_layout` must match the layout used during initialization.
    pub fn alloc_slot(
        slab_ptr_master: NonNull<Self>,
        slab_memory_layout: Layout,
    ) -> Result<NonNull<T>> {
        Self::validate_slab(slab_ptr_master)?;

        let lock_ptr = get_field_ptr!(slab_ptr_master, lock);
        let _guard = lock_ptr.unsafe_ref().lock();

        let master_ptr = Self::try_get_master_ptr(slab_ptr_master)?;

        (master_ptr == slab_ptr_master).on_err(SlabError::FatalError)?;

        let slave_list_ptr = get_field_ptr!(slab_ptr_master, slave_list);

        while let Some(slab_ptr) = IntrusiveDLinkListRaw::<Self>::pop(slave_list_ptr) {
            let slots_ptr = get_field_ptr!(slab_ptr, slots);
            let slots_ = slots_ptr.unsafe_ref().get();

            match unsafe { (*slots_).try_pop_slot() } {
                Ok(slot) => {
                    //
                    let state = unsafe { (*slots_).get_state() };
                    if state == SlotsState::Partial {
                        IntrusiveDLinkListRaw::<Self>::push(slave_list_ptr, slab_ptr);
                    }

                    return Ok(slot);
                }
                Err(err) => match err {
                    SlabError::OutOfMemory => {
                        let state = unsafe { (*slots_).get_state() };
                        if state == SlotsState::Full {
                            IntrusiveDLinkListRaw::<Self>::push(slave_list_ptr, slab_ptr);
                        } else {
                            return Err(SlabError::FatalError);
                        }
                    }
                    _ => {
                        return Err(err);
                    }
                },
            }
        }

        let slave_slab_ptr =
            Self::alloc_slab_ptr(slab_memory_layout, Some(slab_ptr_master), free_slave)?;

        let slots_ptr = get_field_ptr!(slave_slab_ptr, slots);

        let slots_ = slots_ptr.unsafe_ref().get();

        let ptr = unsafe {
            (*slots_)
                .try_pop_slot()
                .map_err(|_| SlabError::FatalError)?
        };

        let state = unsafe { (*slots_).get_state() };
        if state == SlotsState::Full {
            IntrusiveDLinkListRaw::<Self>::remove(slave_list_ptr, slave_slab_ptr);
        }

        Ok(ptr.cast())
    }

    /// Global deallocation entry point.
    ///
    /// Given a pointer to a slot, this method calculates the slab's base address
    /// via alignment, validates the slab's integrity, and dispatches the
    /// deallocation logic to the appropriate function (Master vs Slave).
    ///
    /// # Errors
    /// - `SlabError::InvalidPointer`: If the slot does not belong to a valid Slab
    ///   or the magic header check fails.
    ///
    /// # Safety
    /// - `slot_ptr` must have been allocated by a `Slab` of the same `slab_memory_layout`.
    /// - `slab_memory_layout` must be a power-of-two size that exactly matches
    ///   the alignment of the original slab allocation.
    pub fn free_slot(slot_ptr: NonNull<T>, slab_memory_layout: Layout) -> Result<()> {
        let address = slot_ptr.as_address();
        let slab_address = address.align_down(slab_memory_layout.size());

        let slab_ptr = NonNull::<Slab<T, LOCK>>::from_address(slab_address)
            .ok_or(SlabError::InvalidPointer)?;

        // let slab_ref = slab_ptr.unsafe_ref();
        // slab_ref.validate_slab()?;
        Self::validate_slab(slab_ptr)?;

        let free_behaviour_ptr = get_field_ptr!(slab_ptr, free_behaviour);
        let free_behaviour = unsafe { free_behaviour_ptr.read() };

        (free_behaviour)(slab_ptr, slot_ptr, slab_memory_layout)
    }

    /// Allocates and initializes a new Slab segment.
    ///
    /// This method requests raw memory from the system, initializes the `Slab` header,
    /// sets up the [`Slots`] manager for the remaining area, and links it to the
    /// Master if applicable.
    ///
    /// # Errors
    /// - `SlabError::OutOfMemory`: If the system allocator fails.
    /// - `SlabError::AlignmentMismatch`: If the allocated pointer is not properly aligned.
    ///
    /// # Safety
    /// - `slab_layout` must be a power-of-two size.
    /// - The alignment of `slab_layout` must be at least `Self::SLAB_ALIGNMENT`.
    pub fn alloc_slab_ptr(
        slab_layout: Layout,
        master_ptr: Option<NonNull<Slab<T, LOCK>>>,
        free_behaviour: FreeFunc<T, LOCK>,
    ) -> Result<NonNull<Slab<T, LOCK>>> {
        let _ = Self::_CHECK;

        let buffer_ptr = alloc_slab(slab_layout).ok_or(SlabError::OutOfMemory)?;
        let buffer_size = slab_layout.size();

        buffer_ptr
            .is_aligned_on_pow2(Self::SLAB_ALIGNMENT)
            .on_err(SlabError::AlignmentMismatch)?;

        (Self::SLAB_HEADE_SIZE < buffer_size).on_err(SlabError::OutOfMemory)?;

        let slot_area_ptr = buffer_ptr.unsafe_unsafe_add(Self::SLAB_HEADE_SIZE);
        let slot_area_size = buffer_size - Self::SLAB_HEADE_SIZE;

        let slots = Slots::<T>::new(slot_area_ptr, slot_area_size)?;

        let slab_ptr = buffer_ptr.cast::<Slab<T, LOCK>>();

        let master_fix_ptr = match master_ptr {
            Some(ptr) => {
                let ref_c_ptr = get_field_ptr!(ptr, ref_count);

                unsafe { *ref_c_ptr.unsafe_ref().get() += 1 }

                ptr
            }
            None => slab_ptr,
        };

        let slab = Self {
            magic: UnsafeCell::new(MAGIC_HEADER),
            slots: UnsafeCell::new(slots),
            ref_count: UnsafeCell::new(0),
            slave_list: IntrusiveDLinkListRaw::new(),
            master_ptr: Some(master_fix_ptr),
            next_slab: UnsafeCell::new(None),
            last_slab: UnsafeCell::new(None),
            is_linked: UnsafeCell::new(false),
            lock: LOCK::init(),
            free_behaviour,
        };

        unsafe { slab_ptr.write(slab) };

        IntrusiveDLinkListRaw::init_node(slab_ptr);

        let slave_list_ptr = get_field_ptr!(master_fix_ptr, slave_list);
        IntrusiveDLinkListRaw::<Self>::push(slave_list_ptr, slab_ptr);

        Ok(slab_ptr)
    }

    /// Validates the slab's runtime integrity.
    ///
    /// Verifies that the `magic` header of the current slab and its Master
    /// (if different) match [`MAGIC_HEADER`].
    ///
    /// # Errors
    /// - `SlabError::InvalidPointer`: If any magic check fails, indicating
    ///   corruption or a use-after-free.
    #[inline]
    pub fn validate_slab(slab_ptr: NonNull<Slab<T, LOCK>>) -> Result<()> {
        Self::check_magic_raw(slab_ptr)?;

        let master_field_ptr = get_field_ptr!(slab_ptr, master_ptr);
        let master_ptr_raw = unsafe { master_field_ptr.read() };

        let Some(master_ptr) = master_ptr_raw else {
            return Ok(());
        };

        Self::check_magic_raw(master_ptr)
    }

    /// Performs a raw magic header comparison.
    #[inline]
    fn check_magic_raw(slab_ptr: NonNull<Slab<T, LOCK>>) -> Result<()> {
        let magic_ptr = get_field_ptr!(slab_ptr, magic);
        let magic = unsafe { *magic_ptr.unsafe_ref().get() };
        (magic == MAGIC_HEADER).as_result((), SlabError::InvalidPointer)
    }

    /// Atomically updates the magic header (used for tombstoning).
    fn write_magic(slab_ptr: NonNull<Slab<T, LOCK>>, val: u64) {
        let magic_ptr = get_field_ptr!(slab_ptr, magic);
        unsafe { *magic_ptr.unsafe_ref().get() = val };
    }

    /// Extracts the Master pointer from a slab.
    fn try_get_master_ptr(slab_ptr: NonNull<Slab<T, LOCK>>) -> Result<NonNull<Slab<T, LOCK>>> {
        let master_field_ptr = get_field_ptr!(slab_ptr, master_ptr);
        let master_ptr_raw = unsafe { master_field_ptr.read() };

        if let Some(master_ptr) = master_ptr_raw {
            Ok(master_ptr)
        } else {
            Err(SlabError::InvalidPointer)
        }
    }

    /// Decrements the slab's internal reference count.
    #[inline]
    pub fn ref_down(slab_ptr: NonNull<Slab<T, LOCK>>) {
        let ref_count_ptr = get_field_ptr!(slab_ptr, ref_count);
        unsafe { *ref_count_ptr.unsafe_ref().get() -= 1 };
    }

    /// Thread-safe version of `ref_down`.
    #[inline]
    #[allow(unused)]
    pub fn atomic_ref_down(slab_ptr: NonNull<Slab<T, LOCK>>) -> Result<()> {
        Self::validate_slab(slab_ptr)?;

        let lock_ptr = get_field_ptr!(slab_ptr, lock);
        let _guard = lock_ptr.unsafe_ref().lock();

        let ref_count_ptr = get_field_ptr!(slab_ptr, ref_count);
        unsafe { *ref_count_ptr.unsafe_ref().get() -= 1 };

        Ok(())
    }

    /// Thread-safe increment of the slab's reference count.
    #[inline]
    pub fn atomic_ref_up(slab_ptr: NonNull<Slab<T, LOCK>>) -> Result<()> {
        Self::validate_slab(slab_ptr)?;

        let lock_ptr = get_field_ptr!(slab_ptr, lock);
        let _guard = lock_ptr.unsafe_ref().lock();

        let ref_count_ptr = get_field_ptr!(slab_ptr, ref_count);
        unsafe { *ref_count_ptr.unsafe_ref().get() += 1 };

        Ok(())
    }

    /// Releases an external handle on a Master slab.
    ///
    /// This is typically called by a high-level manager (e.g., a Cache) when it
    /// no longer needs to use this specific Master for new allocations. If the
    /// Master is empty and has no Slaves, it will be destroyed.
    ///
    /// # Safety
    /// - Only call this if you hold a logical "reference" to the Master.
    /// - Once called, the Master may be deallocated if it has no other dependencies.
    pub fn atomic_release_master(
        master_ptr: NonNull<Slab<T, LOCK>>,
        slab_lay: Layout,
    ) -> Result<()> {
        Self::validate_slab(master_ptr)?;

        let lock_ptr = get_field_ptr!(master_ptr, lock);

        let _guard = lock_ptr.unsafe_ref().lock();

        Self::ref_down(master_ptr);

        let can_drop = Self::can_drop(master_ptr);

        if can_drop {
            Self::write_magic(master_ptr, 0);
            drop(_guard);
            free_slab(master_ptr.cast(), slab_lay)?;
        }
        Ok(())
    }
}

impl<T, LOCK> DoublyLinkPointerRaw<Slab<T, LOCK>> for Slab<T, LOCK>
where
    LOCK: LockTrait,
{
    fn get_next(node_ptr: NonNull<Slab<T, LOCK>>) -> Option<NonNull<Slab<T, LOCK>>> {
        unsafe { *node_ptr.unsafe_ref().next_slab.get() }
    }

    fn get_last(node_ptr: NonNull<Slab<T, LOCK>>) -> Option<NonNull<Slab<T, LOCK>>> {
        unsafe { *node_ptr.unsafe_ref().last_slab.get() }
    }

    fn set_next(node_ptr: NonNull<Slab<T, LOCK>>, next_ptr: Option<NonNull<Slab<T, LOCK>>>) {
        unsafe { *node_ptr.unsafe_ref().next_slab.get() = next_ptr }
    }

    fn set_last(node_ptr: NonNull<Slab<T, LOCK>>, last_ptr: Option<NonNull<Slab<T, LOCK>>>) {
        unsafe { *node_ptr.unsafe_ref().last_slab.get() = last_ptr }
    }

    fn set_link_state(node_ptr: NonNull<Slab<T, LOCK>>, state: bool) {
        unsafe { *node_ptr.unsafe_ref().is_linked.get() = state }
    }

    fn is_linked(node_ptr: NonNull<Slab<T, LOCK>>) -> bool {
        unsafe { *node_ptr.unsafe_ref().is_linked.get() }
    }
}

#[cfg(test)]
mod tests {
    extern crate std;
    use crate::{define_allocation_hooks, locks::SpinLock, prelude::*, slab::Slab};
    use std::{alloc::Layout, ptr::NonNull};

    pub fn alloc_i(_: Layout) -> Option<NonNull<u8>> {
        panic!("fake allocation only use for make linker quit");
    }

    pub fn free_i(_: NonNull<u8>, _: Layout) -> Result<()> {
        panic!("fake allocation only use for make linker quit");
    }

    define_allocation_hooks!(alloc_i, free_i);

    #[test]
    fn basic_test() {
        let size = size_of::<Slab<u128, SpinLock>>();
        assert!(size <= 96, "Slab Header Size Test failed");
    }
}