graph_d 1.3.2

//! Batch allocation support for high-performance operations.
//!
//! Satisfies: RT-5 (Memory usage MUST stay within bounds)
//! Satisfies: B3 (Memory usage ≤1GB for 1M documents + 4M edges)
//! Implements: TN2 resolution (Use bumpalo arena allocator for batch allocations)
//! Phase: D (Performance Validation)
//!
//! This module extends the existing memory management with batch-specific
//! optimizations for bulk node/relationship operations.

use crate::error::{GraphError, Result};
use bumpalo::Bump;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;

/// Batch operation memory tracker.
/// Satisfies: RT-5 (track memory during batch operations)
#[derive(Debug)]
pub struct BatchMemoryTracker {
    /// Current allocated bytes in batch operations
    allocated: AtomicU64,

    /// Peak allocated bytes
    peak: AtomicU64,

    /// Number of batch operations
    batch_count: AtomicU64,

    /// Memory limit for batch operations
    limit_bytes: u64,
}

impl BatchMemoryTracker {
    /// Create a new batch memory tracker.
    pub fn new(limit_bytes: u64) -> Self {
        BatchMemoryTracker {
            allocated: AtomicU64::new(0),
            peak: AtomicU64::new(0),
            batch_count: AtomicU64::new(0),
            limit_bytes,
        }
    }

    /// Record bytes allocated in a batch.
    pub fn record_alloc(&self, bytes: u64) -> Result<()> {
        let current = self.allocated.fetch_add(bytes, Ordering::SeqCst) + bytes;

        if current > self.limit_bytes {
            self.allocated.fetch_sub(bytes, Ordering::SeqCst);
            return Err(GraphError::Memory(format!(
                "Batch memory limit exceeded: {} > {} bytes",
                current, self.limit_bytes
            )));
        }

        // Update peak
        let mut peak = self.peak.load(Ordering::SeqCst);
        while current > peak {
            match self.peak.compare_exchange_weak(
                peak,
                current,
                Ordering::SeqCst,
                Ordering::Relaxed,
            ) {
                Ok(_) => break,
                Err(p) => peak = p,
            }
        }

        Ok(())
    }

    /// Record bytes deallocated.
    pub fn record_dealloc(&self, bytes: u64) {
        self.allocated.fetch_sub(bytes, Ordering::SeqCst);
    }

    /// Record a batch operation completed.
    pub fn record_batch_complete(&self) {
        self.batch_count.fetch_add(1, Ordering::SeqCst);
    }

    /// Get current usage.
    pub fn current_usage(&self) -> u64 {
        self.allocated.load(Ordering::SeqCst)
    }

    /// Get peak usage.
    pub fn peak_usage(&self) -> u64 {
        self.peak.load(Ordering::SeqCst)
    }

    /// Get batch count.
    pub fn batch_count(&self) -> u64 {
        self.batch_count.load(Ordering::SeqCst)
    }

    /// Get usage percentage.
    pub fn usage_percent(&self) -> f64 {
        (self.current_usage() as f64 / self.limit_bytes as f64) * 100.0
    }
}

impl Default for BatchMemoryTracker {
    fn default() -> Self {
        // Default to 256MB for batch operations
        Self::new(256 * 1024 * 1024)
    }
}

/// Batch context for a single batch operation.
/// Uses bumpalo for fast allocation with bulk deallocation.
/// Satisfies: TN2 (arena allocator for batch operations)
pub struct BatchContext {
    /// Arena for this batch
    arena: Bump,

    /// Bytes allocated in this batch
    allocated_bytes: u64,

    /// Reference to the tracker
    tracker: Arc<BatchMemoryTracker>,
}

impl BatchContext {
    /// Create a new batch context.
    pub fn new(tracker: Arc<BatchMemoryTracker>) -> Self {
        BatchContext {
            arena: Bump::new(),
            allocated_bytes: 0,
            tracker,
        }
    }

    /// Allocate memory in this batch.
    pub fn alloc<T>(&mut self, value: T) -> Result<&mut T> {
        let size = std::mem::size_of::<T>() as u64;
        self.tracker.record_alloc(size)?;
        self.allocated_bytes += size;
        Ok(self.arena.alloc(value))
    }

    /// Allocate a slice in this batch.
    pub fn alloc_slice<T: Copy>(&mut self, values: &[T]) -> Result<&mut [T]> {
        let size = std::mem::size_of_val(values) as u64;
        self.tracker.record_alloc(size)?;
        self.allocated_bytes += size;
        Ok(self.arena.alloc_slice_copy(values))
    }

    /// Get bytes allocated in this batch.
    pub fn allocated_bytes(&self) -> u64 {
        self.allocated_bytes
    }

    /// Finish this batch, freeing all memory.
    pub fn finish(self) {
        self.tracker.record_dealloc(self.allocated_bytes);
        self.tracker.record_batch_complete();
        // arena is dropped here, freeing all memory
    }
}

/// Batch operation manager.
/// Satisfies: B2 (high throughput batch operations)
pub struct BatchManager {
    /// Memory tracker
    tracker: Arc<BatchMemoryTracker>,
}

impl BatchManager {
    /// Create a new batch manager.
    pub fn new(limit_bytes: u64) -> Self {
        BatchManager {
            tracker: Arc::new(BatchMemoryTracker::new(limit_bytes)),
        }
    }

    /// Start a new batch operation.
    pub fn start_batch(&self) -> BatchContext {
        BatchContext::new(self.tracker.clone())
    }

    /// Get memory usage statistics.
    pub fn get_stats(&self) -> BatchStats {
        BatchStats {
            current_usage: self.tracker.current_usage(),
            peak_usage: self.tracker.peak_usage(),
            batch_count: self.tracker.batch_count(),
            usage_percent: self.tracker.usage_percent(),
        }
    }
}

impl Default for BatchManager {
    fn default() -> Self {
        Self::new(256 * 1024 * 1024) // 256MB default
    }
}

/// Batch operation statistics.
#[derive(Debug, Clone)]
pub struct BatchStats {
    /// Current memory usage
    pub current_usage: u64,

    /// Peak memory usage
    pub peak_usage: u64,

    /// Number of completed batches
    pub batch_count: u64,

    /// Usage as percentage of limit
    pub usage_percent: f64,
}

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

    #[test]
    fn test_batch_context() {
        let manager = BatchManager::new(1024 * 1024); // 1MB limit

        let mut batch = manager.start_batch();

        // Allocate and verify first value
        {
            let val1 = batch.alloc(42u64).unwrap();
            assert_eq!(*val1, 42);
        }

        // Allocate and verify second value
        {
            let val2 = batch.alloc(100u64).unwrap();
            assert_eq!(*val2, 100);
        }

        assert!(batch.allocated_bytes() > 0);

        // Finish batch
        batch.finish();

        // Check stats
        let stats = manager.get_stats();
        assert_eq!(stats.current_usage, 0);
        assert_eq!(stats.batch_count, 1);
    }

    #[test]
    fn test_batch_memory_limit() {
        let manager = BatchManager::new(100); // 100 byte limit

        let mut batch = manager.start_batch();

        // Allocate within limit
        batch.alloc([0u8; 50]).unwrap();

        // This should fail (exceeds limit)
        let result = batch.alloc([0u8; 100]);
        assert!(result.is_err());
    }

    #[test]
    fn test_multiple_batches() {
        let manager = BatchManager::new(1024);

        // First batch
        let batch1 = manager.start_batch();
        batch1.finish();

        // Second batch
        let batch2 = manager.start_batch();
        batch2.finish();

        let stats = manager.get_stats();
        assert_eq!(stats.batch_count, 2);
    }
}