stack_arena/

stack_arena.rs

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

use crate::Allocator;

#[derive(Debug)]
struct Chunk {
    buffer: Box<[u8]>,
    offset: usize,
    len: usize,
}

impl Chunk {
    fn new(capacity: usize) -> Self {
        Self {
            buffer: unsafe { Box::new_uninit_slice(capacity).assume_init() },
            offset: 0,
            len: 0,
        }
    }

    fn capacity(&self) -> usize {
        self.buffer.len()
    }

    fn object(&self) -> NonNull<[u8]> {
        NonNull::slice_from_raw_parts(
            unsafe { NonNull::new_unchecked(self.buffer.as_ptr().add(self.offset).cast_mut()) },
            self.len,
        )
    }
}

impl From<Box<[u8]>> for Chunk {
    fn from(value: Box<[u8]>) -> Self {
        Self {
            buffer: value,
            offset: 0,
            len: 0,
        }
    }
}

impl Into<Box<[u8]>> for Chunk {
    fn into(self) -> Box<[u8]> {
        self.buffer
    }
}

#[derive(Debug)]
pub struct StackArena {
    store: Vec<Box<[u8]>>,
    stack: Vec<NonNull<[u8]>>,
    current: Chunk,
}

impl StackArena {
    pub fn new() -> Self {
        Self {
            store: Vec::new(),
            stack: Vec::new(),
            current: Chunk::new(1024),
        }
    }

    pub fn len(&self) -> usize {
        self.stack.len()
    }

    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    pub fn push<P: AsRef<[u8]>>(&mut self, data: P) -> NonNull<[u8]> {
        let data = data.as_ref();
        let object = unsafe { self.allocate(Layout::for_value(data)) }.unwrap();
        unsafe {
            object.cast().copy_from_nonoverlapping(
                NonNull::new_unchecked(data.as_ptr().cast_mut()),
                data.len(),
            )
        };
        self.finish()
    }

    pub fn pop(&mut self) {
        debug_assert_eq!(
            self.current.len, 0,
            "cannot pop while having partial object"
        );
        let object = self.stack.pop().expect("stack underflow");
        self.current.offset -= object.len();
    }

    pub fn extend<P: AsRef<[u8]>>(&mut self, data: P) {
        let data = data.as_ref();
        unsafe {
            let len = self.current.len;
            let ptr = self.current.buffer.as_mut_ptr().add(self.current.offset);
            let object = self
                .grow(
                    NonNull::new_unchecked(ptr),
                    Layout::array::<u8>(self.current.len).unwrap(),
                    Layout::array::<u8>(self.current.len + data.len()).unwrap(),
                )
                .unwrap();
            object.cast().add(len).copy_from_nonoverlapping(
                NonNull::new_unchecked(data.as_ptr().cast_mut()),
                data.len(),
            )
        }
    }

    pub fn finish(&mut self) -> NonNull<[u8]> {
        let object = self.current.object();
        self.current.offset += object.len();
        self.current.len = 0;
        self.stack.push(object);
        object
    }

    pub fn free(&mut self, data: &[u8]) {
        let data = data.as_ref();
        unsafe {
            self.deallocate(
                NonNull::new_unchecked(data.as_ptr().cast_mut()),
                Layout::for_value(data),
            )
        };
    }
}

impl std::fmt::Write for StackArena {
    fn write_str(&mut self, s: &str) -> std::fmt::Result {
        self.extend(s);
        Ok(())
    }
}

impl Allocator for StackArena {
    unsafe fn allocate(
        &mut self,
        layout: std::alloc::Layout,
    ) -> Result<std::ptr::NonNull<[u8]>, crate::AllocError> {
        debug_assert_eq!(self.current.len, 0, "use grow instead");
        let mut offset = self.current.offset
            + self
                .current
                .buffer
                .as_ptr()
                .add(self.current.offset)
                .align_offset(layout.align());
        if offset + layout.size() > self.current.capacity() {
            let mut capacity = self.current.capacity();
            while capacity < layout.size() {
                capacity *= 2;
            }
            let old_chunk = std::mem::replace(&mut self.current, Chunk::new(capacity));
            if old_chunk.offset != 0 {
                self.store.push(old_chunk.into());
            }
            offset = self.current.buffer.as_ptr().align_offset(layout.align());
        }
        self.current.offset = offset;
        self.current.len = layout.size();
        Ok(self.current.object())
    }

    unsafe fn deallocate(&mut self, ptr: std::ptr::NonNull<u8>, layout: std::alloc::Layout) {
        let object = NonNull::slice_from_raw_parts(ptr, layout.size());
        while let Some(top) = self.stack.pop() {
            // TODO: self.store.pop_if(xxx)
            if std::ptr::eq(object.as_ptr(), top.as_ptr()) {
                break;
            }
        }
    }

    unsafe fn grow(
        &mut self,
        ptr: NonNull<u8>,
        old_layout: std::alloc::Layout,
        new_layout: std::alloc::Layout,
    ) -> Result<NonNull<[u8]>, crate::AllocError> {
        match old_layout.size().cmp(&new_layout.size()) {
            std::cmp::Ordering::Less => {
                debug_assert_eq!(
                    ptr.as_ptr().cast_const(),
                    self.current.buffer.as_ptr().add(self.current.offset)
                );
                debug_assert_eq!(old_layout.size(), self.current.len);
                debug_assert_eq!(old_layout.align(), new_layout.align());

                let mut capacity = self.current.capacity();
                if self.current.offset + new_layout.size() > capacity {
                    while capacity < new_layout.size() {
                        capacity *= 2;
                    }
                    let old_chunk = std::mem::replace(&mut self.current, Chunk::new(capacity));
                    let object = old_chunk.object();
                    self.current.offset = self
                        .current
                        .buffer
                        .as_ptr()
                        .align_offset(new_layout.align());
                    self.current
                        .buffer
                        .as_mut_ptr()
                        .add(self.current.offset)
                        .copy_from_nonoverlapping(object.as_ptr().cast(), object.len());
                    if old_chunk.offset != 0 {
                        self.store.push(old_chunk.into());
                    }
                }
                self.current.len = new_layout.size();
                Ok(self.current.object())
            }
            std::cmp::Ordering::Equal => Ok(NonNull::slice_from_raw_parts(ptr, new_layout.size())),
            std::cmp::Ordering::Greater => panic!("use shrink instead"),
        }
    }

    unsafe fn shrink(
        &mut self,
        ptr: NonNull<u8>,
        old_layout: Layout,
        new_layout: Layout,
    ) -> Result<NonNull<[u8]>, crate::AllocError> {
        match old_layout.size().cmp(&new_layout.size()) {
            std::cmp::Ordering::Less => panic!("use grow instead"),
            std::cmp::Ordering::Equal => Ok(NonNull::slice_from_raw_parts(ptr, new_layout.size())),
            std::cmp::Ordering::Greater => {
                debug_assert_eq!(
                    ptr.as_ptr().cast_const(),
                    self.current.buffer.as_ptr().add(self.current.offset)
                );
                debug_assert_eq!(old_layout.size(), self.current.len);
                debug_assert_eq!(old_layout.align(), new_layout.align());
                self.current.len = new_layout.size();
                Ok(NonNull::slice_from_raw_parts(ptr, new_layout.size()))
            }
        }
    }
}

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

    use super::*;

    #[test]
    fn test_lifecycle() {
        let mut stack = StackArena::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_new() {
        let stack = StackArena::new();
        assert_eq!(stack.len(), 0);
        assert!(stack.is_empty());
        assert_eq!(stack.store.len(), 0);
        assert_eq!(stack.stack.len(), 0);
    }

    #[test]
    fn test_push_pop() {
        let mut stack = StackArena::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 = StackArena::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 = StackArena::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 = StackArena::new();

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

        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 = StackArena::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.free(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 = StackArena::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 = StackArena::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 = StackArena::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 = StackArena::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 = StackArena::new();

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

        // Free all objects by freeing the first one
        let first = unsafe { stack.stack[0].as_ref() };
        stack.free(first);

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

    #[test]
    #[should_panic]
    fn test_free_nonexistent() {
        let mut stack = StackArena::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.free(dummy);

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