memkit 0.1.1-beta.1

//! Ultra Arena - The ultimate hybrid allocator combining all optimizations.
//!
//! Features:
//! - Zero-overhead bump allocation (like MkFastArena)
//! - SIMD-accelerated slice initialization
//! - Branch prediction hints
//! - Optional thread-local mode
//! - Const generic sizing for zero-cost configuration
//! - Arena chaining for unlimited growth

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

use super::hints::{likely, unlikely, cold_path};
use super::simd;

/// A chunk of arena memory.
struct ArenaChunk {
    /// Memory block.
    base: NonNull<u8>,
    /// Size of this chunk.
    size: usize,
    /// Next chunk in the chain.
    next: Option<Box<ArenaChunk>>,
}

impl ArenaChunk {
    fn new(size: usize) -> Option<Self> {
        let layout = Layout::from_size_align(size, 4096).ok()?;
        let ptr = unsafe { alloc(layout) };
        let base = NonNull::new(ptr)?;
        
        Some(Self {
            base,
            size,
            next: None,
        })
    }
}

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

/// The ultimate hybrid allocator.
///
/// Combines all optimization techniques:
/// - **Bump allocation**: O(1) allocation with pointer bump
/// - **SIMD initialization**: AVX2/SSE2 accelerated fills
/// - **Branch hints**: Optimized hot paths
/// - **Arena chaining**: Automatic growth without reallocation
/// - **Const generics**: Zero-cost compile-time configuration
///
/// # Type Parameters
///
/// - `CHUNK_SIZE`: Size of each arena chunk (default 4MB)
/// - `AUTO_GROW`: Whether to automatically allocate new chunks (default true)
///
/// # Example
///
/// ```rust
/// use memkit::MkUltraArena;
///
/// // Default: 4MB chunks, auto-grow enabled
/// let arena = MkUltraArena::<4194304, true>::new();
///
/// // Custom: 1MB chunks, no auto-grow
/// let fixed = MkUltraArena::<1048576, false>::new();
///
/// unsafe {
///     let x = arena.alloc(42u64).unwrap();
///     let slice = arena.alloc_slice_simd(1000, 0.0f32).unwrap();
/// }
/// arena.reset();
/// ```
pub struct MkUltraArena<const CHUNK_SIZE: usize = 4194304, const AUTO_GROW: bool = true> {
    /// Current chunk being allocated from.
    current: Cell<NonNull<u8>>,
    /// End of current chunk.
    end: Cell<*const u8>,
    /// First chunk in the chain.
    chunks: Cell<Option<Box<ArenaChunk>>>,
    /// Total bytes allocated across all chunks.
    total_allocated: Cell<usize>,
    /// Marker for !Send + !Sync.
    _marker: PhantomData<*mut u8>,
}

impl<const CHUNK_SIZE: usize, const AUTO_GROW: bool> MkUltraArena<CHUNK_SIZE, AUTO_GROW> {
    /// Create a new ultra arena.
    #[inline]
    pub fn new() -> Self {
        let chunk = ArenaChunk::new(CHUNK_SIZE).expect("Failed to allocate arena chunk");
        let base = chunk.base;
        let end = unsafe { base.as_ptr().add(chunk.size) };
        
        Self {
            current: Cell::new(base),
            end: Cell::new(end),
            chunks: Cell::new(Some(Box::new(chunk))),
            total_allocated: Cell::new(0),
            _marker: PhantomData,
        }
    }
    
    /// Allocate and initialize a value.
    ///
    /// # Safety
    ///
    /// The caller must ensure that the returned reference does not outlive the arena.
    #[allow(clippy::mut_from_ref)]
    #[inline(always)]
    pub unsafe fn alloc<T>(&self, value: T) -> Option<&mut T> {
        let ptr = self.alloc_raw::<T>()?;
        unsafe {
            ptr.write(value);
            Some(&mut *ptr)
        }
    }
    
    /// Allocate memory without initialization.
    #[inline(always)]
    pub fn alloc_raw<T>(&self) -> Option<*mut T> {
        let layout = Layout::new::<T>();
        self.alloc_layout(layout).map(|p| p as *mut T)
    }
    
    /// Core allocation with branch hints and auto-grow.
    #[inline(always)]
    fn alloc_layout(&self, layout: Layout) -> Option<*mut u8> {
        let current = self.current.get().as_ptr();
        let align = layout.align();
        let size = layout.size();
        
        // Align up
        let aligned = ((current as usize + align - 1) & !(align - 1)) as *mut u8;
        let new_ptr = unsafe { aligned.add(size) };
        
        // Fast path: fits in current chunk
        if likely(new_ptr <= self.end.get() as *mut u8) {
            self.current.set(unsafe { NonNull::new_unchecked(new_ptr) });
            self.total_allocated.set(self.total_allocated.get() + size);
            return Some(aligned);
        }
        
        // Slow path: need to grow
        if AUTO_GROW {
            cold_path(|| unsafe { self.grow_and_alloc(layout) })
        } else {
            None
        }
    }
    
    /// Grow the arena and retry allocation.
    ///
    /// # Safety
    ///
    /// This is an internal method that handles arena growth.
    #[cold]
    unsafe fn grow_and_alloc(&self, layout: Layout) -> Option<*mut u8> {
        // Calculate new chunk size (at least CHUNK_SIZE, or larger if needed)
        let needed = layout.size() + layout.align();
        let new_size = std::cmp::max(CHUNK_SIZE, needed.next_power_of_two());
        
        let mut new_chunk = ArenaChunk::new(new_size)?;
        
        // Take ownership of old chunks
        let old_chunks = self.chunks.take();
        new_chunk.next = old_chunks;
        
        let base = new_chunk.base;
        let end = unsafe { base.as_ptr().add(new_chunk.size) };
        
        self.chunks.set(Some(Box::new(new_chunk)));
        self.current.set(base);
        self.end.set(end);
        
        // Retry allocation
        self.alloc_layout(layout)
    }
    
    /// Allocate a slice and fill with zeros using SIMD.
    ///
    /// # Safety
    ///
    /// The caller must ensure that the returned slice does not outlive the arena.
    #[allow(clippy::mut_from_ref)]
    #[inline(always)]
    pub unsafe fn alloc_slice_zero<T>(&self, len: usize) -> Option<&mut [T]> {
        if unlikely(len == 0) {
            return Some(&mut []);
        }
        
        let layout = Layout::array::<T>(len).ok()?;
        let ptr = self.alloc_layout(layout)?;
        
        // SIMD zero fill
        simd::fill_zero(ptr, layout.size());
        
        Some(unsafe { std::slice::from_raw_parts_mut(ptr as *mut T, len) })
    }
    
    /// Allocate a slice and fill with a value using SIMD when possible.
    ///
    /// # Safety
    ///
    /// The caller must ensure that the returned slice does not outlive the arena.
    #[allow(clippy::mut_from_ref)]
    #[inline(always)]
    pub unsafe fn alloc_slice_simd<T: Copy>(&self, len: usize, value: T) -> Option<&mut [T]> {
        if unlikely(len == 0) {
            return Some(&mut []);
        }
        
        let layout = Layout::array::<T>(len).ok()?;
        let ptr = self.alloc_layout(layout)?;
        let typed_ptr = ptr as *mut T;
        
        // Dispatch to SIMD based on type size
        match std::mem::size_of::<T>() {
            1 => {
                let byte = unsafe { std::mem::transmute_copy::<T, u8>(&value) };
                simd::fill_byte(ptr, byte, len);
            }
            4 => {
                let bits = unsafe { std::mem::transmute_copy::<T, u32>(&value) };
                simd::fill_u32(typed_ptr as *mut u32, bits, len);
            }
            8 => {
                let bits = unsafe { std::mem::transmute_copy::<T, u64>(&value) };
                simd::fill_u64(typed_ptr as *mut u64, bits, len);
            }
            _ => {
                // Scalar fallback for other sizes
                for i in 0..len {
                    unsafe { typed_ptr.add(i).write(value) };
                }
            }
        }
        
        Some(unsafe { std::slice::from_raw_parts_mut(typed_ptr, len) })
    }
    
    /// Allocate f32 slice with SIMD fill.
    ///
    /// # Safety
    ///
    /// The caller must ensure that the returned slice does not outlive the arena.
    #[allow(clippy::mut_from_ref)]
    #[inline(always)]
    pub unsafe fn alloc_f32(&self, len: usize, value: f32) -> Option<&mut [f32]> {
        if unlikely(len == 0) {
            return Some(&mut []);
        }
        
        let layout = Layout::array::<f32>(len).ok()?;
        let ptr = self.alloc_layout(layout)?;
        simd::fill_f32(ptr as *mut f32, value, len);
        Some(unsafe { std::slice::from_raw_parts_mut(ptr as *mut f32, len) })
    }
    
    /// Allocate f64 slice with SIMD fill.
    ///
    /// # Safety
    ///
    /// The caller must ensure that the returned slice does not outlive the arena.
    #[allow(clippy::mut_from_ref)]
    #[inline(always)]
    pub unsafe fn alloc_f64(&self, len: usize, value: f64) -> Option<&mut [f64]> {
        if unlikely(len == 0) {
            return Some(&mut []);
        }
        
        let layout = Layout::array::<f64>(len).ok()?;
        let ptr = self.alloc_layout(layout)?;
        simd::fill_f64(ptr as *mut f64, value, len);
        Some(unsafe { std::slice::from_raw_parts_mut(ptr as *mut f64, len) })
    }
    
    /// Reset the arena, keeping all allocated chunks.
    #[inline(always)]
    pub fn reset(&self) {
        // Reset to first chunk
        if let Some(ref chunks) = unsafe { &*self.chunks.as_ptr() } {
            self.current.set(chunks.base);
            self.end.set(unsafe { chunks.base.as_ptr().add(chunks.size) });
        }
        self.total_allocated.set(0);
    }
    
    /// Reset and release all but the first chunk.
    pub fn reset_and_shrink(&self) {
        if let Some(mut chunks) = self.chunks.take() {
            // Drop all but first chunk
            chunks.next = None;
            
            self.current.set(chunks.base);
            self.end.set(unsafe { chunks.base.as_ptr().add(chunks.size) });
            self.chunks.set(Some(chunks));
        }
        self.total_allocated.set(0);
    }
    
    /// Get total bytes allocated.
    #[inline]
    pub fn allocated(&self) -> usize {
        self.total_allocated.get()
    }
    
    /// Get number of chunks.
    pub fn chunk_count(&self) -> usize {
        let mut count = 0;
        let mut current = unsafe { &*self.chunks.as_ptr() };
        while let Some(ref chunk) = current {
            count += 1;
            current = &chunk.next;
        }
        count
    }
    
    /// Get total capacity across all chunks.
    pub fn total_capacity(&self) -> usize {
        let mut total = 0;
        let mut current = unsafe { &*self.chunks.as_ptr() };
        while let Some(ref chunk) = current {
            total += chunk.size;
            current = &chunk.next;
        }
        total
    }
}

impl<const CHUNK_SIZE: usize, const AUTO_GROW: bool> Default for MkUltraArena<CHUNK_SIZE, AUTO_GROW> {
    fn default() -> Self {
        Self::new()
    }
}

/// Type alias for a fixed-size arena that doesn't auto-grow.
pub type MkFixedArena<const SIZE: usize = 4194304> = MkUltraArena<SIZE, false>;

/// Type alias for a small arena (1MB).
pub type MkSmallArena = MkUltraArena<1048576, true>;

/// Type alias for a large arena (16MB).
pub type MkLargeArena = MkUltraArena<16777216, true>;

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

    #[test]
    fn test_ultra_alloc() {
        let arena: MkUltraArena<4194304, true> = MkUltraArena::new();
        
        unsafe {
            let x = arena.alloc(42u64).unwrap();
            assert_eq!(*x, 42);
            
            let y = arena.alloc(123u32).unwrap();
            assert_eq!(*y, 123);
        }
        
        arena.reset();
    }

    #[test]
    fn test_ultra_slice_zero() {
        let arena: MkUltraArena<4194304, true> = MkUltraArena::new();
        
        unsafe {
            let slice = arena.alloc_slice_zero::<u64>(100).unwrap();
            assert_eq!(slice.len(), 100);
            assert!(slice.iter().all(|&v| v == 0));
        }
        
        arena.reset();
    }

    #[test]
    fn test_ultra_slice_simd() {
        let arena: MkUltraArena<4194304, true> = MkUltraArena::new();
        
        unsafe {
            let slice = arena.alloc_slice_simd(100, 42u64).unwrap();
            assert_eq!(slice.len(), 100);
            assert!(slice.iter().all(|&v| v == 42));
        }
        
        arena.reset();
    }

    #[test]
    fn test_ultra_f32() {
        let arena: MkUltraArena<4194304, true> = MkUltraArena::new();
        
        unsafe {
            let slice = arena.alloc_f32(100, 3.14).unwrap();
            assert_eq!(slice.len(), 100);
            assert!(slice.iter().all(|&v| (v - 3.14).abs() < f32::EPSILON));
        }
        
        arena.reset();
    }

    #[test]
    fn test_ultra_auto_grow() {
        // Small arena that will need to grow
        let arena: MkUltraArena<1024, true> = MkUltraArena::new();
        
        // Allocate more than 1KB
        unsafe {
            for _ in 0..100 {
                let _ = arena.alloc([0u64; 16]).unwrap(); // 128 bytes each
            }
        }
        
        assert!(arena.chunk_count() > 1);
        arena.reset();
    }

    #[test]
    fn test_fixed_arena_no_grow() {
        let arena: MkFixedArena<1024> = MkFixedArena::new();
        
        unsafe {
            // Fill it up
            let _ = arena.alloc([0u8; 512]);
            let _ = arena.alloc([0u8; 400]);
            
            // Should fail - no auto-grow
            let result = arena.alloc([0u8; 256]);
            assert!(result.is_none());
        }
    }

    #[test]
    fn test_reset_and_shrink() {
        let arena: MkUltraArena<1024, true> = MkUltraArena::new();
        
        unsafe {
            // Force growth
            for _ in 0..100 {
                let _ = arena.alloc([0u64; 16]);
            }
        }
        
        let chunks_before = arena.chunk_count();
        assert!(chunks_before > 1);
        
        arena.reset_and_shrink();
        
        assert_eq!(arena.chunk_count(), 1);
        assert_eq!(arena.allocated(), 0);
    }
}