ps-alloc 0.1.0-8

use std::alloc::Layout;

use crate::{
    header::{AllocationHeader, HEADER_SIZE},
    AllocationError,
};

/// A reasonably safe implementation of `alloc`.
///
/// - Memory allocated by this function must be freed by this crate's `free`.
/// - Caller guarantees `free` is called before the returned pointer goes out of scope.
///
/// # Errors
/// - `Err(ArithmeticError)` on integer overflow.
/// - `Err(LayoutError)` if the computed layout is invalid.
/// - `Err(OutOfMemory)` if `alloc()` returns a `nullptr`.
///
/// # Safety
/// This function is safe to call, as it produces a valid pointer.
/// The data buffer being pointed to is, however, not initialized.
pub fn alloc(size: usize) -> Result<*mut u8, AllocationError> {
    // the allocation size is header + size, rounded up
    let size = size
        .checked_add(HEADER_SIZE)
        .ok_or(AllocationError::ArithmeticError)?
        .checked_next_multiple_of(HEADER_SIZE)
        .ok_or(AllocationError::ArithmeticError)?;

    // allocations are aligned to [`HEADER_SIZE`]
    let layout = Layout::from_size_align(size, HEADER_SIZE)?;

    // since we have a valid layout, the safety requirements of [`std::alloc::alloc`] are satisfied.
    let ptr = unsafe { std::alloc::alloc(layout) };

    // note that [`std::alloc::alloc`] is allowed to abort instead
    if ptr.is_null() {
        return Err(AllocationError::OutOfMemory);
    }

    // the first [`HEADER_SIZE`] bytes of the allocation are reserved for the header
    let header: &mut AllocationHeader = unsafe { &mut *ptr.cast() };

    header.mark_as_used();
    header.set_size(size);

    // the user buffer starts at `ptr + HEADER_SIZE`
    Ok(unsafe { ptr.add(HEADER_SIZE) })
}