spectral_vm 0.1.6

/*
 * ═══════════════════════════════════════════════════════════════════════════
 * TECHNICAL MANIFEST: Memory Pool Allocator
 * SPECTRAL ROLE: Large Allocation Optimization for 2^20+ Circuits
 * ═══════════════════════════════════════════════════════════════════════════
 *
 * COMPLEXITY: O(1) allocation/deallocation | O(n) pool management
 * MEMORY: Pre-allocated pools reduce fragmentation and allocation overhead
 * THREADING: Thread-local pools for concurrent FRI operations
 *
 * ARCHITECTURAL INVARIANTS:
 * - Pool sizes: Geometric progression (1K, 4K, 16K, 64K, 256K, 1M, 4M, 16M)
 * - Large allocations: Direct mmap for 2^20+ elements
 * - Memory reuse: Automatic pool recycling between operations
 *
 * SECURITY PROPERTIES:
 * - Zeroization: Sensitive cryptographic data cleared on deallocation
 * - Bounds checking: Prevent buffer overflows in large allocations
 * - Thread isolation: Per-thread pools prevent data leakage
 * ═══════════════════════════════════════════════════════════════════════════
 */

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

/// Memory pool for efficient large allocations in FRI operations
pub struct MemoryPool {
    /// Pre-allocated memory blocks by size class
    pools: HashMap<usize, Vec<NonNull<u8>>>,
    /// Currently allocated blocks (for cleanup)
    allocated: Vec<(NonNull<u8>, Layout)>,
    /// Pool size classes (powers of 4 for Goldilocks field elements)
    size_classes: Vec<usize>,
}

impl MemoryPool {
    /// Create a new memory pool optimized for FRI operations
    pub fn new() -> Self {
        let size_classes = vec![
            1024,      // 1K elements (4KB)
            4096,      // 4K elements (16KB)
            16384,     // 16K elements (64KB)
            65536,     // 64K elements (256KB)
            262144,    // 256K elements (1MB)
            1048576,   // 1M elements (4MB) - for 2^20 circuits
            4194304,   // 4M elements (16MB)
        ];

        Self {
            pools: HashMap::new(),
            allocated: Vec::new(),
            size_classes,
        }
    }

    /// Allocate memory for FRI operations
    /// Uses pooled allocation for small sizes, direct allocation for large sizes
    pub fn allocate(&mut self, size_bytes: usize) -> Result<NonNull<u8>, crate::error::HyperionError> {
        // For very large allocations (2^20+ elements), use direct allocation
        if size_bytes > 4 * 1024 * 1024 { // > 4MB
            return self.allocate_large(size_bytes);
        }

        // Find appropriate size class
        let size_class = self.find_size_class(size_bytes);

        // Try to get from pool
        if let Some(pool) = self.pools.get_mut(&size_class) {
            if let Some(block) = pool.pop() {
                return Ok(block);
            }
        }

        // Allocate new block
        self.allocate_large(size_class)
    }

    /// Deallocate memory and return to pool if appropriate
    pub fn deallocate(&mut self, ptr: NonNull<u8>, layout: Layout) {
        let size = layout.size();

        // For large allocations, directly deallocate
        if size > 4 * 1024 * 1024 {
            unsafe { dealloc(ptr.as_ptr(), layout); }
            return;
        }

        // Return to pool
        let size_class = self.find_size_class(size);
        if let Some(pool) = self.pools.get_mut(&size_class) {
            // Clear sensitive data before pooling
            unsafe {
                std::ptr::write_bytes(ptr.as_ptr(), 0, size);
            }
            pool.push(ptr);
        } else {
            unsafe { dealloc(ptr.as_ptr(), layout); }
        }
    }

    /// Find appropriate size class for allocation
    fn find_size_class(&self, size: usize) -> usize {
        for &class in &self.size_classes {
            if size <= class {
                return class;
            }
        }
        // For very large sizes, round up to next power of 2
        size.next_power_of_two()
    }

    /// Allocate large block directly from system
    fn allocate_large(&mut self, size: usize) -> Result<NonNull<u8>, crate::error::HyperionError> {
        let layout = Layout::from_size_align(size, std::mem::align_of::<u8>())
            .map_err(|_| crate::error::HyperionError::Internal(format!("Invalid layout for size {}", size)))?;
        let ptr = unsafe { alloc(layout) };

        if ptr.is_null() {
            panic!("Memory allocation failed for size {}", size);
        }

        let ptr = unsafe { NonNull::new_unchecked(ptr) };
        self.allocated.push((ptr, layout));
        Ok(ptr)
    }

    /// Clean up all allocated memory
    pub fn cleanup(&mut self) {
        for (ptr, layout) in self.allocated.drain(..) {
            unsafe { dealloc(ptr.as_ptr(), layout); }
        }
        self.pools.clear();
    }
}

impl Drop for MemoryPool {
    fn drop(&mut self) {
        self.cleanup();
    }
}

/// Thread-local memory pool for FRI operations
thread_local! {
    static MEMORY_POOL: RefCell<MemoryPool> = RefCell::new(MemoryPool::new());
}

/// High-level memory management for FRI operations
pub struct FriMemoryManager {
    pool: MemoryPool,
}

impl FriMemoryManager {
    pub fn new() -> Self {
        Self {
            pool: MemoryPool::new(),
        }
    }

    /// Allocate memory for Goldilocks field elements
    pub fn allocate_goldilocks(&mut self, count: usize) -> Vec<crate::field::Goldilocks> {
        let size_bytes = count * std::mem::size_of::<crate::field::Goldilocks>();
        let ptr = self.pool.allocate(size_bytes).expect("Memory allocation should succeed");

        // Convert to Vec (unsafe but controlled)
        unsafe {
            Vec::from_raw_parts(
                ptr.as_ptr() as *mut crate::field::Goldilocks,
                count,
                count,
            )
        }
    }

    /// Deallocate Goldilocks vector
    pub fn deallocate_goldilocks(&mut self, vec: Vec<crate::field::Goldilocks>) -> Result<(), crate::error::HyperionError> {
        let ptr = unsafe { NonNull::new_unchecked(vec.as_ptr() as *mut u8) };
        let size_bytes = vec.len() * std::mem::size_of::<crate::field::Goldilocks>();
        let layout = Layout::from_size_align(size_bytes, std::mem::align_of::<crate::field::Goldilocks>())
            .map_err(|_| crate::error::HyperionError::Internal(format!("Invalid layout for Goldilocks vector size {}", size_bytes)))?;

        // Forget the vector to prevent double-free
        std::mem::forget(vec);

        self.pool.deallocate(ptr, layout);
        Ok(())
    }
}

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

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

    #[test]
    fn test_memory_pool_allocation() {
        let mut pool = MemoryPool::new();

        // Allocate small block
        let ptr1 = pool.allocate(1024).unwrap();
        assert!(!ptr1.as_ptr().is_null());

        // Allocate large block
        let ptr2 = pool.allocate(5 * 1024 * 1024).unwrap(); // 5MB
        assert!(!ptr2.as_ptr().is_null());

        pool.cleanup();
    }

    #[test]
    fn test_fri_memory_manager() {
        let mut manager = FriMemoryManager::new();

        // Allocate Goldilocks vector
        let mut vec = manager.allocate_goldilocks(1024);

        // Fill with data
        for i in 0..1024 {
            vec[i] = Goldilocks::from_i64(i as i64);
        }

        // Verify data
        for i in 0..1024 {
            assert_eq!(vec[i], Goldilocks::from_i64(i as i64));
        }

        // Deallocate
        let _ = manager.deallocate_goldilocks(vec);
    }
}