lambdust 0.1.1

//! Collection algorithms for the parallel garbage collector
//!
//! This module implements the core collection algorithms including:
//! - Stop-the-world copying collector for young generation
//! - Concurrent mark-and-sweep for old generation  
//! - Incremental collection with write barriers
//! - Object promotion logic between generations

use crate::eval::value::Value;
use crate::runtime::gc::generation::{ObjectHeader, GenerationId, CollectionResult};
use crate::runtime::gc::parallel_gc::{SafepointCoordinator, GcStatistics};
use std::collections::{HashMap, HashSet, VecDeque};
use std::sync::{Arc, RwLock, Mutex, atomic::{AtomicBool, AtomicUsize, Ordering}};
use std::time::{Duration, Instant};
use std::thread;

/// Thread-safe wrapper for ObjectHeader pointers
/// SAFETY: These pointers are managed by the garbage collector and are only
/// accessed during stop-the-world collection phases when thread safety is guaranteed.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct GcPtr(*mut ObjectHeader);

impl GcPtr {
    /// Create a new GcPtr from a raw ObjectHeader pointer
    pub fn new(ptr: *mut ObjectHeader) -> Self {
        GcPtr(ptr)
    }
    
    /// Extract the raw pointer from GcPtr
    pub fn as_ptr(self) -> *mut ObjectHeader {
        self.0
    }
}

unsafe impl Send for GcPtr {}
unsafe impl Sync for GcPtr {}

/// Root set for garbage collection - objects that are always reachable
#[derive(Debug)]
pub struct RootSet {
    /// Global roots (e.g., global variables, stack frames)
    global_roots: Arc<RwLock<HashSet<GcPtr>>>,
    /// Thread-local roots by thread ID
    thread_roots: Arc<RwLock<HashMap<thread::ThreadId, HashSet<GcPtr>>>>,
    /// Remembered set for inter-generational pointers
    remembered_set: Arc<RwLock<HashSet<GcPtr>>>,
}

impl RootSet {
    /// Create a new root set
    pub fn new() -> Self {
        RootSet {
            global_roots: Arc::new(RwLock::new(HashSet::new())),
            thread_roots: Arc::new(RwLock::new(HashMap::new())),
            remembered_set: Arc::new(RwLock::new(HashSet::new())),
        }
    }

    /// Add a global root
    pub fn add_global_root(&self, obj: *mut ObjectHeader) -> Result<(), String> {
        let mut roots = self.global_roots.write().map_err(|_| "Failed to write global roots")?;
        roots.insert(GcPtr::new(obj));
        Ok(())
    }

    /// Remove a global root
    pub fn remove_global_root(&self, obj: *mut ObjectHeader) -> Result<(), String> {
        let mut roots = self.global_roots.write().map_err(|_| "Failed to write global roots")?;
        roots.remove(&GcPtr::new(obj));
        Ok(())
    }

    /// Add a thread-local root
    pub fn add_thread_root(&self, thread_id: thread::ThreadId, obj: *mut ObjectHeader) -> Result<(), String> {
        let mut roots = self.thread_roots.write().map_err(|_| "Failed to write thread roots")?;
        roots.entry(thread_id).or_insert_with(HashSet::new).insert(GcPtr::new(obj));
        Ok(())
    }

    /// Remove a thread-local root
    pub fn remove_thread_root(&self, thread_id: thread::ThreadId, obj: *mut ObjectHeader) -> Result<(), String> {
        let mut roots = self.thread_roots.write().map_err(|_| "Failed to write thread roots")?;
        if let Some(thread_roots) = roots.get_mut(&thread_id) {
            thread_roots.remove(&GcPtr::new(obj));
            if thread_roots.is_empty() {
                roots.remove(&thread_id);
            }
        }
        Ok(())
    }

    /// Add to remembered set (for inter-generational pointers)
    pub fn add_to_remembered_set(&self, obj: *mut ObjectHeader) -> Result<(), String> {
        let mut remembered = self.remembered_set.write().map_err(|_| "Failed to write remembered set")?;
        remembered.insert(GcPtr::new(obj));
        Ok(())
    }

    /// Get all roots for collection
    pub fn get_all_roots(&self) -> Result<Vec<*mut ObjectHeader>, String> {
        let mut all_roots = Vec::new();

        // Global roots
        let global_roots = self.global_roots.read().map_err(|_| "Failed to read global roots")?;
        all_roots.extend(global_roots.iter().map(|ptr| ptr.as_ptr()));

        // Thread-local roots
        let thread_roots = self.thread_roots.read().map_err(|_| "Failed to read thread roots")?;
        for thread_roots in thread_roots.values() {
            all_roots.extend(thread_roots.iter().map(|ptr| ptr.as_ptr()));
        }

        // Remembered set
        let remembered = self.remembered_set.read().map_err(|_| "Failed to read remembered set")?;
        all_roots.extend(remembered.iter().map(|ptr| ptr.as_ptr()));

        Ok(all_roots)
    }

    /// Clear remembered set (typically after collection)
    pub fn clear_remembered_set(&self) -> Result<(), String> {
        let mut remembered = self.remembered_set.write().map_err(|_| "Failed to write remembered set")?;
        remembered.clear();
        Ok(())
    }
}

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

/// Write barrier for tracking inter-generational pointers
#[derive(Debug)]
pub struct WriteBarrier {
    /// Dirty cards for tracking modified memory regions
    dirty_cards: Arc<RwLock<HashSet<usize>>>,
    /// Card size (bytes covered by each card)
    card_size: usize,
    /// Whether write barrier is active
    active: AtomicBool,
}

impl WriteBarrier {
    /// Create a new write barrier
    pub fn new(card_size: usize) -> Self {
        WriteBarrier {
            dirty_cards: Arc::new(RwLock::new(HashSet::new())),
            card_size,
            active: AtomicBool::new(true),
        }
    }

    /// Record a write operation (called by mutator)
    pub fn record_write(&self, address: *mut ObjectHeader) -> Result<(), String> {
        if !self.active.load(Ordering::Relaxed) {
            return Ok(());
        }

        let card = (address as usize) / self.card_size;
        let mut dirty_cards = self.dirty_cards.write().map_err(|_| "Failed to write dirty cards")?;
        dirty_cards.insert(card);
        Ok(())
    }

    /// Get and clear dirty cards
    pub fn get_and_clear_dirty_cards(&self) -> Result<Vec<usize>, String> {
        let mut dirty_cards = self.dirty_cards.write().map_err(|_| "Failed to write dirty cards")?;
        let cards: Vec<usize> = dirty_cards.iter().cloned().collect();
        dirty_cards.clear();
        Ok(cards)
    }

    /// Enable/disable write barrier
    pub fn set_active(&self, active: bool) {
        self.active.store(active, Ordering::Relaxed);
    }

    /// Check if write barrier is active
    pub fn is_active(&self) -> bool {
        self.active.load(Ordering::Relaxed)
    }
}

/// Object marker for mark-and-sweep collection
#[derive(Debug)]
pub struct ObjectMarker {
    /// Objects marked as live
    marked_objects: Arc<RwLock<HashSet<GcPtr>>>,
    /// Work queue for marking
    mark_queue: Arc<Mutex<VecDeque<GcPtr>>>,
    /// Whether marking is complete
    marking_complete: AtomicBool,
}

impl ObjectMarker {
    /// Create a new object marker
    pub fn new() -> Self {
        ObjectMarker {
            marked_objects: Arc::new(RwLock::new(HashSet::new())),
            mark_queue: Arc::new(Mutex::new(VecDeque::new())),
            marking_complete: AtomicBool::new(false),
        }
    }

    /// Start marking from roots
    pub fn mark_from_roots(&self, roots: Vec<*mut ObjectHeader>) -> Result<(), String> {
        self.marking_complete.store(false, Ordering::Relaxed);
        
        // Clear previous marking state
        {
            let mut marked = self.marked_objects.write().map_err(|_| "Failed to write marked objects")?;
            marked.clear();
        }
        {
            let mut queue = self.mark_queue.lock().map_err(|_| "Failed to lock mark queue")?;
            queue.clear();
        }

        // Add roots to work queue
        {
            let mut queue = self.mark_queue.lock().map_err(|_| "Failed to lock mark queue")?;
            for root in roots {
                if !root.is_null() {
                    queue.push_back(GcPtr::new(root));
                }
            }
        }

        // Process marking
        self.process_marking()?;
        self.marking_complete.store(true, Ordering::Relaxed);
        Ok(())
    }

    /// Process the marking queue
    fn process_marking(&self) -> Result<(), String> {
        loop {
            let obj = {
                let mut queue = self.mark_queue.lock().map_err(|_| "Failed to lock mark queue")?;
                queue.pop_front()
            };

            match obj {
                Some(obj_ptr) => {
                    let raw_ptr = obj_ptr.as_ptr();
                    if !raw_ptr.is_null() && self.mark_object(raw_ptr)? {
                        // Object was newly marked, scan its references
                        self.scan_object_references(raw_ptr)?;
                    }
                }
                None => break, // Queue is empty
            }
        }
        Ok(())
    }

    /// Mark a single object as live
    fn mark_object(&self, obj: *mut ObjectHeader) -> Result<bool, String> {
        let mut marked = self.marked_objects.write().map_err(|_| "Failed to write marked objects")?;
        let was_new = marked.insert(GcPtr::new(obj));
        
        // Also set the mark bit in the object header
        unsafe {
            if !obj.is_null() {
                (*obj).mark();
            }
        }
        
        Ok(was_new)
    }

    /// Scan object references and add them to the marking queue
    fn scan_object_references(&self, obj: *mut ObjectHeader) -> Result<(), String> {
        // In a real implementation, this would scan the object's references
        // based on its type and add them to the mark queue.
        // For now, we'll simulate with a simplified approach.
        
        unsafe {
            if obj.is_null() {
                return Ok(());
            }

            let header = &*obj;
            
            // Based on the value type, scan for references
            match header.value.as_ref() {
                Value::Pair(car, cdr) => {
                    // For pairs, we would need to mark both car and cdr
                    // This is simplified - in reality we'd need proper object scanning
                    let mut queue = self.mark_queue.lock().map_err(|_| "Failed to lock mark queue")?;
                    
                    // Note: This is a simplified example - in a real GC, 
                    // we'd need proper pointer discovery mechanisms
                    drop(queue);
                }
                Value::Vector(vec) => {
                    // For vectors, mark all contained values
                    let _vec_guard = vec.read().map_err(|_| "Failed to read vector")?;
                    // Similar scanning logic would go here
                }
                _ => {
                    // Other types may not contain references or need different handling
                }
            }
        }

        Ok(())
    }

    /// Check if an object is marked
    pub fn is_marked(&self, obj: *mut ObjectHeader) -> bool {
        if let Ok(marked) = self.marked_objects.read() {
            marked.contains(&GcPtr::new(obj))
        } else {
            false
        }
    }

    /// Get all marked objects
    pub fn get_marked_objects(&self) -> Result<Vec<*mut ObjectHeader>, String> {
        let marked = self.marked_objects.read().map_err(|_| "Failed to read marked objects")?;
        Ok(marked.iter().map(|ptr| ptr.as_ptr()).collect())
    }

    /// Reset marker state
    pub fn reset(&self) -> Result<(), String> {
        {
            let mut marked = self.marked_objects.write().map_err(|_| "Failed to write marked objects")?;
            marked.clear();
        }
        {
            let mut queue = self.mark_queue.lock().map_err(|_| "Failed to lock mark queue")?;
            queue.clear();
        }
        self.marking_complete.store(false, Ordering::Relaxed);
        Ok(())
    }

    /// Check if marking is complete
    pub fn is_marking_complete(&self) -> bool {
        self.marking_complete.load(Ordering::Relaxed)
    }
}

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

/// Stop-the-world copying collector for young generation
#[derive(Debug)]
pub struct CopyingCollector {
    /// Safepoint coordinator for stopping threads
    safepoint: Arc<SafepointCoordinator>,
    /// Root set for collection
    root_set: Arc<RootSet>,
    /// Collection statistics
    statistics: Arc<GcStatistics>,
}

impl CopyingCollector {
    /// Create a new copying collector
    pub fn new(
        safepoint: Arc<SafepointCoordinator>,
        root_set: Arc<RootSet>,
        statistics: Arc<GcStatistics>,
    ) -> Self {
        CopyingCollector {
            safepoint,
            root_set,
            statistics,
        }
    }

    /// Perform a stop-the-world copying collection
    pub fn collect(&self) -> Result<CollectionResult, String> {
        let start_time = Instant::now();
        
        // Request safepoint - stop all mutator threads
        self.safepoint.request_safepoint()?;
        
        let result = self.perform_copying_collection();
        
        // Release safepoint - allow threads to continue
        self.safepoint.release_safepoint();
        
        let collection_time = start_time.elapsed();
        self.statistics.record_minor_collection(collection_time);
        
        match result {
            Ok(mut result) => {
                result.collection_time = collection_time;
                Ok(result)
            }
            Err(e) => Err(e),
        }
    }

    /// Perform the actual copying collection
    fn perform_copying_collection(&self) -> Result<CollectionResult, String> {
        // Get all roots
        let roots = self.root_set.get_all_roots()?;
        
        // In a real copying collector, we would:
        // 1. Scan from roots to identify live objects
        // 2. Copy live objects to to-space
        // 3. Update all pointers to point to new locations
        // 4. Promote objects that are old enough to old generation
        // 5. Flip from-space and to-space
        
        // For now, simulate the collection
        let objects_collected = 100; // Simulated
        let bytes_reclaimed = 8192;  // Simulated
        let objects_promoted = 10;   // Simulated
        let bytes_promoted = 1024;   // Simulated

        Ok(CollectionResult {
            objects_collected,
            bytes_reclaimed,
            objects_promoted,
            bytes_promoted,
            collection_time: Duration::from_millis(0), // Will be set by caller
            success: true,
        })
    }

    /// Copy a single object to to-space
    fn copy_object(&self, _obj: *mut ObjectHeader) -> Result<*mut ObjectHeader, String> {
        // In a real implementation, this would:
        // 1. Allocate space in to-space
        // 2. Copy object data
        // 3. Set forwarding pointer in original object
        // 4. Return new location
        
        // For now, just return the original pointer (no-op)
        Ok(_obj)
    }

    /// Update a pointer to its new location after copying
    fn update_pointer(&self, _ptr: *mut *mut ObjectHeader) -> Result<(), String> {
        // In a real implementation, this would:
        // 1. Check if the object has been moved (has forwarding pointer)
        // 2. Update the pointer to the new location
        
        // For now, no-op
        Ok(())
    }
}

/// Concurrent mark-and-sweep collector for old generation
#[derive(Debug)]
pub struct MarkSweepCollector {
    /// Root set for collection
    root_set: Arc<RootSet>,
    /// Object marker
    marker: Arc<ObjectMarker>,
    /// Write barrier for concurrent collection
    write_barrier: Arc<WriteBarrier>,
    /// Collection statistics
    statistics: Arc<GcStatistics>,
    /// Whether concurrent collection is enabled
    concurrent_enabled: AtomicBool,
}

impl MarkSweepCollector {
    /// Create a new mark-and-sweep collector
    pub fn new(
        root_set: Arc<RootSet>,
        statistics: Arc<GcStatistics>,
    ) -> Self {
        let marker = Arc::new(ObjectMarker::new());
        let write_barrier = Arc::new(WriteBarrier::new(4096)); // 4KB cards
        
        MarkSweepCollector {
            root_set,
            marker,
            write_barrier,
            statistics,
            concurrent_enabled: AtomicBool::new(true),
        }
    }

    /// Perform a mark-and-sweep collection
    pub fn collect(&self, concurrent: bool) -> Result<CollectionResult, String> {
        let start_time = Instant::now();
        
        if concurrent && self.concurrent_enabled.load(Ordering::Relaxed) {
            self.collect_concurrent()
        } else {
            self.collect_stop_the_world()
        }.map(|mut result| {
            let collection_time = start_time.elapsed();
            self.statistics.record_major_collection(collection_time);
            result.collection_time = collection_time;
            result
        })
    }

    /// Perform concurrent mark-and-sweep collection
    fn collect_concurrent(&self) -> Result<CollectionResult, String> {
        // Phase 1: Initial mark (stop-the-world)
        // This phase marks objects directly reachable from roots
        let roots = self.root_set.get_all_roots()?;
        
        // Enable write barrier for concurrent phase
        self.write_barrier.set_active(true);
        
        // Phase 2: Concurrent mark
        // Mark all reachable objects while mutators are running
        self.marker.mark_from_roots(roots)?;
        
        // Phase 3: Final mark (stop-the-world)
        // Process objects modified during concurrent phase
        let dirty_cards = self.write_barrier.get_and_clear_dirty_cards()?;
        self.process_dirty_cards(dirty_cards)?;
        
        // Phase 4: Concurrent sweep
        // Deallocate unmarked objects
        let sweep_result = self.sweep_unmarked_objects()?;
        
        // Disable write barrier
        self.write_barrier.set_active(false);
        
        Ok(sweep_result)
    }

    /// Perform stop-the-world mark-and-sweep collection
    fn collect_stop_the_world(&self) -> Result<CollectionResult, String> {
        // Get all roots and mark from them
        let roots = self.root_set.get_all_roots()?;
        self.marker.mark_from_roots(roots)?;
        
        // Sweep unmarked objects
        self.sweep_unmarked_objects()
    }

    /// Process dirty cards from concurrent marking phase
    fn process_dirty_cards(&self, _dirty_cards: Vec<usize>) -> Result<(), String> {
        // In a real implementation, this would:
        // 1. For each dirty card, scan objects in that memory region
        // 2. Mark any newly discovered objects
        // 3. Process transitively until no new objects are found
        
        // For now, no-op
        Ok(())
    }

    /// Sweep unmarked objects and deallocate them
    fn sweep_unmarked_objects(&self) -> Result<CollectionResult, String> {
        // In a real implementation, this would:
        // 1. Scan through all objects in old generation
        // 2. Deallocate objects that are not marked
        // 3. Reset mark bits on live objects
        // 4. Update free space tracking
        
        // For now, simulate sweep
        let objects_collected = 50;   // Simulated
        let bytes_reclaimed = 16384;  // Simulated

        Ok(CollectionResult {
            objects_collected,
            bytes_reclaimed,
            objects_promoted: 0, // No promotion in old gen collection
            bytes_promoted: 0,
            collection_time: Duration::from_millis(0), // Will be set by caller
            success: true,
        })
    }

    /// Enable/disable concurrent collection
    pub fn set_concurrent_enabled(&self, enabled: bool) {
        self.concurrent_enabled.store(enabled, Ordering::Relaxed);
    }

    /// Get the write barrier
    pub fn get_write_barrier(&self) -> Arc<WriteBarrier> {
        Arc::clone(&self.write_barrier)
    }

    /// Reset collector state
    pub fn reset(&self) -> Result<(), String> {
        self.marker.reset()?;
        self.write_barrier.get_and_clear_dirty_cards()?;
        Ok(())
    }
}

/// Incremental collector that performs collection in small steps
#[derive(Debug)]
pub struct IncrementalCollector {
    /// Copying collector for young generation
    copying_collector: Arc<CopyingCollector>,
    /// Mark-sweep collector for old generation
    mark_sweep_collector: Arc<MarkSweepCollector>,
    /// Current collection state
    state: Arc<RwLock<IncrementalState>>,
    /// Work budget per step (in microseconds)
    step_budget_us: usize,
}

/// State of incremental collection
#[derive(Debug, Clone)]
enum IncrementalState {
    /// No collection in progress
    Idle,
    /// Marking phase in progress
    Marking { progress: f64 },
    /// Sweeping phase in progress
    Sweeping { progress: f64 },
    /// Collection complete, finalizing
    Finalizing,
}

impl IncrementalCollector {
    /// Create a new incremental collector
    pub fn new(
        copying_collector: Arc<CopyingCollector>,
        mark_sweep_collector: Arc<MarkSweepCollector>,
        step_budget_us: usize,
    ) -> Self {
        IncrementalCollector {
            copying_collector,
            mark_sweep_collector,
            state: Arc::new(RwLock::new(IncrementalState::Idle)),
            step_budget_us,
        }
    }

    /// Perform one incremental collection step
    pub fn perform_incremental_step(&self) -> Result<bool, String> {
        let start_time = Instant::now();
        let budget = Duration::from_micros(self.step_budget_us as u64);

        let current_state = {
            let state = self.state.read().map_err(|_| "Failed to read incremental state")?;
            state.clone()
        };

        match current_state {
            IncrementalState::Idle => {
                // Start a new incremental collection
                self.start_incremental_collection()?;
                Ok(false) // Not complete yet
            }
            IncrementalState::Marking { progress } => {
                // Continue marking phase
                let new_progress = self.perform_marking_step(progress, budget)?;
                
                if new_progress >= 1.0 {
                    // Marking complete, move to sweeping
                    let mut state = self.state.write().map_err(|_| "Failed to write incremental state")?;
                    *state = IncrementalState::Sweeping { progress: 0.0 };
                    Ok(false)
                } else {
                    // Update progress
                    let mut state = self.state.write().map_err(|_| "Failed to write incremental state")?;
                    *state = IncrementalState::Marking { progress: new_progress };
                    Ok(false)
                }
            }
            IncrementalState::Sweeping { progress } => {
                // Continue sweeping phase
                let new_progress = self.perform_sweeping_step(progress, budget)?;
                
                if new_progress >= 1.0 {
                    // Sweeping complete, finalize
                    let mut state = self.state.write().map_err(|_| "Failed to write incremental state")?;
                    *state = IncrementalState::Finalizing;
                    Ok(false)
                } else {
                    // Update progress
                    let mut state = self.state.write().map_err(|_| "Failed to write incremental state")?;
                    *state = IncrementalState::Sweeping { progress: new_progress };
                    Ok(false)
                }
            }
            IncrementalState::Finalizing => {
                // Finalize collection
                self.finalize_incremental_collection()?;
                let mut state = self.state.write().map_err(|_| "Failed to write incremental state")?;
                *state = IncrementalState::Idle;
                Ok(true) // Collection complete
            }
        }
    }

    /// Start a new incremental collection
    fn start_incremental_collection(&self) -> Result<(), String> {
        let mut state = self.state.write().map_err(|_| "Failed to write incremental state")?;
        *state = IncrementalState::Marking { progress: 0.0 };
        Ok(())
    }

    /// Perform one step of the marking phase
    fn perform_marking_step(&self, current_progress: f64, _budget: Duration) -> Result<f64, String> {
        // In a real implementation, this would:
        // 1. Mark objects for a limited time budget
        // 2. Track progress through the object graph
        // 3. Return updated progress percentage
        
        // For now, simulate progress
        let progress_increment = 0.1; // 10% per step
        Ok((current_progress + progress_increment).min(1.0))
    }

    /// Perform one step of the sweeping phase
    fn perform_sweeping_step(&self, current_progress: f64, _budget: Duration) -> Result<f64, String> {
        // In a real implementation, this would:
        // 1. Sweep unmarked objects for a limited time budget
        // 2. Track progress through memory regions
        // 3. Return updated progress percentage
        
        // For now, simulate progress
        let progress_increment = 0.2; // 20% per step
        Ok((current_progress + progress_increment).min(1.0))
    }

    /// Finalize the incremental collection
    fn finalize_incremental_collection(&self) -> Result<(), String> {
        // In a real implementation, this would:
        // 1. Reset object mark bits
        // 2. Update heap statistics
        // 3. Clear collection state
        
        // For now, no-op
        Ok(())
    }

    /// Check if incremental collection is in progress
    pub fn is_collection_in_progress(&self) -> bool {
        if let Ok(state) = self.state.read() {
            !matches!(*state, IncrementalState::Idle)
        } else {
            false
        }
    }

    /// Get current collection progress (0.0 to 1.0)
    pub fn get_collection_progress(&self) -> f64 {
        if let Ok(state) = self.state.read() {
            match *state {
                IncrementalState::Idle => 0.0,
                IncrementalState::Marking { progress } => progress * 0.5, // Marking is first half
                IncrementalState::Sweeping { progress } => 0.5 + progress * 0.5, // Sweeping is second half
                IncrementalState::Finalizing => 1.0,
            }
        } else {
            0.0
        }
    }

    /// Force completion of current incremental collection
    pub fn force_complete(&self) -> Result<(), String> {
        while !self.perform_incremental_step()? {
            // Continue until complete
        }
        Ok(())
    }
}