ps-alloc 0.1.0-9

use std::alloc::{Layout, LayoutError};

use crate::{
    error::ValidationError,
    marker::{MARKER_FREE, MARKER_USED},
    AllocationError,
};

/// Metadata header prefixed to every allocation.
///
/// Stores the allocation state via `marker` (`MARKER_USED` or `MARKER_FREE`) and the
/// total allocation size in bytes (including this header and alignment padding) in `size`.
///
/// Aligned to 16 bytes to guarantee all allocations are aligned to 16 bytes.
#[repr(align(16))]
pub struct AllocationHeader {
    marker: [u8; 8],
    size: usize,
}

impl AllocationHeader {
    /// Constructs a header for a used allocation of `size` total bytes.
    #[inline]
    pub const fn new(size: usize) -> Self {
        Self {
            marker: MARKER_USED,
            size,
        }
    }

    /// Returns the total allocation size in bytes stored in this header.
    #[inline]
    pub const fn get_size(&self) -> usize {
        self.size
    }

    /// Checks whether this header is marked as free (deallocated).
    ///
    /// Marker checks are best-effort debugging aids.
    #[inline]
    pub fn is_marked_as_free(&self) -> bool {
        self.marker == MARKER_FREE
    }

    /// Checks whether this header is marked as used (allocated).
    ///
    /// Marker checks are best-effort debugging aids.
    #[inline]
    pub fn is_marked_as_used(&self) -> bool {
        self.marker == MARKER_USED
    }

    /// Marks this header as free (deallocated).
    #[inline]
    pub const fn mark_as_free(&mut self) {
        self.marker = MARKER_FREE;
    }
}

/// The size in bytes of the hidden header prefixed to every allocation; also the
/// alignment of every allocation produced by this crate.
pub const HEADER_SIZE: usize = std::mem::size_of::<AllocationHeader>();

/// The alignment of every allocation produced by this crate.
pub(crate) const ALIGN: usize = std::mem::align_of::<AllocationHeader>();

// The user pointer is offset from the allocation's base by `HEADER_SIZE`, so the two
// constants must be equal for the user pointer to be `ALIGN`-aligned.
const _: () = assert!(HEADER_SIZE == ALIGN);

/// Computes the total allocation size for a user buffer of `size` bytes: the header plus
/// the buffer, rounded up to a multiple of [`ALIGN`].
///
/// # Errors
/// - `Err(ArithmeticError)` on integer overflow.
pub(crate) fn total_size(size: usize) -> Result<usize, AllocationError> {
    size.checked_add(HEADER_SIZE)
        .ok_or(AllocationError::ArithmeticError)?
        .checked_next_multiple_of(ALIGN)
        .ok_or(AllocationError::ArithmeticError)
}

/// Computes the [`Layout`] of an allocation of `total` bytes.
///
/// # Errors
/// - `Err(LayoutError)` if `total` overflows [`isize::MAX`] when rounded up to [`ALIGN`].
pub(crate) fn layout_of(total: usize) -> Result<Layout, LayoutError> {
    Layout::from_size_align(total, ALIGN)
}

/// Writes a used-allocation header of `total` bytes at `base` and returns the user
/// pointer at `base + HEADER_SIZE`.
///
/// # Safety
/// `base` must be non-null, [`ALIGN`]-aligned, and valid for writes of at least `total`
/// bytes, where `total` is at least [`HEADER_SIZE`].
pub(crate) unsafe fn initialize_block(base: *mut u8, total: usize) -> *mut u8 {
    // SAFETY: `base` is valid for writes of at least `HEADER_SIZE` bytes and is
    // sufficiently aligned, per this function's contract.
    unsafe {
        base.cast::<AllocationHeader>()
            .write(AllocationHeader::new(total))
    };

    // SAFETY: the allocation is at least `HEADER_SIZE` bytes, so the offset stays in bounds.
    unsafe { base.add(HEADER_SIZE) }
}

/// Validates a non-null user pointer and returns the pointer to its allocation header.
///
/// # Errors
/// - `Err(ImproperAlignment)` if `ptr` is not [`ALIGN`]-aligned.
/// - `Err(MarkerFree)` if the header is marked as free.
/// - `Err(MarkerCorrupted)` if the header matches neither marker.
///
/// # Safety
/// `ptr` must be a non-null pointer produced by this crate's allocation functions.
/// Unless `ptr` is misaligned, this function dereferences the header below it; passing
/// any other pointer causes undefined behaviour.
pub(crate) unsafe fn validate(ptr: *mut u8) -> Result<*mut AllocationHeader, ValidationError> {
    // SAFETY: the caller guarantees `ptr` was allocated by [`crate::alloc`], so the
    // allocation begins `HEADER_SIZE` bytes below it and the subtraction stays in bounds.
    let header_ptr: *mut AllocationHeader = unsafe { ptr.sub(HEADER_SIZE) }.cast();

    if !header_ptr.is_aligned() {
        return Err(ValidationError::ImproperAlignment);
    }

    // SAFETY: the caller guarantees `ptr` was allocated by [`crate::alloc`], so
    // `header_ptr` points to a live, initialized header.
    if unsafe { &*header_ptr }.is_marked_as_free() {
        return Err(ValidationError::MarkerFree);
    }

    // SAFETY: as above.
    if !unsafe { &*header_ptr }.is_marked_as_used() {
        return Err(ValidationError::MarkerCorrupted);
    }

    Ok(header_ptr)
}