attachable-slab-allocator 0.1.0

//! # Low-Level Utility Extensions
//!
//! This module provides the "glue" that makes safe-ish pointer arithmetic and
//! functional error handling possible in a low-level environment.
//!
//! ## Key Components
//!
//! 1. **`get_field_ptr!` Macro**: Uses `core::ptr::addr_of_mut!` to calculate the
//!    address of a struct field without creating intermediate references. This is
//!    critical for avoiding Undefined Behavior (UB) when dealing with raw memory.
//! 2. **`NonNullEx`**: Adds "memory intelligence" to raw pointers, allowing for
//!    alignment checks, address conversions, and buffer carving.
//! 3. **`NumEx` & `BoolEx`**: Extension traits that bring functional-style logic
//!    (`on_err`, `as_result`) to primitive types, reducing branching boilerplate.
//!
//! ## Safety Warning
//! These utilities are high-power tools. `as_slice_mut` and `unsafe_ref` assume
//! that the caller has already validated memory safety and alignment.

use core::ops::Rem;
use core::{
    ops::{Add, Not},
    ops::{BitAnd, Sub},
    ptr::NonNull,
    slice,
};

/// Safely calculates the address of a field from a [`NonNull`] pointer.
///
/// This macro is essential for the allocator because it allows accessing metadata
/// (like the `magic` or `slots` field) without creating a full Rust reference
/// to the struct, which might violate aliasing rules.
#[macro_export]
macro_rules! get_field_ptr {
    ($non_null_ptr:expr, $field:ident) => {{
        let ptr = $non_null_ptr;

        unsafe {
            core::ptr::NonNull::new_unchecked(core::ptr::addr_of_mut!((*ptr.as_ptr()).$field))
        }
    }};
}

/// Extension trait for [`NonNull`] pointers to simplify low-level memory operations.
pub trait NonNullEx<T> {
    type Type;
    /// Checks if the pointer's address is aligned to the specified power-of-two value.
    fn is_aligned_on_pow2(&self, align: usize) -> bool;

    /// Converts the pointer to its raw integer address.
    fn as_address(&self) -> usize;

    /// Creates a `NonNull` pointer from a raw integer address.
    /// Returns `None` if the address is zero.
    fn from_address(address: usize) -> Option<NonNull<Self::Type>>;

    /// Casts raw memory into a mutable slice.
    /// # Safety
    /// Caller must ensure the memory is valid for `size` elements and not aliased.
    #[allow(unused)]
    fn as_slice_mut<'a>(&mut self, size: usize) -> &'a mut [Self::Type];

    /// Adjusts the pointer forward to the next alignment boundary and
    /// returns the remaining buffer size.
    fn align_up_buffer(&self, align: usize, buffer_size: usize) -> Option<(NonNull<T>, usize)>;

    /// Offset-based pointer addition.
    fn unsafe_unsafe_add(&self, offset: usize) -> NonNull<T>;

    /// Dereferences the pointer into an immutable reference.
    fn unsafe_ref<'a>(&self) -> &'a Self::Type;
    /// Dereferences the pointer into a mutable reference.
    fn unsafe_mut_ref<'a>(&mut self) -> &'a mut Self::Type;

    // fn shrink_buffer(&self, buffer_size: usize, shrink_size: usize) -> Option<(NonNull<T>, usize)>;
}

/// Extension trait for numeric types (primarily `usize`) to provide
/// alignment and error-handling utilities.
pub trait NumEx {
    type Type;
    /// Efficiently checks alignment using `address & (align - 1)`.
    #[allow(unused)]
    fn is_aligned_to_pow2(&self, align: Self::Type) -> bool;

    /// General-purpose alignment check using the modulo operator.
    fn is_aligned_to(&self, align: Self::Type) -> bool;

    /// Rounds a number UP to the nearest multiple of `align`.
    #[allow(unused)]
    fn align_up(self, align: Self::Type) -> Self::Type;

    /// Rounds a number DOWN to the nearest multiple of `align`.
    /// This is the foundation of finding a Slab header from a slot pointer.
    fn align_down(self, align: Self::Type) -> Self::Type;

    /// Functional helper: Converts `0` to an Error, and non-zero to `Ok(())`.
    #[allow(unused)]
    fn zero_err<E>(self, err: E) -> Result<(), E>;
}

impl<T> NonNullEx<T> for NonNull<T> {
    fn is_aligned_on_pow2(&self, align: usize) -> bool {
        let address = self.as_ptr() as usize;

        address & (align - 1) == 0
    }

    type Type = T;

    fn as_address(&self) -> usize {
        self.as_ptr() as usize
    }

    fn from_address(address: usize) -> Option<NonNull<Self::Type>> {
        let ptr = address as *mut Self::Type;
        let n_n = NonNull::new(ptr)?;
        Some(n_n)
    }

    fn as_slice_mut<'a>(&mut self, size: usize) -> &'a mut [Self::Type] {
        unsafe {
            let slice: &'a mut [Self::Type] = slice::from_raw_parts_mut(self.as_ptr().cast(), size);
            slice
        }
    }

    fn align_up_buffer(&self, align: usize, buffer_size: usize) -> Option<(NonNull<T>, usize)> {
        let offset_elements = self.align_offset(align);

        if offset_elements == usize::MAX {
            return None;
        }

        let offset_bytes = offset_elements * size_of::<T>();

        let new_bfs = buffer_size.checked_sub(offset_bytes)?;

        if new_bfs == 0 {
            None
        } else {
            let aligned_ptr = unsafe { self.add(offset_elements) };
            Some((aligned_ptr, new_bfs))
        }
    }

    fn unsafe_unsafe_add(&self, offset: usize) -> NonNull<T> {
        unsafe { self.add(offset) }
    }

    fn unsafe_ref<'a>(&self) -> &'a Self::Type {
        unsafe { self.as_ref() }
    }

    fn unsafe_mut_ref<'a>(&mut self) -> &'a mut Self::Type {
        unsafe { self.as_mut() }
    }
}

impl<T> NumEx for T
where
    T: Copy
        + BitAnd<Output = T>
        + Sub<Output = T>
        + Add<Output = T>
        + From<u8>
        + PartialEq
        + Default
        + Not<Output = T>
        + Rem<Output = T>,
{
    type Type = T;
    fn is_aligned_to_pow2(&self, align: Self::Type) -> bool {
        let one: T = T::from(1u8);
        let zero: T = T::default();
        (*self & (align - one)) == zero
    }

    fn zero_err<E>(self, err: E) -> Result<(), E> {
        let zero = T::default();
        if self == zero { Err(err) } else { Ok(()) }
    }

    fn align_up(self, align: Self::Type) -> Self::Type {
        let one: T = T::from(1u8);
        let mask = align - one;
        (self + mask) & !mask
    }

    fn align_down(self, align: Self::Type) -> Self::Type {
        let one: T = T::from(1u8);
        let mask = align - one;
        self & !mask
    }

    fn is_aligned_to(&self, align: Self::Type) -> bool {
        let zero: T = T::default();
        (*self % align) == zero
    }
}

/// Extension trait for `bool` to simplify functional-style error handling.
pub trait BoolEx {
    /// Maps `true` to `Ok(ok)` and `false` to `Err(err)`.
    fn as_result<O, E>(&self, ok: O, err: E) -> Result<O, E>;

    /// Returns `Ok(())` if true, otherwise returns the provided error.
    /// This is frequently used for assertions in your allocator logic.
    fn on_err<E>(&self, err: E) -> Result<(), E>;
}

impl BoolEx for bool {
    fn as_result<O, E>(&self, ok: O, err: E) -> Result<O, E> {
        match self {
            true => Ok(ok),
            false => Err(err),
        }
    }

    fn on_err<E>(&self, err: E) -> Result<(), E> {
        match self {
            true => Ok(()),
            false => Err(err),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::defs::SlabError;

    #[test]
    fn test_num_ex() {
        // is_aligned_to
        assert!(16usize.is_aligned_to_pow2(8));
        assert!(!15usize.is_aligned_to_pow2(8));
        assert!(0usize.is_aligned_to_pow2(8));

        // align_up
        assert_eq!(13u32.align_up(4), 16);
        assert_eq!(16u32.align_up(4), 16);
        assert_eq!(0u32.align_up(4), 0);

        // align_down
        assert_eq!(13u32.align_down(4), 12);
        assert_eq!(16u32.align_down(4), 16);
        assert_eq!(0u32.align_down(4), 0);

        // zero_err
        assert!(10u32.zero_err(SlabError::FatalError).is_ok());
        assert!(0u32.zero_err(SlabError::FatalError).is_err());
    }

    #[test]
    fn test_bool_ex() {
        assert_eq!(true.as_result(10, 20), Ok(10));
        assert_eq!(false.as_result(10, 20), Err(20));
        assert!(true.on_err(SlabError::FatalError).is_ok());
        assert!(false.on_err(SlabError::FatalError).is_err());
    }

    #[test]
    fn test_non_null_ex() {
        let mut val = 123u64;
        let mut ptr = NonNull::new(&mut val as *mut u64).unwrap();

        assert_eq!(ptr.as_address(), &val as *const u64 as usize);
        assert_eq!(NonNull::from_address(ptr.as_address()), Some(ptr));

        assert!(ptr.is_aligned_on_pow2(1));
        let addr = ptr.as_address();
        assert_eq!(ptr.is_aligned_on_pow2(8), addr % 8 == 0);

        let slice = ptr.as_slice_mut(1);
        assert_eq!(slice.len(), 1);
        assert_eq!(slice[0], 123);
    }

    #[test]
    fn test_align_up_buffer() {
        let mut data = [0u8; 64];
        let ptr = NonNull::new(data.as_mut_ptr()).unwrap();

        if let Some((aligned_ptr, remaining)) = ptr.align_up_buffer(32, 64) {
            assert!(aligned_ptr.as_address() % 32 == 0);
            let offset = aligned_ptr.as_address() - ptr.as_address();
            assert_eq!(remaining, 64 - offset);
        }

        let mem = NonNull::new(data.as_mut_ptr()).unwrap();
        let unaligned_off = mem.align_offset(64) + 1;
        let unaligned_ptr = unsafe { mem.add(unaligned_off) };
        assert!(unaligned_ptr.align_up_buffer(64, 10).is_none());
    }
}