stack-arena 0.12.0

use crate::AllocError;
use std::{alloc::Layout, ptr::NonNull};

/// A trait for memory allocators that manage memory in a stack-like fashion.
///
/// This trait defines the core interface for memory allocation and deallocation
/// operations used by the arena allocator system. Implementors of this trait
/// can be used interchangeably in contexts that require memory allocation.
///
/// # Implementation
///
/// The trait provides a complete interface for memory management:
///
/// - `allocate` - Allocates memory with the specified layout
/// - `allocate_zeroed` - Allocates zeroed memory with the specified layout
/// - `deallocate` - Deallocates previously allocated memory
/// - `grow` - Grows a previously allocated memory block to a larger size
/// - `grow_zeroed` - Grows a memory block and zeroes the new memory
/// - `shrink` - Shrinks a previously allocated memory block to a smaller size
///
/// # Safety
///
/// This trait uses unsafe methods because it deals directly with memory allocation
/// and raw pointers. Implementors must ensure that:
///
/// - Allocated memory is properly aligned and sized according to the layout
/// - Deallocated pointers were previously allocated by the same allocator
/// - Memory is not used after deallocation
/// - Grow and shrink operations maintain the validity of the memory
/// - All memory safety invariants are preserved
///
/// # Examples
///
/// This trait is implemented by [`StackArena`](crate::StackArena) and [`BufferArena`](crate::BufferArena):
///
/// ```
/// use stack_arena::{StackArena, Allocator};
/// use std::alloc::Layout;
///
/// let mut arena = StackArena::new();
/// let layout = Layout::from_size_align(16, 8).unwrap();
///
/// // Allocate memory
/// let ptr = unsafe { arena.allocate(layout).unwrap() };
///
/// // Use the memory...
///
/// // Deallocate when done
/// unsafe { arena.deallocate(ptr.cast(), layout) };
/// ```
pub trait Allocator {
    /// Allocates memory with the specified layout.
    ///
    /// # Safety
    ///
    /// This method is unsafe because it returns a raw pointer that must be used
    /// carefully to avoid memory safety issues.
    ///
    /// # Parameters
    ///
    /// * `layout` - The layout of the memory to allocate
    ///
    /// # Returns
    ///
    /// A non-null pointer to the allocated memory on success, or an error if
    /// allocation fails.
    unsafe fn allocate(&mut self, layout: Layout) -> Result<NonNull<[u8]>, AllocError>;

    /// Allocates zeroed memory with the specified layout.
    ///
    /// This method allocates memory and initializes it with zeros.
    ///
    /// # Safety
    ///
    /// This method is unsafe for the same reasons as `allocate`.
    ///
    /// # Parameters
    ///
    /// * `layout` - The layout of the memory to allocate
    ///
    /// # Returns
    ///
    /// A non-null pointer to the allocated zeroed memory on success, or an error if
    /// allocation fails.
    #[inline]
    unsafe fn allocate_zeroed(&mut self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
        let ptr = unsafe { self.allocate(layout) }?;
        unsafe { std::ptr::write_bytes(ptr.as_ptr() as *mut u8, 0, layout.size()) };
        Ok(ptr)
    }

    /// Deallocates memory previously allocated by this allocator.
    ///
    /// # Safety
    ///
    /// This method is unsafe because it must be called with a pointer returned
    /// by `allocate` and the same layout that was used to allocate it.
    ///
    /// # Parameters
    ///
    /// * `ptr` - A pointer to the memory to deallocate
    /// * `layout` - The layout used to allocate the memory
    unsafe fn deallocate(&mut self, ptr: NonNull<u8>, layout: Layout);

    /// Grows a previously allocated memory block to a larger size.
    ///
    /// # Safety
    ///
    /// This method is unsafe because it must be called with a pointer returned
    /// by `allocate` and the same layout that was used to allocate it.
    ///
    /// # Parameters
    ///
    /// * `ptr` - A pointer to the memory to grow
    /// * `old_layout` - The layout used to allocate the memory
    /// * `new_layout` - The new layout for the memory
    ///
    /// # Returns
    ///
    /// A non-null pointer to the grown memory on success, or an error if
    /// the operation fails.
    unsafe fn grow(
        &mut self,
        ptr: NonNull<u8>,
        old_layout: Layout,
        new_layout: Layout,
    ) -> Result<NonNull<[u8]>, AllocError>;

    /// Grows a previously allocated memory block to a larger size and zeroes the new memory.
    ///
    /// # Safety
    ///
    /// This method is unsafe for the same reasons as `grow`.
    ///
    /// # Parameters
    ///
    /// * `ptr` - A pointer to the memory to grow
    /// * `old_layout` - The layout used to allocate the memory
    /// * `new_layout` - The new layout for the memory
    ///
    /// # Returns
    ///
    /// A non-null pointer to the grown memory on success, or an error if
    /// the operation fails.
    #[inline]
    unsafe fn grow_zeroed(
        &mut self,
        ptr: NonNull<u8>,
        old_layout: Layout,
        new_layout: Layout,
    ) -> Result<NonNull<[u8]>, AllocError> {
        let ptr = unsafe { self.grow(ptr, old_layout, new_layout) }?;
        unsafe {
            ptr.cast::<u8>()
                .write_bytes(0, new_layout.size() - old_layout.size())
        };
        Ok(ptr)
    }

    /// Shrinks a previously allocated memory block to a smaller size.
    ///
    /// # Safety
    ///
    /// This method is unsafe because it must be called with a pointer returned
    /// by `allocate` and the same layout that was used to allocate it.
    ///
    /// # Parameters
    ///
    /// * `ptr` - A pointer to the memory to shrink
    /// * `old_layout` - The layout used to allocate the memory
    /// * `new_layout` - The new layout for the memory
    ///
    /// # Returns
    ///
    /// A non-null pointer to the shrunk memory on success, or an error if
    /// the operation fails.
    unsafe fn shrink(
        &mut self,
        ptr: NonNull<u8>,
        old_layout: Layout,
        new_layout: Layout,
    ) -> Result<NonNull<[u8]>, AllocError>;

    /// Returns a reference to this allocator.
    ///
    /// This method is used to get a reference to the allocator for use in
    /// contexts that require a reference.
    ///
    /// # Returns
    ///
    /// A reference to this allocator.
    #[inline]
    fn by_ref(&self) -> &Self
    where
        Self: Sized,
    {
        self
    }
}