llvm-native-core 0.1.15

//! BumpPtrAllocator — fast bump-pointer allocator for short-lived
//! allocations during compilation.
//!
//! A BumpPtrAllocator allocates memory by bumping a pointer within
//! large pre-allocated slabs. Individual allocations cannot be freed;
//! the entire allocator is reset at once. This is ideal for:
//! - AST node allocation during parsing
//! - Temporary data during optimization passes
//! - SelectionDAG node allocation
//!
//! Clean-room behavioral reconstruction. No LLVM source is consulted.

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

// ═══════════════════════════════════════════════════════════════════════════
// Memory Slab
// ═══════════════════════════════════════════════════════════════════════════

/// A single slab of memory used by the bump allocator.
struct Slab {
    /// Pointer to the start of the slab.
    ptr: NonNull<u8>,
    /// Total size of the slab in bytes.
    size: usize,
    /// Current allocation position within the slab.
    position: usize,
}

impl Slab {
    /// Create a new slab with the given size.
    fn new(size: usize) -> Self {
        let layout = Layout::from_size_align(size, 16).unwrap();
        let ptr = unsafe { alloc(layout) };
        Self {
            ptr: NonNull::new(ptr).expect("allocation failed"),
            size,
            position: 0,
        }
    }

    /// Allocate `size` bytes with `align` alignment from this slab.
    /// Returns None if there isn't enough space.
    fn allocate(&mut self, size: usize, align: usize) -> Option<NonNull<u8>> {
        let aligned_pos = (self.position + align - 1) & !(align - 1);
        if aligned_pos + size > self.size {
            return None;
        }
        let result = unsafe { NonNull::new_unchecked(self.ptr.as_ptr().add(aligned_pos)) };
        self.position = aligned_pos + size;
        Some(result)
    }

    /// Remaining bytes in this slab.
    fn remaining(&self) -> usize {
        self.size - self.position
    }
}

impl Drop for Slab {
    fn drop(&mut self) {
        let layout = Layout::from_size_align(self.size, 16).unwrap();
        unsafe {
            dealloc(self.ptr.as_ptr(), layout);
        }
    }
}

// ═══════════════════════════════════════════════════════════════════════════
// BumpPtrAllocator
// ═══════════════════════════════════════════════════════════════════════════

/// A fast bump-pointer allocator.
///
/// Allocates memory by bumping a pointer. Individual deallocations
/// are not possible; the entire allocator is reset at once by calling
/// `reset()`. New slabs are allocated automatically as needed.
pub struct BumpPtrAllocator {
    /// Currently active slabs.
    slabs: Vec<Slab>,
    /// Default slab size for new slabs.
    slab_size: usize,
    /// Total bytes allocated across all slabs.
    total_allocated: usize,
    /// Number of individual allocations.
    allocation_count: usize,
}

impl BumpPtrAllocator {
    /// Create a new bump allocator with the default slab size.
    pub fn new() -> Self {
        Self::with_slab_size(4096)
    }

    /// Create a new bump allocator with a specific slab size.
    pub fn with_slab_size(slab_size: usize) -> Self {
        Self {
            slabs: Vec::new(),
            slab_size,
            total_allocated: 0,
            allocation_count: 0,
        }
    }

    /// Allocate `size` bytes with the given alignment.
    pub fn allocate(&mut self, size: usize, align: usize) -> NonNull<u8> {
        if size == 0 {
            return NonNull::dangling();
        }

        // Try to allocate from the current slab
        if let Some(slab) = self.slabs.last_mut() {
            if let Some(ptr) = slab.allocate(size, align) {
                self.total_allocated += size;
                self.allocation_count += 1;
                return ptr;
            }
        }

        // Need a new slab — make it at least large enough for this allocation
        let new_slab_size = self.slab_size.max(size + align);
        let mut new_slab = Slab::new(new_slab_size);
        let ptr = new_slab
            .allocate(size, align)
            .expect("new slab should have space");

        self.slabs.push(new_slab);
        self.total_allocated += size;
        self.allocation_count += 1;
        ptr
    }

    /// Allocate an object of type T and return a mutable reference.
    /// The memory is zero-initialized.
    pub fn allocate_obj<T>(&mut self) -> &mut T {
        let size = std::mem::size_of::<T>();
        let align = std::mem::align_of::<T>();
        let ptr = self.allocate(size, align);
        unsafe {
            std::ptr::write_bytes(ptr.as_ptr(), 0, size);
            &mut *ptr.as_ptr().cast::<T>()
        }
    }

    /// Allocate a slice of T with `count` elements.
    pub fn allocate_slice<T>(&mut self, count: usize) -> &mut [T] {
        let size = std::mem::size_of::<T>() * count;
        let align = std::mem::align_of::<T>();
        let ptr = self.allocate(size, align);
        unsafe {
            std::ptr::write_bytes(ptr.as_ptr(), 0, size);
            std::slice::from_raw_parts_mut(ptr.as_ptr().cast::<T>(), count)
        }
    }

    /// Reset the allocator — deallocates all slabs and starts fresh.
    pub fn reset(&mut self) {
        self.slabs.clear();
        self.total_allocated = 0;
        self.allocation_count = 0;
    }

    /// Total bytes allocated (across all slabs, including unused space).
    pub fn total_bytes(&self) -> usize {
        self.slabs.iter().map(|s| s.size).sum()
    }

    /// Total bytes actually allocated to callers.
    pub fn total_allocated(&self) -> usize {
        self.total_allocated
    }

    /// Number of individual allocations.
    pub fn allocation_count(&self) -> usize {
        self.allocation_count
    }

    /// Number of slabs.
    pub fn slab_count(&self) -> usize {
        self.slabs.len()
    }
}

impl Default for BumpPtrAllocator {
    fn default() -> Self {
        Self::new()
    }
}

impl Drop for BumpPtrAllocator {
    fn drop(&mut self) {
        self.slabs.clear();
    }
}

// ═══════════════════════════════════════════════════════════════════════════
// SpecificBumpPtrAllocator — typed bump allocator for a single type T
// ═══════════════════════════════════════════════════════════════════════════

/// A typed bump allocator for a specific type T.
/// More efficient than the generic allocator since alignment and
/// size are known at compile time.
pub struct SpecificBumpPtrAllocator<T> {
    allocator: BumpPtrAllocator,
    _phantom: std::marker::PhantomData<T>,
}

impl<T> SpecificBumpPtrAllocator<T> {
    pub fn new() -> Self {
        Self {
            allocator: BumpPtrAllocator::new(),
            _phantom: std::marker::PhantomData,
        }
    }

    /// Allocate a single T.
    pub fn allocate(&mut self) -> &mut T {
        self.allocator.allocate_obj::<T>()
    }

    /// Allocate an array of T.
    pub fn allocate_array(&mut self, count: usize) -> &mut [T] {
        self.allocator.allocate_slice::<T>(count)
    }

    /// Reset the allocator.
    pub fn reset(&mut self) {
        self.allocator.reset();
    }

    /// Total allocations made.
    pub fn allocation_count(&self) -> usize {
        self.allocator.allocation_count()
    }
}

impl<T> Default for SpecificBumpPtrAllocator<T> {
    fn default() -> Self {
        Self::new()
    }
}

impl BumpPtrAllocator {
    pub fn with_slab_size_and_alignment(slab_size: usize, _alignment: usize) -> Self {
        Self::with_slab_size(slab_size)
    }

    pub fn allocate_raw(&mut self, layout: std::alloc::Layout) -> NonNull<u8> {
        self.allocate(layout.size(), layout.align())
    }

    pub fn deallocate(&mut self, _ptr: NonNull<u8>) {}

    pub fn identify_object(&self, ptr: *const u8) -> bool {
        if ptr.is_null() {
            return false;
        }
        for slab in &self.slabs {
            let start = slab.ptr.as_ptr() as usize;
            let end = start + slab.size;
            let addr = ptr as usize;
            if addr >= start && addr < end {
                return true;
            }
        }
        false
    }

    pub fn contains_ptr(&self, ptr: *const u8) -> bool {
        self.identify_object(ptr)
    }

    pub fn bytes_remaining(&self) -> usize {
        self.slabs.last().map(|s| s.remaining()).unwrap_or(0)
    }

    pub fn print_stats(&self) {
        eprintln!("BumpPtrAllocator stats:");
        eprintln!("  Slabs: {}", self.slab_count());
        eprintln!("  Total slab bytes: {}", self.total_bytes());
        eprintln!("  Total allocated: {}", self.total_allocated);
        eprintln!("  Allocations: {}", self.allocation_count);
        eprintln!("  Bytes remaining: {}", self.bytes_remaining());
        for (i, slab) in self.slabs.iter().enumerate() {
            eprintln!(
                "  Slab {}: size={}, used={}, remaining={}",
                i,
                slab.size,
                slab.position,
                slab.remaining()
            );
        }
    }

    pub fn reset_keep_first(&mut self) {
        if let Some(first) = self.slabs.first_mut() {
            first.position = 0;
        }
        self.slabs.truncate(1);
        self.total_allocated = 0;
        self.allocation_count = 0;
    }

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

    pub fn current_slab_size(&self) -> usize {
        self.slabs.last().map(|s| s.size).unwrap_or(0)
    }

    pub fn current_slab_position(&self) -> usize {
        self.slabs.last().map(|s| s.position).unwrap_or(0)
    }

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

    pub fn utilization(&self) -> f64 {
        let total = self.total_bytes();
        if total == 0 {
            return 0.0;
        }
        self.total_allocated as f64 / total as f64
    }

    pub fn allocate_obj_with<T: Copy>(&mut self, value: T) -> &mut T {
        let obj = self.allocate_obj::<T>();
        *obj = value;
        obj
    }

    fn grow(&mut self, min_size: usize) {
        let new_size = self.slab_size.max(min_size);
        self.slabs.push(Slab::new(new_size));
    }

    pub fn allocate_zeroed(&mut self, size: usize, align: usize) -> NonNull<u8> {
        let ptr = self.allocate(size, align);
        unsafe {
            std::ptr::write_bytes(ptr.as_ptr(), 0, size);
        }
        ptr
    }

    pub fn allocate_copy(&mut self, src: &[u8]) -> &mut [u8] {
        let dest = self.allocate_slice::<u8>(src.len());
        dest.copy_from_slice(src);
        dest
    }
}

impl Slab {
    pub fn contains(&self, ptr: *const u8) -> bool {
        let start = self.ptr.as_ptr() as usize;
        let end = start + self.size;
        let addr = ptr as usize;
        addr >= start && addr < end
    }

    pub fn start(&self) -> *const u8 {
        self.ptr.as_ptr() as *const u8
    }

    pub fn end(&self) -> *const u8 {
        unsafe { (self.ptr.as_ptr() as *const u8).add(self.size) }
    }

    pub fn allocated(&self) -> usize {
        self.position
    }

    pub fn reset(&mut self) {
        self.position = 0;
    }
}

impl<T> SpecificBumpPtrAllocator<T> {
    pub fn with_slab_size(slab_size: usize) -> Self {
        Self {
            allocator: BumpPtrAllocator::with_slab_size(slab_size),
            _phantom: std::marker::PhantomData,
        }
    }

    pub fn contains(&self, ptr: *const T) -> bool {
        self.allocator.identify_object(ptr as *const u8)
    }

    pub fn reset_keep_first(&mut self) {
        self.allocator.reset_keep_first();
    }

    pub fn print_stats(&self) {
        self.allocator.print_stats();
    }

    pub fn allocate_with(&mut self, value: T) -> &mut T
    where
        T: Copy,
    {
        let obj = self.allocate();
        *obj = value;
        obj
    }

    pub fn num_allocated(&self) -> usize {
        self.allocation_count()
    }

    pub fn allocate_array_with(&mut self, count: usize, value: T) -> &mut [T]
    where
        T: Copy,
    {
        let slice = self.allocate_array(count);
        for elem in slice.iter_mut() {
            *elem = value;
        }
        slice
    }
}

// ═══════════════════════════════════════════════════════════════════════════
// SlabPool — manages a pool of pre-allocated slabs for reuse
// ═══════════════════════════════════════════════════════════════════════════

/// A pool of slabs that can be reused after reset, reducing allocation overhead.
pub struct SlabPool {
    slabs: Vec<Slab>,
    free_slabs: Vec<Slab>,
    slab_size: usize,
}

impl SlabPool {
    pub fn new(slab_size: usize) -> Self {
        Self {
            slabs: Vec::new(),
            free_slabs: Vec::new(),
            slab_size,
        }
    }

    pub fn get_slab(&mut self) -> Slab {
        match self.free_slabs.pop() { Some(slab) => {
            slab
        } _ => {
            Slab::new(self.slab_size)
        }}
    }

    pub fn return_slab(&mut self, slab: Slab) {
        self.free_slabs.push(slab);
    }

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

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

// ═══════════════════════════════════════════════════════════════════════════
// BumpPtrAllocator with slab pool support
// ═══════════════════════════════════════════════════════════════════════════

impl BumpPtrAllocator {
    /// Allocate memory and return a pointer along with the size allocated.
    pub fn allocate_sized(&mut self, size: usize, align: usize) -> (NonNull<u8>, usize) {
        (self.allocate(size, align), size)
    }

    /// Allocate a string (copy from &str).
    pub fn allocate_str(&mut self, s: &str) -> &str {
        let bytes = s.as_bytes();
        let dest = self.allocate_slice::<u8>(bytes.len());
        dest.copy_from_slice(bytes);
        unsafe { std::str::from_utf8_unchecked(dest) }
    }

    /// Check if a pointer is allocated within the first slab only.
    pub fn identify_in_first_slab(&self, ptr: *const u8) -> bool {
        self.slabs.first().map(|s| s.contains(ptr)).unwrap_or(false)
    }

    /// Get an iterator over all allocated slabs.
    pub fn slabs(&self) -> &[Slab] {
        &self.slabs
    }

    /// Get the first slab (mutable).
    pub fn first_slab_mut(&mut self) -> Option<&mut Slab> {
        self.slabs.first_mut()
    }

    /// Check if there are no slabs.
    pub fn has_no_slabs(&self) -> bool {
        self.slabs.is_empty()
    }
}

// ═══════════════════════════════════════════════════════════════════════════
// Tests
// ═══════════════════════════════════════════════════════════════════════════

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

    #[test]
    fn test_bump_allocator_basic() {
        let mut alloc = BumpPtrAllocator::new();
        let ptr = alloc.allocate(100, 8);
        assert!(!ptr.as_ptr().is_null());
        assert_eq!(alloc.allocation_count(), 1);
    }

    #[test]
    fn test_bump_allocator_multiple() {
        let mut alloc = BumpPtrAllocator::new();
        let p1 = alloc.allocate(50, 8);
        let p2 = alloc.allocate(50, 8);
        assert_ne!(p1, p2);
        assert_eq!(alloc.allocation_count(), 2);
    }

    #[test]
    fn test_bump_allocator_large_allocation() {
        let mut alloc = BumpPtrAllocator::with_slab_size(256);
        // This should trigger a new slab
        let ptr = alloc.allocate(1024, 8);
        assert!(!ptr.as_ptr().is_null());
        assert!(alloc.slab_count() >= 1);
    }

    #[test]
    fn test_bump_allocator_reset() {
        let mut alloc = BumpPtrAllocator::new();
        alloc.allocate(100, 8);
        alloc.allocate(200, 8);
        assert_eq!(alloc.allocation_count(), 2);
        alloc.reset();
        assert_eq!(alloc.allocation_count(), 0);
        assert_eq!(alloc.slab_count(), 0);
    }

    #[test]
    fn test_bump_allocator_obj() {
        let mut alloc = BumpPtrAllocator::new();
        let obj: &mut u64 = alloc.allocate_obj::<u64>();
        *obj = 42;
        assert_eq!(*obj, 42);
    }

    #[test]
    fn test_bump_allocator_slice() {
        let mut alloc = BumpPtrAllocator::new();
        let slice = alloc.allocate_slice::<u32>(10);
        assert_eq!(slice.len(), 10);
        slice[0] = 123;
        assert_eq!(slice[0], 123);
    }

    #[test]
    fn test_specific_bump_allocator() {
        let mut alloc = SpecificBumpPtrAllocator::<u64>::new();
        // Each allocate() returns a new mutable reference; the previous one
        // must be dropped before allocating again.
        {
            let a = alloc.allocate();
            *a = 100;
        }
        {
            let b = alloc.allocate();
            *b = 200;
        }
        // Verify the count
        assert_eq!(alloc.allocation_count(), 2);
    }

    #[test]
    fn test_bump_allocator_alignment() {
        let mut alloc = BumpPtrAllocator::new();
        let ptr = alloc.allocate(1, 16);
        assert_eq!(ptr.as_ptr() as usize % 16, 0);
    }
}