graph_d 1.3.2

//! Safe memory allocators without unsafe code.
//!
//! This module provides memory allocation strategies that work within
//! Rust's safe memory model, avoiding unsafe code while still providing
//! performance benefits through pooling and arena allocation.

use bumpalo::Bump;
use parking_lot::RwLock;
use std::collections::VecDeque;
use std::sync::Arc;

/// A safe arena allocator using bump allocation.
///
/// This allocator provides fast O(1) allocation by using bump pointer allocation.
/// All allocations are freed together when the arena is reset.
#[allow(clippy::arc_with_non_send_sync)]
pub struct SafeArenaAllocator {
    /// Current arena for allocations
    current_arena: Arc<RwLock<Bump>>,
    /// Backup arena for seamless swapping
    backup_arena: Arc<RwLock<Bump>>,
    /// Size of each arena in bytes
    arena_size: usize,
}

impl SafeArenaAllocator {
    /// Create a new safe arena allocator.
    #[allow(clippy::arc_with_non_send_sync)]
    pub fn new(arena_size: usize) -> Self {
        Self {
            current_arena: Arc::new(RwLock::new(Bump::with_capacity(arena_size))),
            backup_arena: Arc::new(RwLock::new(Bump::with_capacity(arena_size))),
            arena_size,
        }
    }

    /// Reset the current arena and swap it with backup.
    pub fn reset_and_swap(&self) {
        let mut current = self.current_arena.write();
        let mut backup = self.backup_arena.write();

        // Reset the current arena
        current.reset();

        // Swap arenas
        std::mem::swap(&mut *current, &mut *backup);
    }

    /// Get current memory usage.
    pub fn memory_usage(&self) -> usize {
        let current = self.current_arena.read();
        current.allocated_bytes()
    }

    /// Check if arena needs reset (>75% full).
    pub fn needs_reset(&self) -> bool {
        let current = self.current_arena.read();
        current.allocated_bytes() > (self.arena_size * 3) / 4
    }

    /// Execute a closure with automatic arena cleanup.
    pub fn scoped<F, R>(&self, f: F) -> R
    where
        F: FnOnce() -> R,
    {
        let result = f();
        self.reset_and_swap();
        result
    }
}

/// A safe pool allocator for managing reusable objects.
///
/// This allocator maintains pools of boxed objects that can be reused,
/// reducing allocation overhead without requiring unsafe code.
pub struct SafePoolAllocator<T> {
    /// Pool of available objects
    pool: Arc<RwLock<VecDeque<Box<T>>>>,
    /// Factory function for creating new objects
    factory: Arc<dyn Fn() -> T + Send + Sync>,
    /// Maximum pool size
    max_pool_size: usize,
}

impl<T> SafePoolAllocator<T>
where
    T: Default,
{
    /// Create a new safe pool allocator with default factory.
    pub fn new(initial_size: usize, max_size: usize) -> Self {
        let pool = Arc::new(RwLock::new(VecDeque::new()));

        // Pre-populate the pool
        {
            let mut pool_guard = pool.write();
            for _ in 0..initial_size {
                pool_guard.push_back(Box::new(T::default()));
            }
        }

        Self {
            pool,
            factory: Arc::new(|| T::default()),
            max_pool_size: max_size,
        }
    }
}

impl<T> SafePoolAllocator<T> {
    /// Create a new safe pool allocator with custom factory.
    pub fn with_factory<F>(factory: F, initial_size: usize, max_size: usize) -> Self
    where
        F: Fn() -> T + Send + Sync + 'static,
    {
        let pool = Arc::new(RwLock::new(VecDeque::new()));
        let factory = Arc::new(factory);

        // Pre-populate the pool
        {
            let mut pool_guard = pool.write();
            for _ in 0..initial_size {
                pool_guard.push_back(Box::new(factory()));
            }
        }

        Self {
            pool,
            factory,
            max_pool_size: max_size,
        }
    }

    /// Get an object from the pool or create a new one.
    pub fn acquire(&self) -> PooledObject<T> {
        let obj = {
            let mut pool = self.pool.write();
            pool.pop_front()
                .unwrap_or_else(|| Box::new((self.factory)()))
        };

        PooledObject {
            object: Some(obj),
            pool: Arc::clone(&self.pool),
            max_pool_size: self.max_pool_size,
        }
    }

    /// Get current pool statistics.
    pub fn stats(&self) -> PoolStats {
        let pool = self.pool.read();
        PoolStats {
            available_objects: pool.len(),
            max_pool_size: self.max_pool_size,
        }
    }
}

/// Statistics for the safe pool allocator.
#[derive(Debug, Clone)]
pub struct PoolStats {
    /// Number of objects currently available in the pool for reuse
    pub available_objects: usize,
    /// Maximum number of objects the pool can hold
    pub max_pool_size: usize,
}

/// A pooled object that automatically returns to the pool when dropped.
pub struct PooledObject<T> {
    object: Option<Box<T>>,
    pool: Arc<RwLock<VecDeque<Box<T>>>>,
    max_pool_size: usize,
}

impl<T> PooledObject<T> {
    /// Get a reference to the pooled object.
    pub fn get(&self) -> &T {
        self.object.as_ref().unwrap()
    }

    /// Get a mutable reference to the pooled object.
    pub fn get_mut(&mut self) -> &mut T {
        self.object.as_mut().unwrap()
    }
}

impl<T> Drop for PooledObject<T> {
    fn drop(&mut self) {
        if let Some(obj) = self.object.take() {
            let mut pool = self.pool.write();
            if pool.len() < self.max_pool_size {
                pool.push_back(obj);
            }
            // If pool is full, just drop the object
        }
    }
}

/// A safe memory manager that combines arena and pool allocation.
pub struct SafeMemoryManager {
    /// Arena allocator for temporary allocations
    arena: SafeArenaAllocator,
    /// Pool for graph nodes
    node_pool: SafePoolAllocator<Vec<u8>>,
    /// Pool for relationship data
    relationship_pool: SafePoolAllocator<Vec<u8>>,
}

impl SafeMemoryManager {
    /// Create a new safe memory manager.
    pub fn new(arena_size: usize) -> Self {
        Self {
            arena: SafeArenaAllocator::new(arena_size),
            node_pool: SafePoolAllocator::new(1000, 10000),
            relationship_pool: SafePoolAllocator::new(2000, 20000),
        }
    }

    /// Get the arena allocator.
    pub fn arena(&self) -> &SafeArenaAllocator {
        &self.arena
    }

    /// Acquire a pooled buffer for node data.
    pub fn acquire_node_buffer(&self) -> PooledObject<Vec<u8>> {
        self.node_pool.acquire()
    }

    /// Acquire a pooled buffer for relationship data.
    pub fn acquire_relationship_buffer(&self) -> PooledObject<Vec<u8>> {
        self.relationship_pool.acquire()
    }

    /// Get memory usage statistics.
    pub fn stats(&self) -> SafeMemoryStats {
        SafeMemoryStats {
            arena_usage: self.arena.memory_usage(),
            node_pool: self.node_pool.stats(),
            relationship_pool: self.relationship_pool.stats(),
        }
    }

    /// Perform memory cleanup.
    pub fn cleanup(&self) {
        self.arena.reset_and_swap();
    }
}

/// Statistics for the safe memory manager.
#[derive(Debug)]
pub struct SafeMemoryStats {
    /// Current memory usage of the arena allocator in bytes
    pub arena_usage: usize,
    /// Statistics for the node data pool allocator
    pub node_pool: PoolStats,
    /// Statistics for the relationship data pool allocator
    pub relationship_pool: PoolStats,
}

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

    #[test]
    fn test_safe_arena_allocator() {
        let arena = SafeArenaAllocator::new(1024);

        // Test scoped allocation
        let result = arena.scoped(|| {
            // Simulate some work that would use arena allocation
            42
        });

        assert_eq!(result, 42);
        // Arena should be reset after scope
    }

    #[test]
    fn test_safe_pool_allocator() {
        let pool: SafePoolAllocator<Vec<u8>> = SafePoolAllocator::new(5, 20);

        // Acquire some objects
        let mut obj1 = pool.acquire();
        let mut obj2 = pool.acquire();

        // Use the objects
        obj1.get_mut().push(1);
        obj2.get_mut().push(2);

        assert_eq!(obj1.get()[0], 1);
        assert_eq!(obj2.get()[0], 2);

        // Objects are automatically returned to pool when dropped
        drop(obj1);
        drop(obj2);

        let stats = pool.stats();
        assert!(stats.available_objects > 0);
    }

    #[test]
    fn test_pooled_object_return() {
        let pool: SafePoolAllocator<String> = SafePoolAllocator::new(1, 5);

        {
            let mut obj = pool.acquire();
            obj.get_mut().push_str("test");
            assert_eq!(obj.get(), "test");
        } // Object should return to pool here

        let stats = pool.stats();
        assert_eq!(stats.available_objects, 1);
    }

    #[test]
    fn test_safe_memory_manager() {
        let manager = SafeMemoryManager::new(4096);

        // Test node buffer acquisition
        let mut node_buf = manager.acquire_node_buffer();
        node_buf.get_mut().extend_from_slice(&[1, 2, 3, 4]);

        // Test relationship buffer acquisition
        let mut rel_buf = manager.acquire_relationship_buffer();
        rel_buf.get_mut().extend_from_slice(&[5, 6, 7, 8]);

        // Get statistics
        let stats = manager.stats();
        assert!(stats.node_pool.available_objects < 1000); // One is in use
        assert!(stats.relationship_pool.available_objects < 2000); // One is in use

        // Cleanup
        manager.cleanup();
    }

    #[test]
    fn test_arena_memory_tracking() {
        let arena = SafeArenaAllocator::new(1024);

        // Arena pre-allocates capacity, so initial usage may be non-zero
        let initial_usage = arena.memory_usage();

        // Reset and swap - should have roughly the same capacity
        arena.reset_and_swap();
        let after_reset = arena.memory_usage();

        // After reset, the usage should be from the backup arena
        // which also pre-allocates, so we just verify it's reasonable
        assert!(after_reset <= initial_usage || after_reset > 0);
    }
}