ax_allocator/

lib.rs

//! Various allocator algorithms in a unified interface.
//!
//! There are three types of allocators:
//!
//! - [`ByteAllocator`]: Byte-granularity memory allocator. (e.g.,
//!   [`BuddyByteAllocator`], [`SlabByteAllocator`])
//! - [`PageAllocator`]: Page-granularity memory allocator. (e.g.,
//!   [`BitmapPageAllocator`])
//! - [`IdAllocator`]: Used to allocate unique IDs.

#![no_std]
#![cfg_attr(feature = "allocator_api", feature(allocator_api))]

#[cfg(feature = "bitmap")]
mod bitmap;
#[cfg(feature = "bitmap")]
pub use bitmap::BitmapPageAllocator;

#[cfg(feature = "buddy")]
mod buddy;
#[cfg(feature = "buddy")]
pub use buddy::BuddyByteAllocator;

#[cfg(feature = "slab")]
mod slab;
#[cfg(feature = "slab")]
pub use slab::SlabByteAllocator;

#[cfg(feature = "tlsf")]
mod tlsf;
use core::{alloc::Layout, ptr::NonNull};

#[cfg(feature = "ax-errno")]
use ax_errno::AxError;
#[cfg(feature = "tlsf")]
pub use tlsf::TlsfByteAllocator;

/// The error type used for allocation.
#[derive(Debug)]
pub enum AllocError {
    /// Invalid `size` or alignment. (e.g. unaligned)
    InvalidParam,
    /// Memory added by `add_memory` overlapped with existed memory.
    MemoryOverlap,
    /// No enough memory to allocate.
    NoMemory,
    /// Deallocate an unallocated memory region.
    NotAllocated,
}

#[cfg(feature = "ax-errno")]
impl From<AllocError> for AxError {
    fn from(value: AllocError) -> Self {
        match value {
            AllocError::NoMemory => AxError::NoMemory,
            _ => AxError::InvalidInput,
        }
    }
}

/// A [`Result`] type with [`AllocError`] as the error type.
pub type AllocResult<T = ()> = Result<T, AllocError>;

/// The base allocator inherited by other allocators.
pub trait BaseAllocator {
    /// Initialize the allocator with a free memory region.
    fn init(&mut self, start: usize, size: usize);

    /// Add a free memory region to the allocator.
    fn add_memory(&mut self, start: usize, size: usize) -> AllocResult;
}

/// Byte-granularity allocator.
pub trait ByteAllocator: BaseAllocator {
    /// Allocate memory with the given size (in bytes) and alignment.
    fn alloc(&mut self, layout: Layout) -> AllocResult<NonNull<u8>>;

    /// Deallocate memory at the given position, size, and alignment.
    fn dealloc(&mut self, pos: NonNull<u8>, layout: Layout);

    /// Returns total memory size in bytes.
    fn total_bytes(&self) -> usize;

    /// Returns allocated memory size in bytes.
    fn used_bytes(&self) -> usize;

    /// Returns available memory size in bytes.
    fn available_bytes(&self) -> usize;
}

/// Page-granularity allocator.
pub trait PageAllocator: BaseAllocator {
    /// The size of a memory page.
    const PAGE_SIZE: usize;

    /// Allocate contiguous memory pages with given count and byte alignment.
    ///
    /// `align` is the requested alignment in bytes, not a log2/exponent.
    fn alloc_pages(&mut self, num_pages: usize, align: usize) -> AllocResult<usize>;

    /// Deallocate contiguous memory pages with given position and count.
    fn dealloc_pages(&mut self, pos: usize, num_pages: usize);

    /// Allocate contiguous memory pages with given base address, count and byte alignment.
    ///
    /// `align` is the requested alignment in bytes, not a log2/exponent.
    fn alloc_pages_at(&mut self, base: usize, num_pages: usize, align: usize)
    -> AllocResult<usize>;

    /// Returns the total number of memory pages.
    fn total_pages(&self) -> usize;

    /// Returns the number of allocated memory pages.
    fn used_pages(&self) -> usize;

    /// Returns the number of available memory pages.
    fn available_pages(&self) -> usize;
}

/// Used to allocate unique IDs (e.g., thread ID).
pub trait IdAllocator: BaseAllocator {
    /// Allocate contiguous IDs with given count and alignment.
    fn alloc_id(&mut self, count: usize, align_pow2: usize) -> AllocResult<usize>;

    /// Deallocate contiguous IDs with given position and count.
    fn dealloc_id(&mut self, start_id: usize, count: usize);

    /// Whether the given `id` was allocated.
    fn is_allocated(&self, id: usize) -> bool;

    /// Mark the given `id` has been allocated and cannot be reallocated.
    fn alloc_fixed_id(&mut self, id: usize) -> AllocResult;

    /// Returns the maximum number of supported IDs.
    fn size(&self) -> usize;

    /// Returns the number of allocated IDs.
    fn used(&self) -> usize;

    /// Returns the number of available IDs.
    fn available(&self) -> usize;
}

#[inline]
#[allow(dead_code)]
const fn align_down(pos: usize, align: usize) -> usize {
    pos & !(align - 1)
}

#[inline]
#[allow(dead_code)]
const fn align_up(pos: usize, align: usize) -> usize {
    (pos + align - 1) & !(align - 1)
}

/// Checks whether the address has the demanded alignment.
///
/// Equivalent to `addr % align == 0`, but the alignment must be a power of two.
#[inline]
#[allow(dead_code)]
const fn is_aligned(base_addr: usize, align: usize) -> bool {
    base_addr & (align - 1) == 0
}

#[cfg(feature = "allocator_api")]
mod allocator_api {
    extern crate alloc;

    use alloc::rc::Rc;
    use core::{
        alloc::{AllocError, Allocator, Layout},
        cell::RefCell,
        ptr::NonNull,
    };

    use super::ByteAllocator;

    /// A byte-allocator wrapped in [`Rc<RefCell>`] that implements [`core::alloc::Allocator`].
    pub struct AllocatorRc<A: ByteAllocator>(Rc<RefCell<A>>);

    impl<A: ByteAllocator> AllocatorRc<A> {
        /// Creates a new allocator with the given memory pool.
        pub fn new(mut inner: A, pool: &mut [u8]) -> Self {
            inner.init(pool.as_mut_ptr() as usize, pool.len());
            Self(Rc::new(RefCell::new(inner)))
        }
    }

    unsafe impl<A: ByteAllocator> Allocator for AllocatorRc<A> {
        fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
            match layout.size() {
                0 => Ok(NonNull::slice_from_raw_parts(NonNull::dangling(), 0)),
                size => {
                    let raw_addr = self.0.borrow_mut().alloc(layout).map_err(|_| AllocError)?;
                    Ok(NonNull::slice_from_raw_parts(raw_addr, size))
                }
            }
        }

        unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
            self.0.borrow_mut().dealloc(ptr, layout)
        }
    }

    impl<A: ByteAllocator> Clone for AllocatorRc<A> {
        fn clone(&self) -> Self {
            Self(self.0.clone())
        }
    }
}

#[cfg(feature = "allocator_api")]
pub use allocator_api::AllocatorRc;