memkit 0.1.0-beta.1

//! Frame arena - a fast bump allocator with reset capability.

use std::alloc::{alloc, dealloc, Layout};
use std::cell::Cell;
use std::ptr::NonNull;

/// A bump allocator arena for frame-temporary allocations.
///
/// Allocations are O(1) - just bump a pointer.
/// Reset is O(1) - just reset the pointer to the start.
pub struct FrameArena {
    /// Start of the arena memory.
    base: NonNull<u8>,
    /// Total size of the arena.
    size: usize,
    /// Current allocation head (offset from base).
    head: Cell<usize>,
}

impl FrameArena {
    /// Create a new arena with the given size.
    pub fn new(size: usize) -> Option<Self> {
        if size == 0 {
            return None;
        }

        // Allocate the arena memory with page alignment
        let layout = Layout::from_size_align(size, 4096).ok()?;
        let ptr = unsafe { alloc(layout) };
        let base = NonNull::new(ptr)?;

        Some(Self {
            base,
            size,
            head: Cell::new(0),
        })
    }

    /// Allocate memory from the arena.
    ///
    /// Returns null if the arena is full.
    #[inline(always)]
    pub fn alloc(&self, layout: Layout) -> *mut u8 {
        let align = layout.align();
        let size = layout.size();

        // Align the head
        let current = self.head.get();
        let aligned = (current + align - 1) & !(align - 1);

        // Check if we have enough space
        if aligned + size > self.size {
            return std::ptr::null_mut();
        }

        // Bump the head
        self.head.set(aligned + size);

        // Return the pointer
        unsafe { self.base.as_ptr().add(aligned) }
    }

    /// Allocate and initialize a value.
    #[inline(always)]
    pub fn alloc_value<T>(&self, value: T) -> *mut T {
        let layout = Layout::new::<T>();
        let ptr = self.alloc(layout) as *mut T;
        if !ptr.is_null() {
            unsafe { ptr.write(value) };
        }
        ptr
    }

    /// Allocate a slice of uninitialized memory.
    #[inline(always)]
    pub fn alloc_slice<T>(&self, len: usize) -> *mut T {
        if len == 0 {
            return std::ptr::null_mut();
        }
        let layout = Layout::array::<T>(len).unwrap();
        self.alloc(layout) as *mut T
    }

    /// Reset the arena, invalidating all allocations.
    ///
    /// This is O(1) - we just reset the head pointer.
    #[inline(always)]
    pub fn reset(&self) {
        self.head.set(0);
    }

    /// Get the current head position (for checkpointing).
    #[inline(always)]
    pub fn head(&self) -> usize {
        self.head.get()
    }

    /// Reset to a specific head position (for checkpoint restore).
    pub fn reset_to(&self, pos: usize) {
        debug_assert!(pos <= self.head.get(), "Cannot reset forward");
        self.head.set(pos);
    }

    /// Get the number of bytes allocated.
    pub fn allocated(&self) -> usize {
        self.head.get()
    }

    /// Get the number of bytes remaining.
    pub fn remaining(&self) -> usize {
        self.size - self.head.get()
    }

    /// Get the total arena size.
    pub fn capacity(&self) -> usize {
        self.size
    }

    /// Check if the arena is empty.
    pub fn is_empty(&self) -> bool {
        self.head.get() == 0
    }
}

impl Drop for FrameArena {
    fn drop(&mut self) {
        unsafe {
            let layout = Layout::from_size_align_unchecked(self.size, 4096);
            dealloc(self.base.as_ptr(), layout);
        }
    }
}

// FrameArena is not Send/Sync - it's meant to be thread-local
// The MkAllocator handles thread-safety at a higher level

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

    #[test]
    fn test_arena_alloc() {
        let arena = FrameArena::new(4096).unwrap();

        let ptr1 = arena.alloc_value(42u32);
        assert!(!ptr1.is_null());
        assert_eq!(unsafe { *ptr1 }, 42);

        let ptr2 = arena.alloc_value(123u64);
        assert!(!ptr2.is_null());
        assert_eq!(unsafe { *ptr2 }, 123);

        assert!(arena.allocated() > 0);
    }

    #[test]
    fn test_arena_reset() {
        let arena = FrameArena::new(4096).unwrap();

        arena.alloc_value(42u32);
        arena.alloc_value(123u64);
        assert!(arena.allocated() > 0);

        arena.reset();
        assert_eq!(arena.allocated(), 0);
    }

    #[test]
    fn test_arena_checkpoint() {
        let arena = FrameArena::new(4096).unwrap();

        arena.alloc_value(42u32);
        let checkpoint = arena.head();

        arena.alloc_value(123u64);
        arena.alloc_value(456u64);
        assert!(arena.allocated() > checkpoint);

        arena.reset_to(checkpoint);
        assert_eq!(arena.head(), checkpoint);
    }

    #[test]
    fn test_arena_slice() {
        let arena = FrameArena::new(4096).unwrap();

        let slice = arena.alloc_slice::<u32>(100);
        assert!(!slice.is_null());

        // Initialize the slice
        for i in 0..100 {
            unsafe { slice.add(i).write(i as u32) };
        }

        // Verify
        for i in 0..100 {
            assert_eq!(unsafe { *slice.add(i) }, i as u32);
        }
    }

    #[test]
    fn test_arena_full() {
        let arena = FrameArena::new(64).unwrap();

        // Fill it up
        let ptr = arena.alloc_slice::<u8>(64);
        assert!(!ptr.is_null());

        // Should fail now
        let ptr2 = arena.alloc_value(42u32);
        assert!(ptr2.is_null());
    }
}