stack-arena 0.12.0

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

/// A low-level memory chunk that represents a contiguous region of memory.
///
/// `Chunk` is the most basic building block of the arena allocator system,
/// representing a single contiguous region of memory with a base pointer and
/// a limit pointer. It handles the actual memory allocation and deallocation
/// at the lowest level.
///
/// # Memory Management
///
/// - Allocates a contiguous block of memory with the specified capacity
/// - Rounds capacity up to the next power of two for efficient memory use
/// - Provides methods to check if a pointer is contained within the chunk
/// - Automatically deallocates memory when dropped
///
/// # Implementation Details
///
/// - Memory is aligned to 4096 bytes (page size) for optimal performance
/// - Capacity is always a power of two
/// - Uses `std::alloc` functions for low-level memory allocation
#[derive(Debug)]
pub(crate) struct Chunk {
    /// Base pointer to the start of the memory chunk
    base: NonNull<u8>,
    /// Limit pointer to the end of the memory chunk (one past the last byte)
    limit: NonNull<u8>,
}

impl Drop for Chunk {
    #[inline]
    fn drop(&mut self) {
        let layout = unsafe { Layout::from_size_align_unchecked(self.capacity(), 4096) };
        unsafe { std::alloc::dealloc(self.base.as_ptr(), layout) };
    }
}

impl Default for Chunk {
    #[inline]
    fn default() -> Self {
        let layout = unsafe { Layout::from_size_align_unchecked(4096, 4096) };
        let Some(ptr) = NonNull::new(unsafe { std::alloc::alloc(layout) }) else {
            std::alloc::handle_alloc_error(layout);
        };
        Self {
            base: ptr,
            limit: unsafe { ptr.add(layout.size()) },
        }
    }
}

impl Chunk {
    /// Creates a new memory chunk with the specified capacity.
    ///
    /// The capacity is adjusted according to these rules:
    /// - If the capacity is less than 4096 bytes, it will be set to 4096 bytes
    /// - If the capacity is not a power of two, it will be rounded up to the next power of two
    ///
    /// # Parameters
    ///
    /// * `capacity` - The desired capacity in bytes
    ///
    /// # Returns
    ///
    /// A new `Chunk` with the specified (or adjusted) capacity
    #[inline]
    pub(crate) fn with_capacity(mut capacity: usize) -> Self {
        if capacity < 4096 {
            capacity = 4096;
        } else if !capacity.is_power_of_two() {
            capacity = capacity.next_power_of_two();
        }
        let layout = unsafe { Layout::from_size_align_unchecked(capacity, 4096) };
        let Some(base) = NonNull::new(unsafe { std::alloc::alloc(layout) }) else {
            std::alloc::handle_alloc_error(layout);
        };
        Self {
            base,
            limit: unsafe { base.add(layout.size()) },
        }
    }

    #[inline]
    pub(crate) fn base(&self) -> NonNull<u8> {
        self.base
    }

    #[inline]
    pub(crate) fn limit(&self) -> NonNull<u8> {
        self.limit
    }

    /// Returns the total capacity of the chunk in bytes.
    ///
    /// This is the total amount of memory available in the chunk.
    #[inline]
    pub(crate) fn capacity(&self) -> usize {
        unsafe { self.limit.byte_offset_from_unsigned(self.base) }
    }

    /// Checks if the given pointer is contained within this chunk.
    ///
    /// # Parameters
    ///
    /// * `ptr` - The pointer to check
    ///
    /// # Returns
    ///
    /// `true` if the pointer is within the chunk's memory range, `false` otherwise.
    #[inline]
    pub(crate) fn contains(&self, ptr: NonNull<u8>) -> bool {
        ptr >= self.base && ptr < self.limit
    }
}

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

    #[test]
    fn test_alignment() {
        let alignments = [1, 2, 4, 8, 16, 32, 4096];
        let chunk = Chunk::with_capacity(8192);
        let mut offset = 0usize;

        for &align in &alignments {
            // Calculate the next aligned address
            let base_addr = chunk.base.as_ptr() as usize;
            let aligned_addr = (base_addr + offset + (align - 1)) & !(align - 1);
            let rel_offset = aligned_addr - base_addr;
            assert!(rel_offset + 8 <= chunk.capacity());

            // Write data
            let ptr = unsafe { chunk.base.as_ptr().add(rel_offset) };
            unsafe { std::ptr::write_bytes(ptr, align as u8, 8) };

            // Verify data
            let data = unsafe { std::slice::from_raw_parts(ptr, 8) };
            assert_eq!(data, vec![align as u8; 8].as_slice());

            // Move offset forward
            offset = rel_offset + 8;
        }
    }

    #[test]
    fn test_contains() {
        let chunk = Chunk::with_capacity(4096);
        let base = chunk.base;
        let limit = chunk.limit;
        assert!(chunk.contains(base));
        assert!(!chunk.contains(limit));
        let mid = unsafe { base.as_ptr().add(64) };
        assert!(chunk.contains(NonNull::new(mid).unwrap()));
    }

    #[test]
    fn test_capacity() {
        for (capacity, expected) in [
            (0, 4096),
            (1, 4096),
            (32, 4096),
            (128, 4096),
            (1024, 4096),
            (4095, 4096),
            (4097, 8192),
            (8000, 8192),
            (64 * 1024, 64 * 1024),
        ] {
            let chunk = Chunk::with_capacity(capacity);
            assert_eq!(chunk.capacity(), expected);
        }
    }
}