stack_arena/

stack_arena.rs

//! A fast, stack-like arena allocator for efficient memory management.
//!
//! # Overview
//!
//! `StackArena` provides an efficient memory allocation strategy for scenarios where:
//! - Multiple objects need to be allocated in sequence
//! - Objects follow a stack-like (LIFO) allocation/deallocation pattern
//! - Memory usage needs to be minimized with low per-allocation overhead
//! - Performance is critical, such as in parsers, interpreters, or compilers
//!
//! # Features
//!
//! - Efficient allocation of variable-sized objects with minimal overhead
//! - Stack-like (LIFO) allocation and deallocation pattern
//! - Memory is managed in chunks, automatically growing when needed
//! - Implements the [`Allocator`](crate::Allocator) trait for compatibility with allocation APIs
//! - Low-level memory management with minimal overhead
//! - Automatic chunk reuse for improved performance
//! - Significantly faster than the system allocator for small, frequent allocations
//!
//! # Usage Examples
//!
//! Basic usage:
//!
//! ```
//! use stack_arena::{StackArena, Allocator};
//!
//! let mut arena = StackArena::new();
//!
//! // Allocate memory for an object
//! let layout = std::alloc::Layout::from_size_align(10, 1).unwrap();
//! let ptr = unsafe { arena.allocate(layout).unwrap() };
//!
//! // Use the allocated memory
//! unsafe { std::ptr::write_bytes(ptr.as_ptr() as *mut u8, 0xAA, 10) };
//!
//! // Deallocate when done
//! unsafe { arena.deallocate(ptr.cast(), layout) };
//! ```
//!
//! Creating with a custom chunk size:
//!
//! ```
//! use stack_arena::StackArena;
//!
//! // Create an arena with 8KB chunks
//! let mut arena = StackArena::with_chunk_size(8192);
//! assert!(arena.is_empty());
//! ```
//!
//! For higher-level operations, use [`ObjectStack`](crate::ObjectStack) instead:
//!
//! ```
//! use stack_arena::ObjectStack;
//!
//! let mut stack = ObjectStack::new();
//! stack.extend("Hello, ");
//! stack.extend("world!");
//! let ptr = stack.finish(); // Finalizes the object
//! ```
//!
//! # Memory Management
//!
//! `StackArena` manages memory in chunks:
//!
//! - When the current chunk is full, a new larger chunk is allocated
//! - Old chunks are stored for later reuse when objects are deallocated
//! - Memory is allocated contiguously within chunks for better cache locality
//! - The LIFO pattern allows for efficient memory reuse
//!
//! # Safety
//!
//! - All returned pointers are valid until the corresponding object is deallocated
//! - Objects are stored in contiguous memory chunks
//! - Unsafe operations are used internally for performance
//! - The library follows LIFO (Last-In-First-Out) allocation pattern for efficiency
//! - Users must ensure proper memory safety when working with raw pointers

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

use crate::{Allocator, BufferArena, chunk::Chunk};

#[derive(Debug)]
/// A stack-like arena allocator that efficiently manages memory in chunks.
///
/// `StackArena` provides a low-level memory allocation strategy where objects are allocated
/// in a stack-like fashion (Last-In-First-Out). It manages memory in chunks, automatically
/// allocating new chunks when needed, which reduces allocation overhead and improves performance.
///
/// # Memory Management
///
/// - Memory is organized in chunks, with a current active chunk for allocations
/// - When the current chunk is full, a new larger chunk is allocated
/// - Old chunks are stored to keep previous allocations valid until explicitly deallocated
/// - Objects are allocated contiguously within chunks for better cache locality
/// - Follows LIFO (Last-In-First-Out) allocation/deallocation pattern
///
/// # Use Cases
///
/// `StackArena` is particularly well-suited for:
///
/// - Parsers and compilers that need to allocate many temporary objects
/// - Interpreters that need efficient memory management for short-lived objects
/// - Data processing pipelines with predictable allocation patterns
/// - Any scenario where many small allocations are made in sequence
///
/// # Thread Safety
///
/// `StackArena` is not thread-safe and should not be shared between threads
/// without external synchronization.
///
/// # Performance
///
/// This allocator is optimized for scenarios with many small allocations that follow
/// a stack-like pattern. Benchmarks show it significantly outperforms the system allocator:
///
/// - Up to 10x faster for consecutive small allocations
/// - Minimal overhead for allocation and deallocation
/// - Efficient memory reuse with the LIFO pattern
pub struct StackArena {
    /// Stores old chunks to keep previous allocations valid until explicitly deallocated.
    /// This ensures that pointers returned by `allocate` remain valid even after new chunks are allocated.
    store: Vec<Chunk>,

    /// Tracks all allocated objects in LIFO order for proper deallocation.
    stack: Vec<NonNull<[u8]>>,

    /// The current active chunk used for new allocations.
    current: BufferArena,

    /// Default chunk size used when allocating new chunks.
    /// When a new chunk is needed, its size will be at least this value.
    default_chunk_size: usize,
}

impl Default for StackArena {
    fn default() -> Self {
        let default_chunk_size = 1 << 12;
        Self {
            store: Vec::with_capacity(16),
            stack: Vec::with_capacity(256),
            current: Default::default(),
            default_chunk_size,
        }
    }
}

impl StackArena {
    /// Creates a new empty `StackArena` with a default initial capacity.
    ///
    /// The initial chunk size is 1024 bytes.
    ///
    /// # Examples
    ///
    /// ```
    /// use stack_arena::StackArena;
    ///
    /// let arena = StackArena::new();
    /// assert!(arena.is_empty());
    /// ```
    #[inline(always)]
    pub fn new() -> Self {
        Self::with_chunk_size(4096)
    }

    /// Creates a new empty `StackArena` with a specified chunk size.
    ///
    /// # Parameters
    ///
    /// * `chunk_size` - The size in bytes to use for new chunks
    ///
    /// # Examples
    ///
    /// ```
    /// use stack_arena::StackArena;
    ///
    /// // Create an arena with 4096-byte chunks
    /// let arena = StackArena::with_chunk_size(4096);
    /// assert!(arena.is_empty());
    /// ```
    #[inline]
    pub fn with_chunk_size(chunk_size: usize) -> Self {
        let current = BufferArena::with_capacity(chunk_size);
        Self {
            store: Vec::with_capacity(4),
            stack: Vec::with_capacity(32),
            current,
            default_chunk_size: chunk_size,
        }
    }

    /// Returns the number of objects currently on the stack.
    ///
    /// # Examples
    ///
    /// ```
    /// use stack_arena::{StackArena, Allocator};
    /// use std::alloc::Layout;
    ///
    /// let mut arena = StackArena::new();
    /// assert_eq!(arena.len(), 0);
    ///
    /// // Allocate memory for an object
    /// let layout = Layout::from_size_align(5, 1).unwrap();
    /// let ptr = unsafe { arena.allocate(layout).unwrap() };
    /// assert_eq!(arena.len(), 1);
    /// ```
    #[inline]
    pub fn len(&self) -> usize {
        self.stack.len()
    }

    /// Returns `true` if the arena contains no objects.
    ///
    /// # Examples
    ///
    /// ```
    /// use stack_arena::{StackArena, Allocator};
    /// use std::alloc::Layout;
    ///
    /// let mut arena = StackArena::new();
    /// assert!(arena.is_empty());
    ///
    /// // Allocate memory for an object
    /// let layout = Layout::from_size_align(5, 1).unwrap();
    /// let ptr = unsafe { arena.allocate(layout).unwrap() };
    /// assert!(!arena.is_empty());
    /// ```
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    /// Removes the most recently pushed object from the stack.
    ///
    /// This method follows the LIFO (Last-In-First-Out) principle.
    /// After popping, any pointers to the popped object become invalid.
    ///
    /// # Panics
    ///
    /// Panics if the stack is empty or if there is a partial object
    /// being built (i.e., if `extend` has been called but `finish` has not).
    ///
    /// # Examples
    ///
    /// ```
    /// use stack_arena::{StackArena, Allocator};
    /// use std::alloc::Layout;
    ///
    /// let mut arena = StackArena::new();
    ///
    /// // Allocate first object
    /// let layout1 = Layout::from_size_align(5, 1).unwrap();
    /// let ptr1 = unsafe { arena.allocate(layout1).unwrap() };
    ///
    /// // Allocate second object
    /// let layout2 = Layout::from_size_align(5, 1).unwrap();
    /// let ptr2 = unsafe { arena.allocate(layout2).unwrap() };
    /// assert_eq!(arena.len(), 2);
    ///
    /// arena.pop();
    /// assert_eq!(arena.len(), 1);
    /// ```
    #[inline]
    pub fn pop(&mut self) -> Option<NonNull<[u8]>> {
        let ptr = *self.stack.last().unwrap();
        let layout = Layout::for_value(unsafe { ptr.as_ref() });
        unsafe { self.deallocate(ptr.cast(), layout) };
        Some(ptr)
    }

    /// Returns the top object on the stack without removing it.
    ///
    /// This method returns a reference to the most recently allocated object
    /// on the stack.
    ///
    /// # Returns
    ///
    /// An optional non-null pointer to the top object, or None if the stack is empty.
    ///
    /// # Examples
    ///
    /// ```
    /// use stack_arena::{StackArena, Allocator};
    /// use std::alloc::Layout;
    ///
    /// let mut arena = StackArena::new();
    /// let layout = Layout::from_size_align(8, 1).unwrap();
    /// let ptr = unsafe { arena.allocate(layout).unwrap() };
    ///
    /// // Get the top object
    /// let top = arena.top();
    /// assert!(top.is_some());
    /// ```
    #[inline]
    pub fn top(&mut self) -> Option<NonNull<[u8]>> {
        self.stack.last().map(|&ptr| ptr)
    }

    /// Rolls back to a specific object, freeing it and all objects allocated after it.
    ///
    /// This method allows for rolling back to a specific point in the allocation
    /// history by providing a reference to an object on the stack.
    ///
    /// # Parameters
    ///
    /// * `data` - A reference to the object to free, along with all objects
    ///   allocated after it.
    ///
    /// # Examples
    ///
    /// ```
    /// use stack_arena::{StackArena, Allocator};
    /// use std::alloc::Layout;
    ///
    /// let mut arena = StackArena::new();
    ///
    /// // Allocate some objects
    /// let layout1 = Layout::from_size_align(8, 1).unwrap();
    /// let ptr1 = unsafe { arena.allocate(layout1).unwrap() };
    ///
    /// let layout2 = Layout::from_size_align(16, 1).unwrap();
    /// let ptr2 = unsafe { arena.allocate(layout2).unwrap() };
    ///
    /// // Roll back to the first object
    /// unsafe {
    ///     arena.rollback(ptr1.as_ref());
    /// }
    /// assert_eq!(arena.len(), 0); // All objects are freed
    /// ```
    #[inline]
    pub fn rollback(&mut self, data: &[u8]) {
        let data = data.as_ref();
        unsafe {
            self.deallocate(
                NonNull::new_unchecked(data.as_ptr().cast_mut()),
                Layout::for_value(data),
            )
        };
    }
}

/// Implementation of the `Allocator` trait for `StackArena`.
///
/// This implementation allows `StackArena` to be used with APIs that require
/// an allocator. The implementation follows the stack-like allocation pattern,
/// where memory is allocated in chunks and objects are allocated contiguously.
///
/// # Safety
///
/// The implementation uses unsafe code internally to manage memory efficiently.
/// Users should follow the LIFO (Last-In-First-Out) pattern when deallocating
/// memory to ensure proper behavior.
///
/// # Performance
///
/// This allocator is optimized for scenarios with many small allocations that follow
/// a stack-like pattern. It significantly outperforms the system allocator in these cases
/// as demonstrated in the benchmarks.
impl Allocator for StackArena {
    /// Allocates memory with the specified layout.
    ///
    /// This method allocates a new object with the given layout in the current chunk.
    /// If the current chunk doesn't have enough space, a new chunk is allocated.
    ///
    /// # Safety
    ///
    /// This method is unsafe because it returns a raw pointer that must be used
    /// carefully to avoid memory safety issues.
    unsafe fn allocate(
        &mut self,
        layout: std::alloc::Layout,
    ) -> Result<std::ptr::NonNull<[u8]>, crate::AllocError> {
        if !self.current.sufficient_for(layout) {
            let capacity = layout.size().max(self.default_chunk_size);
            let prev = std::mem::replace(&mut self.current, BufferArena::with_capacity(capacity));
            self.store.push(prev.into());
        }

        unsafe {
            let object = self.current.allocate(layout)?;
            self.stack.push(object);
            Ok(object)
        }
    }

    /// Deallocates memory previously allocated by `allocate`.
    ///
    /// This method frees the specified object and all objects allocated after it,
    /// following the stack-like (LIFO) deallocation pattern.
    ///
    /// # 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.
    #[inline]
    unsafe fn deallocate(&mut self, ptr: std::ptr::NonNull<u8>, layout: std::alloc::Layout) {
        let object = NonNull::slice_from_raw_parts(ptr, layout.size());
        // match self.stack.binary_search(&object) {
        //     Ok(pos) => todo!(),
        //     Err(_) => todo!(),
        // };
        while let Some(top) = self.stack.pop() {
            if std::ptr::eq(object.as_ptr(), top.as_ptr()) {
                break;
            }
        }
        if self.current.contains(ptr) {
            unsafe { self.current.deallocate(ptr, layout) };
            return;
        }
        while let Some(_chunk) = self.store.pop_if(|chunk| !chunk.contains(ptr)) {
            // eprintln!("free chunk: {chunk:?}");
        }
        if let Some(chunk) = self.store.pop() {
            debug_assert!(chunk.contains(ptr));
            self.current = chunk.into();
        }
        self.current.ptr = ptr;
    }

    /// Grows a previously allocated memory block to a larger size.
    ///
    /// This method is used to extend an existing allocation to a larger size.
    /// It's used internally by the `extend` method. It **only supports growing
    /// the last allocation** (following LIFO pattern). Attempting to grow any
    /// other allocation will trigger a debug assertion failure.
    ///
    /// # 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.
    unsafe fn grow(
        &mut self,
        ptr: NonNull<u8>,
        old_layout: std::alloc::Layout,
        new_layout: std::alloc::Layout,
    ) -> Result<NonNull<[u8]>, crate::AllocError> {
        let top = self.stack.pop().unwrap();
        debug_assert_eq!(
            top.cast().as_ptr(),
            ptr.as_ptr(),
            "this operation is only supported for the last allocation"
        );
        let object = if self.current.remaining() < new_layout.size() - old_layout.size() {
            let capacity = new_layout.size().max(self.default_chunk_size);
            let prev = std::mem::replace(&mut self.current, BufferArena::with_capacity(capacity));
            self.store.push(prev.into());
            unsafe {
                let object = self.current.allocate(new_layout)?;
                object
                    .cast()
                    .copy_from_nonoverlapping(ptr, old_layout.size());
                object
            }
        } else {
            unsafe { self.current.grow(ptr, old_layout, new_layout) }?
        };
        self.stack.push(object);
        Ok(object)
    }

    /// Shrinks a previously allocated memory block to a smaller size.
    ///
    /// This method is used to reduce the size of an existing allocation.
    ///
    /// # 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.
    #[inline]
    unsafe fn shrink(
        &mut self,
        ptr: NonNull<u8>,
        old_layout: Layout,
        new_layout: Layout,
    ) -> Result<NonNull<[u8]>, crate::AllocError> {
        let top = self.stack.pop().unwrap();
        debug_assert_eq!(
            top.cast().as_ptr(),
            ptr.as_ptr(),
            "this operation is only supported for the last allocation"
        );
        debug_assert_eq!(top.cast().as_ptr(), ptr.as_ptr());
        debug_assert_eq!(self.current.ptr, unsafe { ptr.add(old_layout.size()) });
        let object = unsafe { self.current.shrink(ptr, old_layout, new_layout) }?;
        self.stack.push(object);
        Ok(object)
    }
}

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

    #[test]
    fn test_new() {
        let stack = StackArena::new();
        assert_eq!(stack.len(), 0);
        assert!(stack.is_empty());
    }

    #[test]
    fn test_allocate_deallocate() {
        let mut stack = StackArena::new();

        // Allocate some memory
        let layout = Layout::from_size_align(16, 8).unwrap();
        let ptr = unsafe { stack.allocate(layout) }.unwrap();

        // Write some data
        unsafe { std::ptr::write_bytes(ptr.as_ptr() as *mut u8, 0xAA, 16) };

        // Verify data
        let data = unsafe { std::slice::from_raw_parts(ptr.as_ptr() as *const u8, 16) };
        assert_eq!(data, [0xAA; 16].as_slice());

        // Deallocate
        unsafe { stack.deallocate(ptr.cast(), layout) };
        assert_eq!(stack.len(), 0);
    }

    #[test]
    fn test_multi_allocations() {
        let mut stack = StackArena::new();
        let n = 10;
        let mut allocations = Vec::with_capacity(n);

        // Allocate multiple objects
        for i in 0..n {
            let layout = Layout::array::<u8>(i + 1).unwrap();
            let ptr = unsafe { stack.allocate(layout) }.unwrap();

            // Write some data
            unsafe { std::ptr::write_bytes(ptr.as_ptr() as *mut u8, i as u8, i + 1) };

            allocations.push((ptr, layout));
        }

        assert_eq!(stack.len(), n);

        // Verify data
        for (i, (ptr, _)) in allocations.iter().enumerate() {
            let data = unsafe { std::slice::from_raw_parts(ptr.as_ptr() as *const u8, i + 1) };
            assert_eq!(data, vec![i as u8; i + 1].as_slice());
        }

        // Deallocate in LIFO order
        for (ptr, layout) in allocations.into_iter().rev() {
            unsafe { stack.deallocate(ptr.cast(), layout) };
        }

        assert_eq!(stack.len(), 0);
    }

    #[test]
    fn test_grow_shrink() {
        let mut stack = StackArena::new();

        // Allocate initial memory
        let initial_layout = Layout::from_size_align(8, 8).unwrap();
        let ptr = unsafe { stack.allocate(initial_layout) }.unwrap();

        // Write initial data
        unsafe { std::ptr::write_bytes(ptr.as_ptr() as *mut u8, 0xAA, 8) };

        // Grow the allocation
        let new_layout = Layout::from_size_align(16, 8).unwrap();
        let grown_ptr = unsafe { stack.grow(ptr.cast(), initial_layout, new_layout) }.unwrap();

        // Verify initial data is preserved
        let data = unsafe { std::slice::from_raw_parts(grown_ptr.as_ptr() as *const u8, 8) };
        assert_eq!(data, [0xAA; 8].as_slice());

        // Write to the new space
        unsafe { std::ptr::write_bytes((grown_ptr.as_ptr() as *mut u8).add(8), 0xBB, 8) };

        // Shrink back to original size
        let shrunk_ptr =
            unsafe { stack.shrink(grown_ptr.cast(), new_layout, initial_layout) }.unwrap();

        // Verify data
        let final_data = unsafe { std::slice::from_raw_parts(shrunk_ptr.as_ptr() as *const u8, 8) };
        assert_eq!(final_data, [0xAA; 8].as_slice());

        // Deallocate
        unsafe { stack.deallocate(shrunk_ptr.cast(), initial_layout) };
        assert_eq!(stack.len(), 0);
    }

    /// Helper function to allocate memory and write data to it
    #[inline]
    fn allocate_and_write<A: Allocator>(allocator: &mut A, size: usize) -> (NonNull<[u8]>, Layout) {
        let layout = Layout::from_size_align(size, 8).unwrap();
        let ptr = unsafe { allocator.allocate(layout).unwrap() };

        // Write some data
        unsafe { std::ptr::write_bytes(ptr.as_ptr() as *mut u8, 0xAA, size) };

        (ptr, layout)
    }

    /// Helper function to perform consecutive allocations
    #[inline]
    fn perform_consecutive_allocations<A: Allocator>(
        allocator: &mut A,
        count: usize,
    ) -> Vec<(NonNull<[u8]>, Layout)> {
        let mut ptrs = Vec::with_capacity(count);

        for i in 0..count {
            // Use different small sizes for each allocation
            let size = (i % 8) + 1; // Sizes from 1 to 8 bytes
            let (ptr, layout) = allocate_and_write(allocator, size);
            ptrs.push((ptr, layout));
        }

        ptrs
    }

    #[test]
    fn test_consecutive_allocations() {
        let mut stack = StackArena::new();
        let allocations = perform_consecutive_allocations(&mut stack, 10000);

        // assert_eq!(stack.len(), 10);

        // Deallocate in LIFO order
        for (ptr, layout) in allocations.into_iter().rev() {
            unsafe { stack.deallocate(ptr.cast(), layout) };
        }

        assert_eq!(stack.len(), 0);
    }

    #[test]
    fn test_custom_chunk_size() {
        // Create a stack arena with a custom chunk size
        let mut stack = StackArena::with_chunk_size(256);

        // Allocate an object that fits within the chunk
        let small_layout = Layout::from_size_align(64, 8).unwrap();
        let small_ptr = unsafe { stack.allocate(small_layout) }.unwrap();

        // Allocate an object that exceeds the default chunk size (1024) but is less than
        // our custom chunk size (256)
        let medium_layout = Layout::from_size_align(128, 8).unwrap();
        let medium_ptr = unsafe { stack.allocate(medium_layout) }.unwrap();

        // Allocate an object that exceeds our custom chunk size
        let large_layout = Layout::from_size_align(512, 8).unwrap();
        let large_ptr = unsafe { stack.allocate(large_layout) }.unwrap();

        // Verify all allocations succeeded
        assert_eq!(stack.len(), 3);

        // Clean up
        unsafe {
            stack.deallocate(large_ptr.cast(), large_layout);
            stack.deallocate(medium_ptr.cast(), medium_layout);
            stack.deallocate(small_ptr.cast(), small_layout);
        }

        assert_eq!(stack.len(), 0);
    }

    #[test]
    fn test_cross_chunk_allocation_deallocation() {
        // Create a stack arena with a very small chunk size to easily trigger new chunk allocations
        let mut stack = StackArena::with_chunk_size(32);

        // Allocate objects that will span across multiple chunks
        let mut allocations = Vec::new();

        let specs = [
            (16, 8, 0xAA),
            (8, 8, 0xBB),
            (16, 8, 0xCC),
            (8, 8, 0xDD),
            (39, 4, 0xEE),
        ];
        for (size, align, fill) in specs {
            let layout = unsafe { Layout::from_size_align_unchecked(size, align) };
            let ptr = unsafe { stack.allocate(layout).unwrap() };
            unsafe { ptr.cast::<u8>().write_bytes(fill, layout.size()) };
            allocations.push((ptr, layout));
        }

        // Verify all allocations
        assert_eq!(stack.len(), specs.len());

        for ((size, _align, fill), (ptr, _layout)) in std::iter::zip(specs, &allocations) {
            let data = unsafe { ptr.as_ref() };
            let expected = vec![fill; size];
            assert_eq!(data, &expected);
        }

        // Deallocate in LIFO order to test chunk reuse
        for (ptr, layout) in allocations.into_iter().rev() {
            unsafe { stack.deallocate(ptr.cast(), layout) };
        }

        // Verify all memory has been deallocated
        assert_eq!(stack.len(), 0);

        // Allocate again to test chunk reuse
        let layout5 = Layout::from_size_align(24, 8).unwrap();
        let ptr5 = unsafe { stack.allocate(layout5) }.unwrap();
        unsafe { ptr5.cast::<u8>().write_bytes(0xFF, layout5.size()) };

        // Verify the new allocation
        assert_eq!(stack.len(), 1);
        let data5 =
            unsafe { std::slice::from_raw_parts(ptr5.as_ptr() as *const u8, layout5.size()) };
        assert_eq!(data5, [0xFF; 24].as_slice());
    }

    #[test]
    fn test_stack_arena_grow_with_new_chunk() {
        // Create a stack arena with a very small chunk size
        let mut stack = StackArena::with_chunk_size(32);

        // First allocation - should fit in the first chunk
        let layout1 = Layout::from_size_align(8, 8).unwrap();
        let ptr1 = unsafe { stack.allocate(layout1) }.unwrap();
        unsafe { ptr1.cast::<u8>().write_bytes(0xAA, layout1.size()) };

        // Allocate a second object to fill the first chunk
        let layout2 = Layout::from_size_align(16, 8).unwrap();
        let ptr2 = unsafe { stack.allocate(layout2) }.unwrap();
        unsafe { ptr2.cast::<u8>().write_bytes(0xBB, layout2.size()) };

        // Verify data integrity
        let data1 =
            unsafe { std::slice::from_raw_parts(ptr1.as_ptr() as *const u8, layout1.size()) };
        assert_eq!(data1, [0xAA; 8].as_slice());

        let data2 =
            unsafe { std::slice::from_raw_parts(ptr2.as_ptr() as *const u8, layout2.size()) };
        assert_eq!(data2, [0xBB; 16].as_slice());

        // Deallocate the second object
        unsafe { stack.deallocate(ptr2.cast(), layout2) };

        // Now grow the first allocation to trigger a new chunk allocation
        // The new size won't fit in the first chunk, so it should be copied to a new chunk
        let layout1_grown = Layout::from_size_align(24, 8).unwrap();

        let ptr1_grown = unsafe { stack.grow(ptr1.cast(), layout1, layout1_grown) }.unwrap();

        // Write to the newly grown portion
        unsafe {
            ptr1_grown
                .cast::<u8>()
                .add(layout1.size())
                .write_bytes(0xCC, layout1_grown.size() - layout1.size())
        };

        // Verify the grown object has correct data
        let data1_grown = unsafe {
            std::slice::from_raw_parts(ptr1_grown.as_ptr() as *const u8, layout1_grown.size())
        };
        let mut expected1_grown = vec![0xAA; layout1.size()];
        expected1_grown.extend_from_slice(&vec![0xCC; layout1_grown.size() - layout1.size()]);
        assert_eq!(data1_grown, expected1_grown.as_slice());

        // Final cleanup
        unsafe { stack.deallocate(ptr1_grown.cast(), layout1_grown) };
        assert_eq!(stack.len(), 0);
    }
}