stack_arena/

object_stack.rs

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

use crate::{Allocator, StackArena};

/// A high-level stack-based object builder that uses `StackArena` for memory management.
///
/// `ObjectStack` provides a more user-friendly interface on top of [`StackArena`](crate::StackArena) for
/// building and managing objects in a stack-like fashion. It supports incremental
/// object construction through the `extend` and `finish` methods, as well as
/// direct formatting through the `std::fmt::Write` trait.
///
/// # Features
///
/// - Push complete objects onto the stack
/// - Build objects incrementally using `extend` and `finish` methods
/// - Implements `std::fmt::Write` for string formatting directly into objects
/// - Stack-like (LIFO) allocation and deallocation pattern
/// - Memory-efficient with minimal overhead
/// - Automatic memory management with chunk reuse
///
/// # Use Cases
///
/// `ObjectStack` is particularly useful for:
///
/// - String building and text processing
/// - Constructing complex objects incrementally
/// - Serialization operations
/// - Any scenario where objects need to be built in stages
///
/// # Examples
///
/// ```
/// use stack_arena::ObjectStack;
/// use std::fmt::Write;
///
/// let mut stack = ObjectStack::new();
///
/// // Push a complete object
/// stack.push(b"hello");
///
/// // Build an object incrementally
/// stack.extend("world ");
/// write!(&mut stack, "from {}", "Rust").unwrap();
/// let ptr = stack.finish();
///
/// // Access the object
/// let message = unsafe { std::str::from_utf8_unchecked(ptr.as_ref()) };
/// assert_eq!(message, "world from Rust");
///
/// // Pop the object when done
/// stack.pop();
/// ```
#[derive(Debug)]
pub struct ObjectStack {
    arena: StackArena,
    partial: bool,
}

impl ObjectStack {
    /// Creates a new empty `ObjectStack` with a default initial capacity.
    ///
    /// The initial chunk size is 1024 bytes, managed by the underlying `StackArena`.
    ///
    /// # Examples
    ///
    /// ```
    /// use stack_arena::ObjectStack;
    ///
    /// let stack = ObjectStack::new();
    /// assert!(stack.is_empty());
    /// ```
    #[inline]
    pub fn new() -> Self {
        Self {
            arena: StackArena::new(),
            partial: false,
        }
    }

    /// Returns the number of objects currently on the stack.
    ///
    /// # Examples
    ///
    /// ```
    /// use stack_arena::ObjectStack;
    ///
    /// let mut stack = ObjectStack::new();
    /// assert_eq!(stack.len(), 0);
    ///
    /// stack.push(b"hello");
    /// assert_eq!(stack.len(), 1);
    /// ```
    #[inline]
    pub fn len(&self) -> usize {
        self.arena.len()
    }

    /// Returns `true` if the stack contains no objects.
    ///
    /// # Examples
    ///
    /// ```
    /// use stack_arena::ObjectStack;
    ///
    /// let mut stack = ObjectStack::new();
    /// assert!(stack.is_empty());
    ///
    /// stack.push(b"hello");
    /// assert!(!stack.is_empty());
    /// ```
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.arena.is_empty()
    }

    /// Pushes a complete object onto the stack.
    ///
    /// This method allocates memory for the object, copies the data,
    /// and returns a pointer to the allocated memory.
    ///
    /// # Parameters
    ///
    /// * `object` - The data to push onto the stack, which can be any type
    ///   that can be converted to a byte slice.
    ///
    /// # Returns
    ///
    /// A non-null pointer to the allocated memory containing the data.
    ///
    /// # Examples
    ///
    /// ```
    /// use stack_arena::ObjectStack;
    ///
    /// let mut stack = ObjectStack::new();
    /// let ptr = stack.push(b"hello world");
    /// // The pointer is valid until the object is popped or freed
    /// ```
    #[inline]
    pub fn push<P: AsRef<[u8]>>(&mut self, object: P) -> NonNull<[u8]> {
        let data = object.as_ref();
        let layout = Layout::for_value(data);
        let ptr = unsafe { self.arena.allocate(layout).unwrap() };
        unsafe {
            ptr.cast().copy_from_nonoverlapping(
                NonNull::new_unchecked(data.as_ptr().cast_mut()),
                data.len(),
            )
        };
        ptr
    }

    /// 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::ObjectStack;
    ///
    /// let mut stack = ObjectStack::new();
    /// stack.push(b"hello");
    /// stack.push(b"world");
    /// assert_eq!(stack.len(), 2);
    ///
    /// stack.pop();
    /// assert_eq!(stack.len(), 1);
    /// ```
    #[inline]
    pub fn pop(&mut self) {
        self.arena.pop();
        self.partial = false;
    }

    /// Extends the current object being built with additional data.
    ///
    /// This method is used for incrementally building objects. Multiple calls to
    /// `extend` can be made before finalizing the object with `finish`. This method
    /// **only supports extending the last allocation** (following LIFO pattern),
    /// as it uses the underlying arena's grow functionality.
    ///
    /// # Parameters
    ///
    /// * `value` - The data to append to the current object, which can be any type
    ///   that can be converted to a byte slice.
    ///
    /// # Examples
    ///
    /// ```
    /// use stack_arena::ObjectStack;
    ///
    /// let mut stack = ObjectStack::new();
    /// stack.extend("Hello, ");
    /// stack.extend("world!");
    /// let ptr = stack.finish();
    /// // ptr now points to "Hello, world!"
    /// ```
    #[inline]
    pub fn extend<P: AsRef<[u8]>>(&mut self, value: P) {
        let data = value.as_ref();
        if self.partial {
            let partial_object = self.arena.top().unwrap();
            let old_layout = Layout::for_value(unsafe { partial_object.as_ref() });
            let new_layout = unsafe {
                Layout::from_size_align_unchecked(
                    old_layout.size() + data.len(),
                    old_layout.align(),
                )
            };
            let store = unsafe {
                self.arena
                    .grow(partial_object.cast(), old_layout, new_layout)
            }
            .unwrap();
            unsafe {
                store
                    .cast::<u8>()
                    .add(old_layout.size())
                    .copy_from_nonoverlapping(
                        NonNull::new_unchecked(data.as_ptr().cast_mut()),
                        data.len(),
                    );
            }
        } else {
            let store = unsafe { self.arena.allocate(Layout::for_value(data)) }.unwrap();
            unsafe {
                store.cast::<u8>().copy_from_nonoverlapping(
                    NonNull::new_unchecked(data.as_ptr().cast_mut()),
                    data.len(),
                )
            };
        }
        self.partial = true;
    }

    /// Finalizes the current object being built and adds it to the stack.
    ///
    /// This method should be called after one or more calls to `extend` to
    /// finalize the object and make it available on the stack.
    ///
    /// # Returns
    ///
    /// A non-null pointer to the finalized object.
    ///
    /// # Examples
    ///
    /// ```
    /// use stack_arena::ObjectStack;
    ///
    /// let mut stack = ObjectStack::new();
    /// stack.extend("Hello");
    /// stack.extend(" world");
    /// let ptr = stack.finish();
    /// // ptr now points to "Hello world"
    /// ```
    #[inline]
    pub fn finish(&mut self) -> NonNull<[u8]> {
        debug_assert!(self.partial);
        self.partial = false;
        self.arena.top().unwrap()
    }

    /// 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::ObjectStack;
    ///
    /// let mut stack = ObjectStack::new();
    /// stack.push(b"first");
    /// let second = stack.push(b"second");
    /// stack.push(b"third");
    ///
    /// // Free "second" and "third"
    /// stack.rollback(unsafe { second.as_ref() });
    /// assert_eq!(stack.len(), 1); // Only "first" remains
    /// ```
    #[inline]
    pub fn rollback(&mut self, data: &[u8]) {
        let data = data.as_ref();
        unsafe {
            self.arena.deallocate(
                NonNull::new_unchecked(data.as_ptr().cast_mut()),
                Layout::for_value(data),
            )
        };
    }
}

/// Implementation of the `std::fmt::Write` trait for `ObjectStack`.
///
/// This allows using the `write!` macro and other formatting utilities
/// to write formatted text directly into the object being built.
///
/// # Examples
///
/// ```
/// use std::fmt::Write;
/// use stack_arena::ObjectStack;
///
/// let mut stack = ObjectStack::new();
/// write!(&mut stack, "Hello, {}!", "world").unwrap();
/// let formatted = stack.finish();
/// // formatted now points to "Hello, world!"
/// ```
impl std::fmt::Write for ObjectStack {
    /// Writes a string into the arena.
    ///
    /// This method extends the current object with the given string.
    #[inline]
    fn write_str(&mut self, s: &str) -> std::fmt::Result {
        self.extend(s);
        Ok(())
    }
}

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

    #[test]
    fn test_lifecycle() {
        let mut stack = ObjectStack::new();
        write!(&mut stack, "ab").expect("write");
        let s = "c";
        stack.extend(s);
        let p = unsafe { stack.finish().as_ref() };
        assert_eq!(p, b"abc");
    }

    #[test]
    fn test_push_pop() {
        let mut stack = ObjectStack::new();

        // Test push
        stack.push(b"hello");
        assert_eq!(stack.len(), 1);
        assert!(!stack.is_empty());

        // Test push multiple items
        stack.push(b"world");
        assert_eq!(stack.len(), 2);

        // Test pop
        stack.pop();
        assert_eq!(stack.len(), 1);

        // Test pop to empty
        stack.pop();
        assert_eq!(stack.len(), 0);
        assert!(stack.is_empty());
    }

    #[test]
    fn test_extend_small_data() {
        let mut stack = ObjectStack::new();

        // Extend with data smaller than chunk capacity
        stack.extend(b"hello");

        // Extend again with small data
        stack.extend(b" world");

        // Finish and verify
        let data = unsafe { stack.finish().as_ref() };
        assert_eq!(data, b"hello world");
    }

    #[test]
    fn test_extend_large_data() {
        let mut stack = ObjectStack::new();

        // Extend with data larger than chunk capacity
        let large_data = vec![b'x'; 20];
        stack.extend(&large_data);

        // Finish and verify
        let data = unsafe { stack.finish().as_ref() };
        assert_eq!(data, &large_data[..]);
    }

    #[test]
    fn test_extend_after_finish() {
        let mut stack = ObjectStack::new();

        // First object
        stack.extend("first");
        let first = unsafe { stack.finish().as_ref() };
        assert_eq!(first, b"first");
        // Second object
        stack.extend(b"second");
        let second = unsafe { stack.finish().as_ref() };

        assert_eq!(second, b"second");

        // Verify both objects are still valid
        assert_eq!(first, b"first");
        assert_eq!(second, b"second");
    }

    #[test]
    fn test_free() {
        let mut stack = ObjectStack::new();

        // Create multiple objects
        stack.push(b"first");
        stack.extend(b"second");
        let second = unsafe { stack.finish().as_ref() };
        stack.extend(b"third");
        let _third = unsafe { stack.finish().as_ref() };

        // Free up to second object
        stack.rollback(second);

        // Verify stack state
        assert_eq!(stack.len(), 1); // Only "first" remains

        // Add a new object
        stack.extend(b"fourth");
        let fourth = unsafe { stack.finish().as_ref() };
        assert_eq!(fourth, b"fourth");

        // Verify stack state
        assert_eq!(stack.len(), 2); // "first" and "fourth"
    }

    #[test]
    fn test_write_trait() {
        let mut stack = ObjectStack::new();

        // Test write_str via the Write trait
        write!(&mut stack, "Hello, {}!", "world").unwrap();
        // Finish and verify
        let data = unsafe { stack.finish().as_ref() };
        assert_eq!(data, b"Hello, world!");
    }

    #[test]
    fn test_empty_data() {
        let mut stack = ObjectStack::new();

        // Test with empty data
        stack.extend(b"");
        let data = unsafe { stack.finish().as_ref() };
        assert_eq!(data, b"");

        // Test push with empty data
        stack.push(b"");
        assert_eq!(stack.len(), 2);
    }

    #[test]
    fn test_multiple_operations() {
        let mut stack = ObjectStack::new();

        // Mix of operations
        stack.push(b"item1");
        stack.extend(b"item2-part1");
        stack.extend(b"-part2");
        let item2 = unsafe { stack.finish().as_ref() };
        write!(&mut stack, "item3").unwrap();
        let item3 = unsafe { stack.finish().as_ref() };

        // Verify
        assert_eq!(item2, b"item2-part1-part2");
        assert_eq!(item3, b"item3");
        assert_eq!(stack.len(), 3);

        // Pop and verify
        stack.pop();
        assert_eq!(stack.len(), 2);
    }

    #[test]
    fn test_extend_exact_capacity() {
        let mut stack = ObjectStack::new();

        // Fill exactly to capacity
        let data = vec![b'x'; 10]; // Same as chunk_size
        stack.extend(&data);

        // Add more data to trigger new allocation
        stack.extend(b"more");

        // Finish and verify
        let result = unsafe { stack.finish().as_ref() };
        let mut expected = data.clone();
        expected.extend_from_slice(b"more");
        assert_eq!(result, expected.as_slice());
    }

    #[test]
    fn test_free_all() {
        let mut stack = ObjectStack::new();

        // Create multiple objects
        let first = stack.push(b"first");
        stack.extend(b"second");
        let _second = unsafe { stack.finish().as_ref() };

        // Free all objects by freeing the first one
        stack.rollback(unsafe { first.as_ref() });

        // Verify stack is empty
        assert_eq!(stack.len(), 0);
        assert!(stack.is_empty());
    }

    #[test]
    #[should_panic]
    fn test_free_nonexistent() {
        let mut stack = ObjectStack::new();

        // Create an object
        stack.push(b"object");
        assert_eq!(stack.len(), 1);

        // Try to free a non-existent object
        let dummy = b"nonexistent";
        stack.rollback(dummy);

        // Stack should remain unchanged
        assert_eq!(stack.len(), 1);
    }

    #[test]
    fn test_cross_chunk_allocation_deallocation() {
        // Create an ObjectStack with a custom StackArena that has a very small chunk size
        let mut stack = ObjectStack {
            arena: StackArena::with_chunk_size(8),
            partial: false,
        };

        let small1 = stack.push("small1");
        stack.push("small2");
        assert_eq!(stack.len(), 2);

        // Pop the second object to maintain LIFO order
        stack.pop();
        assert_eq!(stack.len(), 1);

        let large = "start- middle- this-is-a-longer-string-to-trigger-new-chunk";
        for part in large.split(' ') {
            stack.extend(part);
        }
        assert_eq!(stack.len(), 2);

        // Finish the object
        let large = stack.finish();
        assert_eq!(stack.len(), 2);
        unsafe {
            assert_eq!(large.as_ref(), large.as_ref());
        }

        // Verify data integrity
        assert_eq!(unsafe { small1.as_ref() }, b"small1");

        // Pop in LIFO order
        stack.pop(); // Pop large
        assert_eq!(stack.len(), 1);

        stack.pop(); // Pop small1
        assert!(stack.is_empty());

        // Build a new object with a single extend
        stack.extend("single-extend-object");

        // Finish the object
        let single = stack.finish();

        // Verify the object
        assert_eq!(unsafe { single.as_ref() }, b"single-extend-object");

        // Clean up
        stack.pop();
        assert_eq!(stack.len(), 0);
    }
}