graph_d 1.3.2

//! Memory pools for efficient allocation of graph elements.
//!
//! This module provides specialized memory pools for different types of graph data,
//! optimizing allocation patterns and reducing memory fragmentation.
//!
//! Note: This module uses unsafe code for performance-critical allocations.
//! For safe alternatives, see the allocator_safe module.

use crate::error::{GraphError, Result};
use parking_lot::RwLock;
use std::collections::VecDeque;
use std::ptr::NonNull;
use std::sync::Arc;

/// A memory pool for fixed-size allocations.
///
/// This pool maintains a free list of pre-allocated blocks, providing O(1) allocation
/// and deallocation for objects of a specific size.
pub struct FixedSizePool {
    /// Free blocks available for allocation
    free_blocks: Arc<RwLock<VecDeque<NonNull<u8>>>>,
    /// Size of each block in bytes
    block_size: usize,
    /// Alignment requirement for blocks
    alignment: usize,
    /// Total number of blocks allocated by this pool
    total_blocks: Arc<RwLock<usize>>,
    /// Maximum number of blocks this pool can allocate
    max_blocks: usize,
}

// SAFETY: FixedSizePool uses Arc<RwLock<_>> for all mutable state,
// which provides thread-safe access.
#[allow(unsafe_code)]
unsafe impl Send for FixedSizePool {}
#[allow(unsafe_code)]
unsafe impl Sync for FixedSizePool {}

impl FixedSizePool {
    /// Create a new fixed-size pool.
    #[allow(clippy::arc_with_non_send_sync)]
    pub fn new(
        block_size: usize,
        alignment: usize,
        initial_blocks: usize,
        max_blocks: usize,
    ) -> Result<Self> {
        if block_size == 0 || alignment == 0 || !alignment.is_power_of_two() {
            return Err(GraphError::Memory("Invalid pool parameters".to_string()));
        }

        let pool = Self {
            free_blocks: Arc::new(RwLock::new(VecDeque::new())),
            block_size,
            alignment,
            total_blocks: Arc::new(RwLock::new(0)),
            max_blocks,
        };

        // Pre-allocate initial blocks
        pool.expand_pool(initial_blocks)?;

        Ok(pool)
    }

    /// Allocate a block from the pool.
    pub fn allocate(&self) -> Result<NonNull<u8>> {
        // Try to get a block from the free list
        {
            let mut free_blocks = self.free_blocks.write();
            if let Some(block) = free_blocks.pop_front() {
                return Ok(block);
            }
        }

        // No free blocks available, expand the pool if possible
        let total_blocks = *self.total_blocks.read();
        if total_blocks < self.max_blocks {
            self.expand_pool(std::cmp::min(64, self.max_blocks - total_blocks))?;

            // Try again after expansion
            let mut free_blocks = self.free_blocks.write();
            if let Some(block) = free_blocks.pop_front() {
                return Ok(block);
            }
        }

        Err(GraphError::Memory("Pool exhausted".to_string()))
    }

    /// Return a block to the pool.
    #[allow(unsafe_code)]
    pub fn deallocate(&self, ptr: NonNull<u8>) {
        // SAFETY: Zero out the memory for security using write_bytes.
        unsafe {
            std::ptr::write_bytes(ptr.as_ptr(), 0, self.block_size);
        }

        let mut free_blocks = self.free_blocks.write();
        free_blocks.push_back(ptr);
    }

    /// Expand the pool by allocating more blocks.
    #[allow(unsafe_code)]
    fn expand_pool(&self, count: usize) -> Result<()> {
        let layout = std::alloc::Layout::from_size_align(self.block_size, self.alignment)
            .map_err(|e| GraphError::Memory(format!("Invalid layout: {}", e)))?;

        let mut free_blocks = self.free_blocks.write();
        let mut total_blocks = self.total_blocks.write();

        for _ in 0..count {
            if *total_blocks >= self.max_blocks {
                break;
            }

            unsafe {
                let ptr = std::alloc::alloc(layout);
                if ptr.is_null() {
                    return Err(GraphError::Memory("Failed to allocate memory".to_string()));
                }

                free_blocks.push_back(NonNull::new_unchecked(ptr));
                *total_blocks += 1;
            }
        }

        Ok(())
    }

    /// Get pool statistics.
    pub fn stats(&self) -> PoolStats {
        PoolStats {
            block_size: self.block_size,
            total_blocks: *self.total_blocks.read(),
            free_blocks: self.free_blocks.read().len(),
            allocated_blocks: *self.total_blocks.read() - self.free_blocks.read().len(),
            max_blocks: self.max_blocks,
        }
    }

    /// Shrink the pool by deallocating unused blocks.
    #[allow(unsafe_code)]
    pub fn shrink(&self, target_free: usize) -> Result<usize> {
        let layout = std::alloc::Layout::from_size_align(self.block_size, self.alignment)
            .map_err(|e| GraphError::Memory(format!("Invalid layout: {}", e)))?;

        let mut free_blocks = self.free_blocks.write();
        let mut total_blocks = self.total_blocks.write();
        let mut deallocated = 0;

        while free_blocks.len() > target_free && !free_blocks.is_empty() {
            if let Some(ptr) = free_blocks.pop_back() {
                unsafe {
                    std::alloc::dealloc(ptr.as_ptr(), layout);
                }
                *total_blocks -= 1;
                deallocated += 1;
            }
        }

        Ok(deallocated)
    }
}

impl Drop for FixedSizePool {
    #[allow(unsafe_code)]
    fn drop(&mut self) {
        let layout = std::alloc::Layout::from_size_align(self.block_size, self.alignment).unwrap();

        let mut free_blocks = self.free_blocks.write();
        // SAFETY: We deallocate pointers that were allocated with the same layout.
        unsafe {
            while let Some(ptr) = free_blocks.pop_front() {
                std::alloc::dealloc(ptr.as_ptr(), layout);
            }
        }
    }
}

/// Statistics for a memory pool.
#[derive(Debug, Clone)]
pub struct PoolStats {
    /// Size of each block in bytes
    pub block_size: usize,
    /// Total number of blocks allocated by the pool
    pub total_blocks: usize,
    /// Number of free blocks available for allocation
    pub free_blocks: usize,
    /// Number of currently allocated blocks
    pub allocated_blocks: usize,
    /// Maximum number of blocks this pool can allocate
    pub max_blocks: usize,
}

impl PoolStats {
    /// Calculate utilization ratio (0.0 - 1.0).
    pub fn utilization(&self) -> f64 {
        if self.total_blocks == 0 {
            0.0
        } else {
            self.allocated_blocks as f64 / self.total_blocks as f64
        }
    }

    /// Check if the pool is nearly full.
    pub fn is_nearly_full(&self) -> bool {
        self.utilization() > 0.9
    }

    /// Check if the pool has too many free blocks.
    pub fn has_excess_free(&self) -> bool {
        self.free_blocks > self.allocated_blocks && self.free_blocks > 100
    }
}

/// A multi-size pool manager for variable-sized allocations.
///
/// This manager maintains multiple fixed-size pools for different size classes,
/// automatically selecting the appropriate pool for each allocation.
pub struct MultiSizePool {
    /// Pools for different size classes
    pools: Vec<FixedSizePool>,
    /// Size class boundaries
    size_classes: Vec<usize>,
}

impl MultiSizePool {
    /// Create a new multi-size pool with standard size classes.
    pub fn new() -> Result<Self> {
        // Standard size classes: 32, 64, 128, 256, 512, 1024, 2048, 4096 bytes
        let size_classes = vec![32, 64, 128, 256, 512, 1024, 2048, 4096];
        let mut pools = Vec::new();

        for &size in &size_classes {
            let initial_blocks = match size {
                32..=128 => 1000,  // Many small objects
                129..=512 => 500,  // Medium objects
                513..=2048 => 200, // Fewer large objects
                _ => 50,           // Very few huge objects
            };

            let max_blocks = initial_blocks * 10; // Allow 10x growth
            let alignment = if size >= 8 { 8 } else { size };

            pools.push(FixedSizePool::new(
                size,
                alignment,
                initial_blocks,
                max_blocks,
            )?);
        }

        Ok(Self {
            pools,
            size_classes,
        })
    }

    /// Find the appropriate pool for a given size.
    fn find_pool_index(&self, size: usize) -> Option<usize> {
        self.size_classes
            .iter()
            .position(|&class_size| class_size >= size)
    }

    /// Allocate memory of the specified size.
    pub fn allocate(&self, size: usize) -> Result<(NonNull<u8>, usize)> {
        if let Some(pool_index) = self.find_pool_index(size) {
            let ptr = self.pools[pool_index].allocate()?;
            let actual_size = self.size_classes[pool_index];
            Ok((ptr, actual_size))
        } else {
            Err(GraphError::Memory(format!(
                "Size {} too large for pools",
                size
            )))
        }
    }

    /// Deallocate memory, returning it to the appropriate pool.
    pub fn deallocate(&self, ptr: NonNull<u8>, size: usize) -> Result<()> {
        if let Some(pool_index) = self.find_pool_index(size) {
            self.pools[pool_index].deallocate(ptr);
            Ok(())
        } else {
            Err(GraphError::Memory(format!(
                "Invalid size {} for deallocation",
                size
            )))
        }
    }

    /// Get statistics for all pools.
    pub fn all_stats(&self) -> Vec<PoolStats> {
        self.pools.iter().map(|pool| pool.stats()).collect()
    }

    /// Perform maintenance on all pools (shrink excess free blocks).
    pub fn maintain(&self) -> Result<usize> {
        let mut total_freed = 0;

        for pool in &self.pools {
            let stats = pool.stats();
            if stats.has_excess_free() {
                let target_free = std::cmp::max(10, stats.allocated_blocks / 2);
                total_freed += pool.shrink(target_free)?;
            }
        }

        Ok(total_freed)
    }
}

/// Specialized pool for graph nodes.
pub struct NodePool {
    /// Pool for small nodes (≤ 256 bytes)
    small_pool: FixedSizePool,
    /// Pool for medium nodes (257-1024 bytes)
    medium_pool: FixedSizePool,
    /// Pool for large nodes (1025-4096 bytes)
    large_pool: FixedSizePool,
    /// Fallback allocator for very large nodes
    fallback: MultiSizePool,
}

impl NodePool {
    /// Create a new node pool optimized for typical node sizes.
    pub fn new() -> Result<Self> {
        Ok(Self {
            small_pool: FixedSizePool::new(256, 8, 2000, 20000)?, // Most nodes are small
            medium_pool: FixedSizePool::new(1024, 8, 1000, 10000)?, // Some medium nodes
            large_pool: FixedSizePool::new(4096, 8, 200, 2000)?,  // Few large nodes
            fallback: MultiSizePool::new()?,
        })
    }

    /// Allocate space for a node with the given data size.
    pub fn allocate_node(&self, data_size: usize) -> Result<(NonNull<u8>, usize)> {
        if data_size <= 256 {
            let ptr = self.small_pool.allocate()?;
            Ok((ptr, 256))
        } else if data_size <= 1024 {
            let ptr = self.medium_pool.allocate()?;
            Ok((ptr, 1024))
        } else if data_size <= 4096 {
            let ptr = self.large_pool.allocate()?;
            Ok((ptr, 4096))
        } else {
            self.fallback.allocate(data_size)
        }
    }

    /// Deallocate node space.
    pub fn deallocate_node(&self, ptr: NonNull<u8>, allocated_size: usize) -> Result<()> {
        match allocated_size {
            256 => {
                self.small_pool.deallocate(ptr);
                Ok(())
            }
            1024 => {
                self.medium_pool.deallocate(ptr);
                Ok(())
            }
            4096 => {
                self.large_pool.deallocate(ptr);
                Ok(())
            }
            _ => self.fallback.deallocate(ptr, allocated_size),
        }
    }

    /// Get comprehensive statistics for the node pool.
    pub fn stats(&self) -> NodePoolStats {
        NodePoolStats {
            small: self.small_pool.stats(),
            medium: self.medium_pool.stats(),
            large: self.large_pool.stats(),
            fallback: self.fallback.all_stats(),
        }
    }
}

/// Statistics for the node pool.
#[derive(Debug)]
pub struct NodePoolStats {
    /// Statistics for small node allocations
    pub small: PoolStats,
    /// Statistics for medium node allocations
    pub medium: PoolStats,
    /// Statistics for large node allocations
    pub large: PoolStats,
    /// Statistics for fallback pool allocations
    pub fallback: Vec<PoolStats>,
}

impl NodePoolStats {
    /// Calculate total memory usage across all pools.
    pub fn total_memory_bytes(&self) -> usize {
        let main_memory = self.small.total_blocks * self.small.block_size
            + self.medium.total_blocks * self.medium.block_size
            + self.large.total_blocks * self.large.block_size;

        let fallback_memory: usize = self
            .fallback
            .iter()
            .map(|stats| stats.total_blocks * stats.block_size)
            .sum();

        main_memory + fallback_memory
    }

    /// Calculate overall utilization.
    pub fn overall_utilization(&self) -> f64 {
        let total_blocks =
            self.small.total_blocks + self.medium.total_blocks + self.large.total_blocks;
        let allocated_blocks = self.small.allocated_blocks
            + self.medium.allocated_blocks
            + self.large.allocated_blocks;

        if total_blocks == 0 {
            0.0
        } else {
            allocated_blocks as f64 / total_blocks as f64
        }
    }
}

/// Specialized pool for graph relationships.
pub struct RelationshipPool {
    /// Most relationships are small (just source, target, type)
    standard_pool: FixedSizePool,
    /// Some relationships have properties
    extended_pool: FixedSizePool,
}

impl RelationshipPool {
    /// Create a new relationship pool.
    pub fn new() -> Result<Self> {
        Ok(Self {
            standard_pool: FixedSizePool::new(128, 8, 4000, 40000)?, // Most relationships
            extended_pool: FixedSizePool::new(512, 8, 1000, 10000)?, // Relationships with properties
        })
    }

    /// Allocate space for a relationship.
    pub fn allocate_relationship(&self, has_properties: bool) -> Result<(NonNull<u8>, usize)> {
        if has_properties {
            let ptr = self.extended_pool.allocate()?;
            Ok((ptr, 512))
        } else {
            let ptr = self.standard_pool.allocate()?;
            Ok((ptr, 128))
        }
    }

    /// Deallocate relationship space.
    pub fn deallocate_relationship(&self, ptr: NonNull<u8>, allocated_size: usize) {
        match allocated_size {
            128 => self.standard_pool.deallocate(ptr),
            512 => self.extended_pool.deallocate(ptr),
            _ => panic!("Invalid relationship size: {}", allocated_size),
        }
    }

    /// Get statistics for the relationship pool.
    pub fn stats(&self) -> (PoolStats, PoolStats) {
        (self.standard_pool.stats(), self.extended_pool.stats())
    }
}

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

    #[test]
    fn test_fixed_size_pool() {
        let pool = FixedSizePool::new(64, 8, 10, 100).unwrap();

        // Allocate some blocks
        let ptr1 = pool.allocate().unwrap();
        let ptr2 = pool.allocate().unwrap();

        let stats = pool.stats();
        assert_eq!(stats.allocated_blocks, 2);
        assert_eq!(stats.free_blocks, 8);

        // Deallocate
        pool.deallocate(ptr1);
        pool.deallocate(ptr2);

        let stats = pool.stats();
        assert_eq!(stats.allocated_blocks, 0);
        assert_eq!(stats.free_blocks, 10);
    }

    #[test]
    fn test_multi_size_pool() {
        let pool = MultiSizePool::new().unwrap();

        // Allocate different sizes
        let (ptr1, size1) = pool.allocate(30).unwrap();
        let (ptr2, size2) = pool.allocate(100).unwrap();
        let (ptr3, size3) = pool.allocate(500).unwrap();

        assert_eq!(size1, 32); // Rounded up to next size class
        assert_eq!(size2, 128);
        assert_eq!(size3, 512);

        // Deallocate
        pool.deallocate(ptr1, size1).unwrap();
        pool.deallocate(ptr2, size2).unwrap();
        pool.deallocate(ptr3, size3).unwrap();
    }

    #[test]
    fn test_node_pool() {
        let pool = NodePool::new().unwrap();

        // Allocate nodes of different sizes
        let (ptr1, size1) = pool.allocate_node(100).unwrap();
        let (ptr2, size2) = pool.allocate_node(800).unwrap();
        let (ptr3, size3) = pool.allocate_node(3000).unwrap();

        assert_eq!(size1, 256);
        assert_eq!(size2, 1024);
        assert_eq!(size3, 4096);

        // Deallocate
        pool.deallocate_node(ptr1, size1).unwrap();
        pool.deallocate_node(ptr2, size2).unwrap();
        pool.deallocate_node(ptr3, size3).unwrap();

        let stats = pool.stats();
        assert!(stats.overall_utilization() == 0.0);
    }

    #[test]
    fn test_relationship_pool() {
        let pool = RelationshipPool::new().unwrap();

        // Allocate relationships
        let (ptr1, size1) = pool.allocate_relationship(false).unwrap();
        let (ptr2, size2) = pool.allocate_relationship(true).unwrap();

        assert_eq!(size1, 128); // Standard relationship
        assert_eq!(size2, 512); // Relationship with properties

        // Deallocate
        pool.deallocate_relationship(ptr1, size1);
        pool.deallocate_relationship(ptr2, size2);

        let (standard_stats, extended_stats) = pool.stats();
        assert_eq!(standard_stats.allocated_blocks, 0);
        assert_eq!(extended_stats.allocated_blocks, 0);
    }

    #[test]
    fn test_pool_maintenance() {
        let pool = MultiSizePool::new().unwrap();

        // Allocate and then deallocate many blocks to create excess free blocks
        let mut ptrs = Vec::new();
        for _ in 0..100 {
            let (ptr, size) = pool.allocate(64).unwrap();
            ptrs.push((ptr, size));
        }

        // Deallocate all
        for (ptr, size) in ptrs {
            pool.deallocate(ptr, size).unwrap();
        }

        // Perform maintenance
        let freed = pool.maintain().unwrap();
        assert!(freed > 0);
    }
}