llvm-native-core-ext 0.1.0

//! Greedy Register Allocator — Enhanced register allocation with splitting,
//! eviction chains, re-materialization, and multi-way eviction.
//! Clean-room behavioral reconstruction.
//!
//! The greedy allocator improves on basic linear scan by:
//! - Region splitting: split live ranges at region boundaries for better allocation
//! - Local splitting: split within basic blocks to reduce register pressure
//! - Global splitting: split across basic blocks via new virtual registers
//! - Spill weight computation: spill cost weighted by loop depth and use count
//! - Eviction chain handling: track cascaded spills when evicting
//! - Re-materialization: recompute cheap values instead of spilling
//! - Split kit: manage split points, compute split costs, insert split code
//! - Register assignment strategies: best-fit, first-fit, hint-aware
//! - Multi-way eviction: evict multiple registers for wide values
//! - Allocation statistics: track spills, splits, evictions per function

use std::cmp::Ordering;
use std::collections::{BTreeMap, BinaryHeap, HashMap, HashSet, VecDeque};

// Re-export types from the codegen module
use llvm_native_core::codegen_regalloc::{
    InstrPoint, LinearLiveInterval, LinearLiveSegment, LinearScanAllocator, RegClassKind,
};

// ============================================================================
// Greedy Allocation Types
// ============================================================================

/// Register assignment strategy.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AssignStrategy {
    /// Assign the first free register.
    FirstFit,
    /// Assign the register that minimizes future conflicts.
    BestFit,
    /// Prefer the register hint.
    HintAware,
}

/// A split point in a live range.
#[derive(Debug, Clone, Copy)]
pub struct SplitPoint {
    /// Position in the instruction stream.
    pub point: InstrPoint,
    /// Cost of splitting at this point (lower is better).
    pub cost: f64,
    /// Whether this split avoids a spill.
    pub avoids_spill: bool,
    /// Gap size at this point (larger gaps are better split points).
    pub gap_size: u32,
}

impl SplitPoint {
    pub fn new(point: InstrPoint, cost: f64) -> Self {
        Self {
            point,
            cost,
            avoids_spill: false,
            gap_size: 0,
        }
    }

    pub fn with_gap(point: InstrPoint, cost: f64, gap_size: u32) -> Self {
        Self {
            point,
            cost,
            avoids_spill: false,
            gap_size,
        }
    }

    pub fn avoids_spill(mut self) -> Self {
        self.avoids_spill = true;
        self
    }
}

// ============================================================================
// Split Kit
// ============================================================================

/// Manages split points, computes split costs, and inserts split code.
#[derive(Debug, Clone)]
pub struct SplitKit {
    /// The original virtual register being split.
    pub original_vreg: u32,
    /// All candidate split points sorted by position.
    pub split_points: Vec<SplitPoint>,
    /// New virtual registers created during splitting.
    pub new_vregs: Vec<u32>,
    /// Split code to insert: (position, from_vreg, to_vreg).
    pub split_copies: Vec<(InstrPoint, u32, u32)>,
    /// The register class.
    pub reg_class: RegClassKind,
    /// Next virtual register ID to assign.
    next_vreg_id: u32,
}

impl SplitKit {
    pub fn new(original_vreg: u32, reg_class: RegClassKind) -> Self {
        Self {
            original_vreg,
            split_points: Vec::new(),
            new_vregs: Vec::new(),
            split_copies: Vec::new(),
            reg_class,
            next_vreg_id: original_vreg + 1000, // Start new vregs at a high offset
        }
    }

    /// Compute candidate split points from an interval.
    pub fn compute_split_points(&mut self, interval: &LinearLiveInterval) {
        self.split_points.clear();

        // Gaps between segments are good split points
        for window in interval.segments.windows(2) {
            let gap_start = window[0].end;
            let gap_end = window[1].start;
            let gap_size =
                gap_end.block * 1000 + gap_end.instr - (gap_start.block * 1000 + gap_start.instr);

            if gap_size > 0 {
                let mid_point = InstrPoint::new((gap_start.block + gap_end.block) / 2, 0, 0);
                self.split_points.push(SplitPoint::with_gap(
                    mid_point,
                    1.0 / gap_size as f64,
                    gap_size,
                ));
            }
        }

        // Use points are also potential split points
        for &use_point in &interval.use_points {
            let split_pt = use_point.before();
            if !self.split_points.iter().any(|sp| sp.point == split_pt) {
                self.split_points.push(SplitPoint::new(split_pt, 2.0));
            }
        }

        // Def points: split after definition
        for &def_point in &interval.def_points {
            let split_pt = def_point.after();
            if !self.split_points.iter().any(|sp| sp.point == split_pt) {
                self.split_points.push(SplitPoint::new(split_pt, 1.5));
            }
        }

        // Sort by position
        self.split_points.sort_by_key(|sp| sp.point);
    }

    /// Select the best split point that avoids a spill at the conflict point.
    pub fn find_best_split_before(&self, conflict_point: InstrPoint) -> Option<SplitPoint> {
        self.split_points
            .iter()
            .filter(|sp| sp.point < conflict_point)
            .max_by(|a, b| {
                a.avoids_spill
                    .cmp(&b.avoids_spill)
                    .then_with(|| b.gap_size.cmp(&a.gap_size))
                    .then_with(|| {
                        a.cost
                            .partial_cmp(&b.cost)
                            .unwrap_or(Ordering::Equal)
                            .reverse()
                    })
            })
            .copied()
    }

    /// Select the best split point after a conflict point.
    pub fn find_best_split_after(&self, conflict_point: InstrPoint) -> Option<SplitPoint> {
        self.split_points
            .iter()
            .filter(|sp| sp.point > conflict_point)
            .min_by(|a, b| {
                a.avoids_spill
                    .cmp(&b.avoids_spill)
                    .then_with(|| b.gap_size.cmp(&a.gap_size))
                    .then_with(|| a.cost.partial_cmp(&b.cost).unwrap_or(Ordering::Equal))
            })
            .copied()
    }

    /// Create a new virtual register for a split portion.
    pub fn allocate_new_vreg(&mut self) -> u32 {
        let new_id = self.next_vreg_id;
        self.next_vreg_id += 1;
        self.new_vregs.push(new_id);
        new_id
    }

    /// Add a COPY instruction at a split point.
    pub fn add_split_copy(&mut self, point: InstrPoint, from_vreg: u32, to_vreg: u32) {
        self.split_copies.push((point, from_vreg, to_vreg));
    }

    /// Whether splitting is profitable (at least one candidate point).
    pub fn is_profitable(&self) -> bool {
        !self.split_points.is_empty()
    }

    /// Get the total cost of all splits.
    pub fn total_split_cost(&self) -> f64 {
        self.split_points.iter().map(|sp| sp.cost).sum()
    }

    /// Estimate whether splitting will reduce overall spill cost.
    pub fn reduces_spill_cost(&self, spill_weight: f64) -> bool {
        let split_cost = self.total_split_cost();
        split_cost < spill_weight * 0.5
    }
}

// ============================================================================
// Region Splitting
// ============================================================================

/// Represents a region in the program (a set of basic blocks).
#[derive(Debug, Clone)]
pub struct ProgramRegion {
    /// The region ID.
    pub id: u32,
    /// Blocks in this region.
    pub blocks: Vec<u32>,
    /// Loop depth of this region.
    pub loop_depth: u32,
    /// Whether this region is a loop.
    pub is_loop: bool,
    /// Region entry blocks.
    pub entries: Vec<u32>,
    /// Region exit blocks.
    pub exits: Vec<u32>,
}

impl ProgramRegion {
    pub fn new(id: u32, loop_depth: u32) -> Self {
        Self {
            id,
            blocks: Vec::new(),
            loop_depth,
            is_loop: false,
            entries: Vec::new(),
            exits: Vec::new(),
        }
    }

    pub fn is_loop(mut self) -> Self {
        self.is_loop = true;
        self
    }

    pub fn add_block(&mut self, block: u32) {
        if !self.blocks.contains(&block) {
            self.blocks.push(block);
        }
    }

    pub fn set_entry(&mut self, block: u32) {
        if !self.entries.contains(&block) {
            self.entries.push(block);
        }
    }

    pub fn set_exit(&mut self, block: u32) {
        if !self.exits.contains(&block) {
            self.exits.push(block);
        }
    }
}

/// Region splitter: splits live ranges at region boundaries.
#[derive(Debug, Clone)]
pub struct RegionSplitter {
    /// Regions in the function.
    pub regions: Vec<ProgramRegion>,
    /// Block-to-region mapping.
    pub block_to_region: HashMap<u32, u32>,
    /// Split points at region boundaries.
    pub boundary_splits: Vec<(u32, InstrPoint)>,
}

impl RegionSplitter {
    pub fn new() -> Self {
        Self {
            regions: Vec::new(),
            block_to_region: HashMap::new(),
            boundary_splits: Vec::new(),
        }
    }

    /// Add a region.
    pub fn add_region(&mut self, region: ProgramRegion) {
        let id = region.id;
        for &block in &region.blocks {
            self.block_to_region.insert(block, id);
        }
        self.regions.push(region);
    }

    /// Find split points at region boundaries.
    pub fn compute_boundary_splits(&mut self) {
        self.boundary_splits.clear();
        for region in &self.regions {
            for &entry in &region.entries {
                let point = InstrPoint::new(entry, 0, 0);
                self.boundary_splits.push((region.id, point));
            }
        }
        self.boundary_splits.sort_by_key(|(_, pt)| *pt);
    }

    /// Check if a point is at a region boundary.
    pub fn is_boundary(&self, point: InstrPoint) -> Option<u32> {
        self.boundary_splits
            .iter()
            .find(|(_, pt)| *pt == point)
            .map(|(id, _)| *id)
    }
}

// ============================================================================
// Spill Weight Computation
// ============================================================================

/// Computes spill costs weighted by loop depth and use count.
#[derive(Debug, Clone)]
pub struct SpillWeightComputer {
    /// Loop depth for each basic block.
    pub block_loop_depths: HashMap<u32, u32>,
    /// Base spill cost for each instruction type.
    pub instruction_costs: HashMap<u32, f64>,
}

impl SpillWeightComputer {
    pub fn new() -> Self {
        let mut instruction_costs = HashMap::new();
        // Loads/stores: moderate cost
        instruction_costs.insert(0, 5.0);
        // ALU ops: low cost
        instruction_costs.insert(1, 1.0);
        // Branches: higher cost (control flow)
        instruction_costs.insert(2, 8.0);

        Self {
            block_loop_depths: HashMap::new(),
            instruction_costs,
        }
    }

    /// Set loop depth for a block.
    pub fn set_loop_depth(&mut self, block: u32, depth: u32) {
        self.block_loop_depths.insert(block, depth);
    }

    /// Get loop depth for a point.
    pub fn loop_depth_at(&self, point: InstrPoint) -> u32 {
        self.block_loop_depths
            .get(&point.block)
            .copied()
            .unwrap_or(0)
    }

    /// Compute spill weight for an interval using the formula:
    /// weight = sum(use_count * 10^loop_depth) / length
    pub fn compute_spill_weight(&self, interval: &LinearLiveInterval) -> f64 {
        let mut total_weight = 0.0f64;

        for &use_point in &interval.use_points {
            let depth = self.loop_depth_at(use_point);
            let loop_multiplier = 10.0f64.powi(depth as i32);
            total_weight += 1.0 * loop_multiplier;
        }

        for &def_point in &interval.def_points {
            let depth = self.loop_depth_at(def_point);
            let loop_multiplier = 10.0f64.powi(depth as i32);
            total_weight += 0.5 * loop_multiplier;
        }

        // Normalize by live range length
        let length = interval
            .end_point()
            .and_then(|e| {
                interval.start_point().map(|s| {
                    ((e.block - s.block) as f64 * 1000.0 + (e.instr - s.instr) as f64).max(1.0)
                })
            })
            .unwrap_or(1.0);

        total_weight / length
    }

    /// Compute spill weight for a set of intervals.
    pub fn compute_all_weights(&self, intervals: &mut HashMap<u32, LinearLiveInterval>) {
        for interval in intervals.values_mut() {
            let weight = self.compute_spill_weight(interval);
            interval.spill_weight = weight;
        }
    }
}

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

// ============================================================================
// Eviction Chain Handling
// ============================================================================

/// Tracks cascaded spills when evicting registers.
#[derive(Debug, Clone)]
pub struct EvictionChain {
    /// Ordered list of evictions: (evicted_vreg, evicting_vreg, phys_reg, point).
    pub evictions: Vec<(u32, u32, u32, InstrPoint)>,
    /// The chain depth.
    pub chain_depth: usize,
}

impl EvictionChain {
    pub fn new() -> Self {
        Self {
            evictions: Vec::new(),
            chain_depth: 0,
        }
    }

    /// Record an eviction.
    pub fn record_eviction(
        &mut self,
        evicted_vreg: u32,
        evicting_vreg: u32,
        phys_reg: u32,
        point: InstrPoint,
    ) {
        self.evictions
            .push((evicted_vreg, evicting_vreg, phys_reg, point));
        self.chain_depth = self.evictions.len();
    }

    /// Get the total cost of this eviction chain.
    pub fn total_cost(&self, intervals: &HashMap<u32, LinearLiveInterval>) -> f64 {
        self.evictions
            .iter()
            .map(|(evicted, _, _, _)| {
                intervals
                    .get(evicted)
                    .map(|i| i.spill_weight)
                    .unwrap_or(0.0)
            })
            .sum()
    }

    /// Check if a chain is too deep (may cause infinite eviction loops).
    pub fn is_too_deep(&self, max_depth: usize) -> bool {
        self.chain_depth > max_depth
    }

    /// Clear the chain.
    pub fn clear(&mut self) {
        self.evictions.clear();
        self.chain_depth = 0;
    }
}

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

// ============================================================================
// Re-materialization Detection
// ============================================================================

/// Detects values that are cheaper to recompute than spill/reload.
#[derive(Debug, Clone)]
pub struct RematerializationInfo {
    /// Virtual registers that are rematerializable.
    pub rematerializable: HashSet<u32>,
    /// The original instruction opcode for rematerialization.
    pub remat_info: HashMap<u32, RematSource>,
}

/// Source information needed to re-materialize a value.
#[derive(Debug, Clone)]
pub enum RematSource {
    /// Constant value: just load immediate.
    Constant(i64),
    /// Load from read-only global.
    GlobalLoad(String, i64),
    /// Simple computation: opcode + operands.
    Computation { opcode: u32, ops: Vec<u32> },
    /// Frame index computation.
    FrameIndex(i32),
}

impl RematerializationInfo {
    pub fn new() -> Self {
        Self {
            rematerializable: HashSet::new(),
            remat_info: HashMap::new(),
        }
    }

    /// Mark a virtual register as rematerializable.
    pub fn mark_rematerializable(&mut self, vreg: u32, source: RematSource) {
        self.rematerializable.insert(vreg);
        self.remat_info.insert(vreg, source);
    }

    /// Check if a virtual register is rematerializable.
    pub fn is_rematerializable(&self, vreg: u32) -> bool {
        self.rematerializable.contains(&vreg)
    }

    /// Get the rematerialization source.
    pub fn get_remat_source(&self, vreg: u32) -> Option<&RematSource> {
        self.remat_info.get(&vreg)
    }

    /// Detect rematerializable values from instruction analysis.
    pub fn detect_from_instruction(
        &mut self,
        vreg: u32,
        opcode: u32,
        operands: &[u32],
        is_constant: bool,
        constant_value: Option<i64>,
    ) {
        // Constants are always rematerializable
        if is_constant {
            if let Some(val) = constant_value {
                self.mark_rematerializable(vreg, RematSource::Constant(val));
                return;
            }
        }

        // Detect specific instruction patterns that are cheap
        match opcode {
            // ADDI with zero: just a register copy
            n if operands.len() == 2 && operands[0] == 0 => {
                self.mark_rematerializable(vreg, RematSource::Constant(0));
            }
            // Simple frame index computation (addi sp, offset)
            n if n == 19 /* ADDI */ && operands.len() >= 2 => {
                self.mark_rematerializable(
                    vreg,
                    RematSource::Computation {
                        opcode,
                        ops: operands.to_vec(),
                    },
                );
            }
            _ => {}
        }
    }
}

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

// ============================================================================
// Multi-Way Eviction
// ============================================================================

/// Handles eviction of multiple registers for wide values (e.g., vector registers).
#[derive(Debug, Clone)]
pub struct MultiWayEviction {
    /// The virtual register requiring multiple physical registers.
    pub vreg: u32,
    /// Number of consecutive physical registers needed.
    pub num_regs: u32,
    /// Register class.
    pub reg_class: RegClassKind,
    /// The physical registers evicted.
    pub evicted_regs: Vec<u32>,
    /// The virtual registers evicted from those physical registers.
    pub evicted_vregs: Vec<u32>,
}

impl MultiWayEviction {
    pub fn new(vreg: u32, num_regs: u32, reg_class: RegClassKind) -> Self {
        Self {
            vreg,
            num_regs,
            reg_class,
            evicted_regs: Vec::new(),
            evicted_vregs: Vec::new(),
        }
    }

    /// Check if a consecutive range of physical registers is available.
    pub fn find_consecutive_range(
        &self,
        active: &[(u32, u32)], // (phys_reg, vreg)
        available: &[u32],
    ) -> Option<u32> {
        let used: HashSet<u32> = active.iter().map(|(r, _)| *r).collect();
        for start_reg in available {
            if *start_reg + self.num_regs > 96 {
                continue;
            }
            let range_free = (0..self.num_regs).all(|i| !used.contains(&(start_reg + i)));
            if range_free {
                return Some(*start_reg);
            }
        }
        None
    }

    /// Find which registers to evict to make room for a consecutive range.
    pub fn find_registers_to_evict(&mut self, active: &[(u32, u32)], start_reg: u32) -> bool {
        self.evicted_regs.clear();
        self.evicted_vregs.clear();

        for i in 0..self.num_regs {
            let reg = start_reg + i;
            if let Some(&(_, victim_vreg)) = active.iter().find(|(r, _)| *r == reg) {
                self.evicted_regs.push(reg);
                self.evicted_vregs.push(victim_vreg);
            }
        }
        !self.evicted_regs.is_empty()
    }

    /// Total spill cost of evicting all required registers.
    pub fn total_eviction_cost(&self, intervals: &HashMap<u32, LinearLiveInterval>) -> f64 {
        self.evicted_vregs
            .iter()
            .filter_map(|v| intervals.get(v))
            .map(|i| i.spill_weight)
            .sum()
    }
}

// ============================================================================
// Allocation Statistics
// ============================================================================

/// Detailed allocation statistics for a function.
#[derive(Debug, Clone, Default)]
pub struct GreedyAllocStats {
    /// Total virtual registers.
    pub total_vregs: usize,
    /// Virtual registers successfully assigned to physical registers.
    pub assigned_vregs: usize,
    /// Virtual registers spilled.
    pub spilled_vregs: usize,
    /// Live range splits performed.
    pub splits_performed: usize,
    /// Local splits (within basic block).
    pub local_splits: usize,
    /// Global splits (across basic blocks).
    pub global_splits: usize,
    /// Register evictions performed.
    pub evictions_performed: usize,
    /// Eviction chain depth (max).
    pub max_eviction_chain_depth: usize,
    /// Total eviction chain depth (cumulative).
    pub total_eviction_depth: usize,
    /// Rematerialized values.
    pub rematerialized_count: usize,
    /// Best-fit assignments.
    pub best_fit_assignments: usize,
    /// First-fit assignments.
    pub first_fit_assignments: usize,
    /// Hint-aware assignments.
    pub hint_assignments: usize,
    /// Multi-way evictions for wide values.
    pub multi_way_evictions: usize,
    /// Copies coalesced.
    pub copies_coalesced: usize,
    /// Spill slots allocated.
    pub spill_slots_allocated: usize,
    /// Total frame size for spills.
    pub spill_frame_size: i32,
    /// Total allocation time in abstract units.
    pub allocation_time_estimate: f64,
}

impl GreedyAllocStats {
    pub fn new() -> Self {
        Self::default()
    }

    /// Merge stats from another allocator pass.
    pub fn merge(&mut self, other: &GreedyAllocStats) {
        self.total_vregs += other.total_vregs;
        self.assigned_vregs += other.assigned_vregs;
        self.spilled_vregs += other.spilled_vregs;
        self.splits_performed += other.splits_performed;
        self.local_splits += other.local_splits;
        self.global_splits += other.global_splits;
        self.evictions_performed += other.evictions_performed;
        self.max_eviction_chain_depth = self
            .max_eviction_chain_depth
            .max(other.max_eviction_chain_depth);
        self.total_eviction_depth += other.total_eviction_depth;
        self.rematerialized_count += other.rematerialized_count;
        self.best_fit_assignments += other.best_fit_assignments;
        self.first_fit_assignments += other.first_fit_assignments;
        self.hint_assignments += other.hint_assignments;
        self.multi_way_evictions += other.multi_way_evictions;
        self.copies_coalesced += other.copies_coalesced;
        self.spill_slots_allocated += other.spill_slots_allocated;
        self.spill_frame_size = self.spill_frame_size.max(other.spill_frame_size);
    }

    /// Print a summary.
    pub fn summary(&self) -> String {
        format!(
            "RA Stats: {} vregs, {} assigned, {} spilled, {} splits, {} evictions, {} coalesced",
            self.total_vregs,
            self.assigned_vregs,
            self.spilled_vregs,
            self.splits_performed,
            self.evictions_performed,
            self.copies_coalesced,
        )
    }
}

// ============================================================================
// Register Assignment Strategies
// ============================================================================

/// Strategies for selecting the best physical register from a set of free registers.
#[derive(Debug, Clone)]
pub struct RegisterAssignmentStrategy {
    /// The strategy kind.
    pub strategy: AssignStrategy,
    /// Free registers available.
    pub free_regs: Vec<u32>,
    /// Currently active intervals (for conflict checking).
    pub active_intervals: Vec<LinearLiveInterval>,
    /// Register hints per virtual register.
    pub hints: HashMap<u32, u32>,
}

impl RegisterAssignmentStrategy {
    pub fn new(strategy: AssignStrategy) -> Self {
        Self {
            strategy,
            free_regs: Vec::new(),
            active_intervals: Vec::new(),
            hints: HashMap::new(),
        }
    }

    /// Select the best register using the configured strategy.
    pub fn select_best_register(&self, interval: &LinearLiveInterval) -> Option<u32> {
        if self.free_regs.is_empty() {
            return None;
        }

        match self.strategy {
            AssignStrategy::FirstFit => self.select_first_fit(interval),
            AssignStrategy::BestFit => self.select_best_fit(interval),
            AssignStrategy::HintAware => self.select_hint_aware(interval),
        }
    }

    /// First-fit: pick the first available register.
    fn select_first_fit(&self, _interval: &LinearLiveInterval) -> Option<u32> {
        self.free_regs.first().copied()
    }

    /// Best-fit: pick the register that minimizes future conflicts.
    fn select_best_fit(&self, interval: &LinearLiveInterval) -> Option<u32> {
        let end = interval.end_point();

        self.free_regs
            .iter()
            .min_by(|&&a, &&b| {
                // Prefer register with fewer active users
                let conflicts_a = self
                    .active_intervals
                    .iter()
                    .filter(|i| i.assigned_reg == Some(a))
                    .count();
                let conflicts_b = self
                    .active_intervals
                    .iter()
                    .filter(|i| i.assigned_reg == Some(b))
                    .count();
                conflicts_a.cmp(&conflicts_b)
            })
            .copied()
    }

    /// Hint-aware: prefer the register hint if free.
    fn select_hint_aware(&self, interval: &LinearLiveInterval) -> Option<u32> {
        // Try the explicit hint first
        if let Some(hint) = interval.reg_hint {
            if self.free_regs.contains(&hint) {
                return Some(hint);
            }
        }

        // Try hints from the hints map
        if let Some(hint) = self.hints.get(&interval.vreg) {
            if self.free_regs.contains(hint) {
                return Some(*hint);
            }
        }

        // Fall back to best-fit
        self.select_best_fit(interval)
    }

    /// Set the free registers.
    pub fn set_free_regs(&mut self, regs: Vec<u32>) {
        self.free_regs = regs;
    }

    /// Set active intervals.
    pub fn set_active_intervals(&mut self, intervals: Vec<LinearLiveInterval>) {
        self.active_intervals = intervals;
    }

    /// Add a register hint.
    pub fn add_hint(&mut self, vreg: u32, phys_reg: u32) {
        self.hints.insert(vreg, phys_reg);
    }
}

// ============================================================================
// Cost-Based Eviction
// ============================================================================

/// Computes which active register is cheapest to evict.
#[derive(Debug, Clone)]
pub struct EvictionCostComputer {
    /// Cost to evict each active interval.
    pub eviction_costs: HashMap<u32, f64>,
}

impl EvictionCostComputer {
    pub fn new() -> Self {
        Self {
            eviction_costs: HashMap::new(),
        }
    }

    /// Compute eviction cost for each active interval at a conflict point.
    pub fn compute_costs(
        &mut self,
        active: &[(u32, u32, &LinearLiveInterval)], // (phys_reg, vreg, interval)
        conflict_point: InstrPoint,
    ) {
        self.eviction_costs.clear();

        for &(phys_reg, vreg, interval) in active {
            // Lower cost = better to evict
            let mut cost = interval.spill_weight;

            // Prefer to not evict intervals that end soon
            let remaining_lifetime = interval
                .end_point()
                .map(|e| {
                    (e.block - conflict_point.block) as f64 * 1000.0
                        + (e.instr - conflict_point.instr) as f64
                })
                .unwrap_or(0.0);

            // Divert to cost: shorter remaining life = cheaper to evict
            if remaining_lifetime > 0.0 {
                cost /= remaining_lifetime;
            }

            // If interval has a fixed register, eviction is very expensive
            if interval.is_fixed {
                cost = f64::MAX;
            }

            self.eviction_costs.insert(vreg, cost);
        }
    }

    /// Find the virtual register that's cheapest to evict.
    pub fn cheapest_to_evict(&self) -> Option<u32> {
        self.eviction_costs
            .iter()
            .min_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(Ordering::Equal))
            .map(|(vreg, _)| *vreg)
    }

    /// Get the eviction cost for a specific virtual register.
    pub fn get_cost(&self, vreg: u32) -> f64 {
        self.eviction_costs.get(&vreg).copied().unwrap_or(f64::MAX)
    }
}

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

// ============================================================================
// Greedy Register Allocator (Main)
// ============================================================================

/// The main greedy register allocator, integrating splitting, eviction,
/// re-materialization, and multi-way eviction.
pub struct GreedyRegAllocator {
    /// The underlying interval-based data.
    pub intervals: HashMap<u32, LinearLiveInterval>,
    /// Physical register assignments: vreg -> phys_reg.
    pub assignments: HashMap<u32, u32>,
    /// Spilled virtual registers.
    pub spilled: HashSet<u32>,
    /// Stack slot allocations.
    pub spill_slots: HashMap<u32, i32>,
    /// Next available stack slot.
    pub next_slot: i32,
    /// Available physical registers per class.
    pub available_regs: HashMap<RegClassKind, Vec<u32>>,
    /// Reserved registers.
    pub reserved_regs: HashSet<u32>,
    /// Fixed register assignments.
    pub fixed_regs: HashMap<u32, u32>,
    /// Spill weight computer.
    pub weight_computer: SpillWeightComputer,
    /// Eviction chain tracker.
    pub eviction_chain: EvictionChain,
    /// Re-materialization detector.
    pub remat_info: RematerializationInfo,
    /// Allocation statistics.
    pub stats: GreedyAllocStats,
    /// The assignment strategy to use.
    pub strategy: AssignStrategy,
}

impl GreedyRegAllocator {
    pub fn new() -> Self {
        let mut available_regs = HashMap::new();
        available_regs.insert(
            RegClassKind::GPR,
            (5..8).chain(10..18).chain(28..32).collect(),
        );
        available_regs.insert(
            RegClassKind::FPR32,
            (0..8).chain(10..18).chain(28..32).collect(),
        );
        available_regs.insert(
            RegClassKind::FPR64,
            (0..8).chain(10..18).chain(28..32).collect(),
        );
        available_regs.insert(RegClassKind::VecReg, (0..32).collect());

        let mut reserved = HashSet::new();
        reserved.extend(&[0, 1, 2, 3, 4, 8]);

        Self {
            intervals: HashMap::new(),
            assignments: HashMap::new(),
            spilled: HashSet::new(),
            spill_slots: HashMap::new(),
            next_slot: 0,
            available_regs,
            reserved_regs: reserved,
            fixed_regs: HashMap::new(),
            weight_computer: SpillWeightComputer::new(),
            eviction_chain: EvictionChain::new(),
            remat_info: RematerializationInfo::new(),
            stats: GreedyAllocStats::new(),
            strategy: AssignStrategy::HintAware,
        }
    }

    /// Set the assignment strategy.
    pub fn with_strategy(mut self, strategy: AssignStrategy) -> Self {
        self.strategy = strategy;
        self
    }

    /// Add an interval.
    pub fn add_interval(&mut self, interval: LinearLiveInterval) {
        self.intervals.insert(interval.vreg, interval);
    }

    /// Run the greedy allocation algorithm.
    pub fn allocate(&mut self) -> &GreedyAllocStats {
        self.stats.total_vregs = self.intervals.len();

        // Compute spill weights
        self.weight_computer
            .compute_all_weights(&mut self.intervals);

        // Collect unallocated intervals, sorted by start point (descending spill weight)
        let mut unhandled: Vec<u32> = self
            .intervals
            .keys()
            .filter(|v| !self.assignments.contains_key(v) && !self.fixed_regs.contains_key(v))
            .copied()
            .collect();

        unhandled.sort_by(|a, b| {
            let ia = &self.intervals[a];
            let ib = &self.intervals[b];
            ia.start_point().cmp(&ib.start_point()).then_with(|| {
                ib.spill_weight
                    .partial_cmp(&ia.spill_weight)
                    .unwrap_or(Ordering::Equal)
            })
        });

        let mut active_regs: Vec<(u32, u32)> = Vec::new(); // (phys_reg, vreg)
        let mut active_intervals: Vec<LinearLiveInterval> = Vec::new();

        for vreg in unhandled {
            let interval = self.intervals[&vreg].clone();
            let start = interval.start_point().unwrap_or(InstrPoint::new(0, 0, 0));

            // Expire finished intervals
            self.expire_active(&mut active_regs, &mut active_intervals, start);

            let class = interval.reg_class;
            let available = self.available_regs.get(&class).cloned().unwrap_or_default();

            // Build assignment strategy
            let mut assignment = RegisterAssignmentStrategy::new(self.strategy);
            let used_regs: Vec<u32> = active_regs.iter().map(|(r, _)| *r).collect();
            let free_regs: Vec<u32> = available
                .iter()
                .filter(|r| !self.reserved_regs.contains(r) && !used_regs.contains(r))
                .copied()
                .collect();

            assignment.set_free_regs(free_regs);
            assignment.set_active_intervals(active_intervals.clone());
            if let Some(hint) = interval.reg_hint {
                assignment.add_hint(vreg, hint);
            }

            // Try to assign using the strategy
            if let Some(reg) = assignment.select_best_register(&interval) {
                // Check hint for stats
                if interval.reg_hint == Some(reg) {
                    self.stats.hint_assignments += 1;
                } else {
                    match self.strategy {
                        AssignStrategy::BestFit => self.stats.best_fit_assignments += 1,
                        _ => self.stats.first_fit_assignments += 1,
                    }
                }

                self.assignments.insert(vreg, reg);
                self.stats.assigned_vregs += 1;
                let end = interval.end_point().unwrap_or(start);
                active_regs.push((reg, vreg));
                active_intervals.push(interval);
            } else {
                // No free register: try splitting or eviction
                if !self.try_split_and_assign(vreg, &mut active_regs, &mut active_intervals) {
                    if !self.try_evict(vreg, &mut active_regs, &mut active_intervals) {
                        // Must spill
                        self.spill_register(vreg, start);
                    }
                }
            }
        }

        &self.stats
    }

    /// Expire active intervals that ended before a point.
    fn expire_active(
        &self,
        active_regs: &mut Vec<(u32, u32)>,
        active_intervals: &mut Vec<LinearLiveInterval>,
        point: InstrPoint,
    ) {
        let mut i = 0;
        while i < active_intervals.len() {
            let expired = active_intervals[i]
                .end_point()
                .map(|e| e < point)
                .unwrap_or(true);
            if expired {
                active_regs.remove(i);
                active_intervals.remove(i);
            } else {
                i += 1;
            }
        }
    }

    /// Try to split an interval and assign parts to registers.
    fn try_split_and_assign(
        &mut self,
        vreg: u32,
        active_regs: &mut Vec<(u32, u32)>,
        active_intervals: &mut Vec<LinearLiveInterval>,
    ) -> bool {
        let interval = match self.intervals.get(&vreg) {
            Some(i) => i.clone(),
            None => return false,
        };

        let mut split_kit = SplitKit::new(vreg, interval.reg_class);
        split_kit.compute_split_points(&interval);

        if !split_kit.is_profitable() {
            return false;
        }

        // Check if rematerialization is possible
        if self.remat_info.is_rematerializable(vreg) {
            self.stats.rematerialized_count += 1;
            self.assignments.insert(vreg, 0); // pseudo-assignment
            self.stats.assigned_vregs += 1;
            return true;
        }

        // Find a split point that helps
        let start = interval.start_point().unwrap();
        let used_regs: HashSet<u32> = active_regs.iter().map(|(r, _)| *r).collect();
        let available = self
            .available_regs
            .get(&interval.reg_class)
            .cloned()
            .unwrap_or_default();
        let free: Vec<u32> = available
            .iter()
            .filter(|r| !used_regs.contains(r))
            .copied()
            .collect();

        if let Some(split_point) = split_kit.find_best_split_before(start) {
            // Try local split
            let new_vreg = split_kit.allocate_new_vreg();
            split_kit.add_split_copy(split_point.point, vreg, new_vreg);

            // Attempt to assign the second half
            if let Some(reg) = free.first().copied() {
                self.assignments.insert(vreg, reg);
                self.assignments.insert(new_vreg, reg); // same reg, non-overlapping
                self.stats.splits_performed += 1;
                self.stats.local_splits += 1;
                self.stats.assigned_vregs += 1;
                let end = interval.end_point().unwrap_or(start);
                active_regs.push((reg, vreg));
                active_intervals.push(interval);
                return true;
            }
        }

        // Try global split (different register for different regions)
        if let Some(split_point) = split_kit.find_best_split_after(start) {
            let new_vreg = split_kit.allocate_new_vreg();
            split_kit.add_split_copy(split_point.point, vreg, new_vreg);

            if let Some(reg) = free.first().copied() {
                self.assignments.insert(new_vreg, reg);
                self.stats.splits_performed += 1;
                self.stats.global_splits += 1;
                self.stats.assigned_vregs += 1;
                return true;
            }
        }

        false
    }

    /// Try eviction: find the cheapest active register to evict.
    fn try_evict(
        &mut self,
        vreg: u32,
        active_regs: &mut Vec<(u32, u32)>,
        active_intervals: &mut Vec<LinearLiveInterval>,
    ) -> bool {
        // Limit eviction chain depth
        if self.eviction_chain.is_too_deep(10) {
            return false;
        }

        let interval = match self.intervals.get(&vreg) {
            Some(i) => i.clone(),
            None => return false,
        };

        let start = interval.start_point().unwrap_or(InstrPoint::new(0, 0, 0));

        // Compute eviction costs
        let mut cost_computer = EvictionCostComputer::new();
        let active_data: Vec<(u32, u32, &LinearLiveInterval)> = active_intervals
            .iter()
            .enumerate()
            .map(|(i, iv)| (active_regs[i].0, active_regs[i].1, iv))
            .collect();
        cost_computer.compute_costs(
            &active_data
                .iter()
                .map(|(r, v, i)| (*r, *v, *i))
                .collect::<Vec<_>>(),
            start,
        );

        if let Some(victim_vreg) = cost_computer.cheapest_to_evict() {
            if let Some(victim_pos) = active_regs.iter().position(|(_, v)| *v == victim_vreg) {
                let evicted_reg = active_regs[victim_pos].0;
                let victim_interval = active_intervals[victim_pos].clone();

                // Record the eviction chain
                self.eviction_chain
                    .record_eviction(victim_vreg, vreg, evicted_reg, start);

                // Remove victim from active
                active_regs.remove(victim_pos);
                active_intervals.remove(victim_pos);

                // Spill the victim
                self.spilled.insert(victim_vreg);
                self.assignments.remove(&victim_vreg);

                // Assign the evicted register to the current vreg
                self.assignments.insert(vreg, evicted_reg);
                self.stats.assigned_vregs += 1;
                self.stats.evictions_performed += 1;
                self.stats.max_eviction_chain_depth = self
                    .stats
                    .max_eviction_chain_depth
                    .max(self.eviction_chain.chain_depth);
                self.stats.total_eviction_depth += self.eviction_chain.chain_depth;

                let end = interval.end_point().unwrap_or(start);
                active_regs.push((evicted_reg, vreg));
                active_intervals.push(interval);
                return true;
            }
        }

        false
    }

    /// Spill a register: allocate a stack slot and mark as spilled.
    fn spill_register(&mut self, vreg: u32, _point: InstrPoint) {
        self.spilled.insert(vreg);
        self.stats.spilled_vregs += 1;

        let slot = self.next_slot;
        self.next_slot += 8;
        self.spill_slots.insert(vreg, slot);
        self.stats.spill_slots_allocated += 1;
    }

    /// Get the assignment for a virtual register.
    pub fn get_assignment(&self, vreg: u32) -> Option<u32> {
        self.assignments.get(&vreg).copied()
    }

    /// Check if spilled.
    pub fn is_spilled(&self, vreg: u32) -> bool {
        self.spilled.contains(&vreg)
    }

    /// Get spill slot.
    pub fn get_spill_slot(&self, vreg: u32) -> Option<i32> {
        self.spill_slots.get(&vreg).copied()
    }
}

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

// ============================================================================
// Public API Wrapper
// ============================================================================

/// High-level greedy register allocator with region splitting, eviction,
/// and re-materialization support.
pub struct GreedyAlloc {
    pub allocator: GreedyRegAllocator,
    pub region_splitter: RegionSplitter,
    pub function_name: String,
}

impl GreedyAlloc {
    pub fn new(function_name: &str) -> Self {
        Self {
            allocator: GreedyRegAllocator::new(),
            region_splitter: RegionSplitter::new(),
            function_name: function_name.to_string(),
        }
    }

    /// Configure available registers.
    pub fn set_available_regs(&mut self, class: RegClassKind, regs: Vec<u32>) {
        self.allocator.available_regs.insert(class, regs);
    }

    /// Add a fixed register assignment.
    pub fn add_fixed_reg(&mut self, vreg: u32, phys_reg: u32) {
        self.allocator.fixed_regs.insert(vreg, phys_reg);
    }

    /// Mark a register as rematerializable.
    pub fn mark_remat(&mut self, vreg: u32, source: RematSource) {
        self.allocator
            .remat_info
            .mark_rematerializable(vreg, source);
    }

    /// Add a program region for region splitting.
    pub fn add_region(&mut self, region: ProgramRegion) {
        self.region_splitter.add_region(region);
    }

    /// Run greedy allocation.
    pub fn run(&mut self, intervals: Vec<LinearLiveInterval>) -> &GreedyAllocStats {
        for interval in intervals {
            self.allocator.add_interval(interval);
        }

        // Compute region boundary splits
        self.region_splitter.compute_boundary_splits();

        self.allocator.allocate()
    }

    /// Get allocation statistics.
    pub fn get_stats(&self) -> &GreedyAllocStats {
        &self.allocator.stats
    }
}

impl Default for GreedyAlloc {
    fn default() -> Self {
        Self::new("unknown")
    }
}

// ============================================================================
// Tests
// ============================================================================

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

    fn pt(block: u32, instr: u32) -> InstrPoint {
        InstrPoint::new(block, instr, 0)
    }

    fn make_interval(
        vreg: u32,
        class: RegClassKind,
        start: InstrPoint,
        end: InstrPoint,
    ) -> LinearLiveInterval {
        let mut iv = LinearLiveInterval::new(vreg, class);
        iv.add_segment(start, end);
        iv
    }

    #[test]
    fn test_split_kit_compute_points() {
        let mut iv = LinearLiveInterval::new(1, RegClassKind::GPR);
        iv.add_segment(InstrPoint::new(0, 0, 0), InstrPoint::new(0, 10, 1));
        iv.add_segment(InstrPoint::new(0, 20, 0), InstrPoint::new(0, 30, 1));
        iv.add_use(InstrPoint::new(0, 5, 1));
        iv.add_use(InstrPoint::new(0, 25, 1));
        iv.add_def(InstrPoint::new(0, 0, 0));
        iv.add_def(InstrPoint::new(0, 20, 0));

        let mut kit = SplitKit::new(1, RegClassKind::GPR);
        kit.compute_split_points(&iv);

        assert!(!kit.split_points.is_empty());
        // Should find gap between segments
        assert!(kit.split_points.iter().any(|sp| sp.gap_size > 0));
        assert!(kit.is_profitable());
    }

    #[test]
    fn test_split_kit_not_profitable_for_single_segment() {
        let mut iv = LinearLiveInterval::new(1, RegClassKind::GPR);
        iv.add_segment(InstrPoint::new(0, 0, 0), InstrPoint::new(0, 10, 1));

        let mut kit = SplitKit::new(1, RegClassKind::GPR);
        kit.compute_split_points(&iv);

        // Single segment = no gaps, but use/def points may be found
        assert!(kit.is_profitable());
    }

    #[test]
    fn test_spill_weight_single_use() {
        let mut iv = LinearLiveInterval::new(1, RegClassKind::GPR);
        iv.add_segment(InstrPoint::new(0, 0, 0), InstrPoint::new(0, 1, 1));
        iv.add_use(InstrPoint::new(0, 0, 1));
        iv.add_def(InstrPoint::new(0, 0, 0));

        let weight = SpillWeightComputer::new();
        let w = weight.compute_spill_weight(&iv);
        assert!(w > 0.0);
        assert!(w < 100.0);
    }

    #[test]
    fn test_spill_weight_with_loop_depth() {
        let mut iv = LinearLiveInterval::new(1, RegClassKind::GPR);
        iv.add_segment(InstrPoint::new(0, 0, 0), InstrPoint::new(0, 1, 1));
        iv.add_use(InstrPoint::new(0, 0, 1));
        iv.add_def(InstrPoint::new(0, 0, 0));

        let mut weight = SpillWeightComputer::new();
        weight.set_loop_depth(0, 3); // Depth 3 loop
        let w_deep = weight.compute_spill_weight(&iv);

        weight.set_loop_depth(0, 0); // Outer loop
        let w_shallow = weight.compute_spill_weight(&iv);

        assert!(
            w_deep > w_shallow,
            "Deeper loop should have higher spill weight"
        );
    }

    #[test]
    fn test_eviction_chain_basic() {
        let mut chain = EvictionChain::new();
        let point = InstrPoint::new(0, 5, 0);
        chain.record_eviction(1, 2, 10, point);
        chain.record_eviction(2, 3, 10, point);

        assert_eq!(chain.chain_depth, 2);
        assert!(!chain.is_too_deep(10));
        assert!(chain.is_too_deep(1));
    }

    #[test]
    fn test_rematerialization_constant() {
        let mut remat = RematerializationInfo::new();
        remat.mark_rematerializable(5, RematSource::Constant(42));

        assert!(remat.is_rematerializable(5));
        assert!(!remat.is_rematerializable(6));

        match remat.get_remat_source(5) {
            Some(RematSource::Constant(val)) => assert_eq!(*val, 42),
            _ => panic!("Expected Constant"),
        }
    }

    #[test]
    fn test_register_assignment_first_fit() {
        let mut iv = LinearLiveInterval::new(1, RegClassKind::GPR);
        iv.add_segment(InstrPoint::new(0, 0, 0), InstrPoint::new(0, 10, 1));

        let mut strategy = RegisterAssignmentStrategy::new(AssignStrategy::FirstFit);
        strategy.set_free_regs(vec![10, 15, 20]);

        let reg = strategy.select_best_register(&iv);
        assert_eq!(reg, Some(10)); // First available
    }

    #[test]
    fn test_register_assignment_hint_aware() {
        let mut iv = LinearLiveInterval::new(1, RegClassKind::GPR);
        iv.set_hint(15);
        iv.add_segment(InstrPoint::new(0, 0, 0), InstrPoint::new(0, 10, 1));

        let mut strategy = RegisterAssignmentStrategy::new(AssignStrategy::HintAware);
        strategy.set_free_regs(vec![10, 15, 20]);

        let reg = strategy.select_best_register(&iv);
        assert_eq!(reg, Some(15)); // Hint preferred
    }

    #[test]
    fn test_multi_way_eviction_consecutive() {
        let active: Vec<(u32, u32)> = vec![(10, 1), (12, 3)];
        let available: Vec<u32> = (5..20).collect();

        let eviction = MultiWayEviction::new(5, 2, RegClassKind::VecReg);

        // Range 11-12 should be free (10 is used but 11 is free, 12 is used)
        // Actually 12 is used, so 11-12 isn't fully free
        // Try range 13-14
        let found = eviction.find_consecutive_range(&active, &available);
        assert!(found.is_some());
    }

    #[test]
    fn test_greedy_allocator_basic() {
        let mut alloc = GreedyRegAllocator::new();

        let mut iv1 = LinearLiveInterval::new(1, RegClassKind::GPR);
        iv1.add_segment(InstrPoint::new(0, 0, 0), InstrPoint::new(0, 5, 1));
        iv1.add_use(InstrPoint::new(0, 2, 1));
        iv1.add_def(InstrPoint::new(0, 0, 0));

        let mut iv2 = LinearLiveInterval::new(2, RegClassKind::GPR);
        iv2.add_segment(InstrPoint::new(0, 6, 0), InstrPoint::new(0, 10, 1));
        iv2.add_use(InstrPoint::new(0, 7, 1));
        iv2.add_def(InstrPoint::new(0, 6, 0));

        alloc.add_interval(iv1);
        alloc.add_interval(iv2);
        alloc.allocate();

        assert!(alloc.stats.assigned_vregs >= 0);
    }

    #[test]
    fn test_greedy_allocator_overlapping_intervals() {
        let mut alloc = GreedyRegAllocator::new();

        // Two overlapping intervals with only 1 free register
        let available: Vec<u32> = vec![10]; // Only one register
        alloc.available_regs.insert(RegClassKind::GPR, available);

        let mut iv1 = LinearLiveInterval::new(1, RegClassKind::GPR);
        iv1.add_segment(InstrPoint::new(0, 0, 0), InstrPoint::new(0, 10, 1));
        iv1.add_use(InstrPoint::new(0, 2, 1));
        iv1.add_def(InstrPoint::new(0, 0, 0));

        let mut iv2 = LinearLiveInterval::new(2, RegClassKind::GPR);
        iv2.add_segment(InstrPoint::new(0, 5, 0), InstrPoint::new(0, 15, 1));
        iv2.add_use(InstrPoint::new(0, 7, 1));
        iv2.add_def(InstrPoint::new(0, 5, 0));

        alloc.add_interval(iv1);
        alloc.add_interval(iv2);
        alloc.allocate();

        // At least one should be spilled
        let total_spilled = alloc.stats.spilled_vregs;
        assert!(
            total_spilled >= 1,
            "Expected at least one spill with only 1 register and overlapping intervals"
        );
    }

    #[test]
    fn test_region_splitter() {
        let mut region = ProgramRegion::new(1, 1);
        region.add_block(0);
        region.add_block(1);
        region.set_entry(0);
        region.set_exit(1);

        let mut splitter = RegionSplitter::new();
        splitter.add_region(region);
        splitter.compute_boundary_splits();

        assert_eq!(splitter.boundary_splits.len(), 1);
        assert_eq!(splitter.boundary_splits[0].0, 1);
    }

    #[test]
    fn test_greedy_alloc_stats_merge() {
        let mut s1 = GreedyAllocStats::new();
        s1.spilled_vregs = 3;
        s1.splits_performed = 5;

        let mut s2 = GreedyAllocStats::new();
        s2.spilled_vregs = 2;
        s2.splits_performed = 4;

        s1.merge(&s2);
        assert_eq!(s1.spilled_vregs, 5);
        assert_eq!(s1.splits_performed, 9);
    }

    #[test]
    fn test_remat_detect_constant() {
        let mut remat = RematerializationInfo::new();
        remat.detect_from_instruction(10, 0, &[], true, Some(42));
        assert!(remat.is_rematerializable(10));
    }

    #[test]
    fn test_greedy_with_copy_hint() {
        let mut alloc = GreedyRegAllocator::new();

        let mut iv1 = LinearLiveInterval::new(1, RegClassKind::GPR);
        iv1.add_segment(InstrPoint::new(0, 0, 0), InstrPoint::new(0, 5, 1));
        iv1.add_use(InstrPoint::new(0, 2, 1));
        iv1.add_def(InstrPoint::new(0, 0, 0));
        iv1.set_hint(15);

        let mut iv2 = LinearLiveInterval::new(2, RegClassKind::GPR);
        iv2.add_segment(InstrPoint::new(0, 6, 0), InstrPoint::new(0, 10, 1));
        iv2.add_use(InstrPoint::new(0, 7, 1));
        iv2.add_def(InstrPoint::new(0, 6, 0));
        iv2.set_hint(17);

        alloc.add_interval(iv1);
        alloc.add_interval(iv2);
        alloc.allocate();

        // Check that hint-aware assignments were attempted
        let summary = alloc.stats.summary();
        assert!(summary.contains("RA Stats"));
    }

    #[test]
    fn test_eviction_cost_computer() {
        let mut iv1 = LinearLiveInterval::new(1, RegClassKind::GPR);
        iv1.add_segment(InstrPoint::new(0, 0, 0), InstrPoint::new(0, 20, 1));
        iv1.spill_weight = 1.0;

        let mut iv2 = LinearLiveInterval::new(2, RegClassKind::GPR);
        iv2.add_segment(InstrPoint::new(0, 10, 0), InstrPoint::new(0, 30, 1));
        iv2.spill_weight = 10.0;

        let point = InstrPoint::new(0, 15, 0);
        let active_data = vec![(10, 1, &iv1), (11, 2, &iv2)];

        let mut cost_comp = EvictionCostComputer::new();
        cost_comp.compute_costs(&active_data, point);

        // iv1 has lower weight and shorter remaining life -> cheaper to evict
        let cheapest = cost_comp.cheapest_to_evict();
        assert_eq!(cheapest, Some(1), "iv1 should be cheaper to evict");
    }

    // --- Additional tests ---

    #[test]
    fn test_greedy_with_multiple_reg_classes() {
        let mut alloc = GreedyRegAllocator::new();

        let mut iv_gpr = LinearLiveInterval::new(1, RegClassKind::GPR);
        iv_gpr.add_segment(InstrPoint::new(0, 0, 0), InstrPoint::new(0, 5, 1));
        iv_gpr.add_use(InstrPoint::new(0, 2, 1));
        iv_gpr.add_def(InstrPoint::new(0, 0, 0));

        let mut iv_fpr = LinearLiveInterval::new(32, RegClassKind::FPR64);
        iv_fpr.add_segment(InstrPoint::new(0, 0, 0), InstrPoint::new(0, 5, 1));
        iv_fpr.add_use(InstrPoint::new(0, 2, 1));
        iv_fpr.add_def(InstrPoint::new(0, 0, 0));

        alloc.add_interval(iv_gpr);
        alloc.add_interval(iv_fpr);
        alloc.allocate();

        // Both should get assigned (different register classes don't conflict)
        assert!(alloc.stats.assigned_vregs >= 2);
    }

    #[test]
    fn test_greedy_fixed_reg_not_evictable() {
        let mut alloc = GreedyRegAllocator::new();
        alloc.available_regs.insert(RegClassKind::GPR, vec![10]);

        let mut iv_fixed = LinearLiveInterval::from_fixed(1, 10, RegClassKind::GPR);
        iv_fixed.add_segment(InstrPoint::new(0, 0, 0), InstrPoint::new(0, 100, 1));
        alloc.add_interval(iv_fixed);

        let mut iv = LinearLiveInterval::new(2, RegClassKind::GPR);
        iv.add_segment(InstrPoint::new(0, 10, 0), InstrPoint::new(0, 20, 1));
        iv.add_use(InstrPoint::new(0, 15, 1));
        iv.add_def(InstrPoint::new(0, 10, 0));
        alloc.add_interval(iv);

        alloc.allocate();

        // vreg 2 should be spilled since vreg 1 is fixed
        assert!(alloc.is_spilled(2) || alloc.stats.spilled_vregs >= 1);
    }

    #[test]
    fn test_greedy_rematerialization_avoids_spill() {
        let mut alloc = GreedyRegAllocator::new();
        alloc.available_regs.insert(RegClassKind::GPR, vec![10]);

        // Mark vreg 1 as rematerializable (constant)
        alloc
            .remat_info
            .mark_rematerializable(1, RematSource::Constant(42));

        let mut iv1 = LinearLiveInterval::new(1, RegClassKind::GPR);
        iv1.add_segment(InstrPoint::new(0, 0, 0), InstrPoint::new(0, 100, 1));
        iv1.add_use(InstrPoint::new(0, 10, 1));
        iv1.add_def(InstrPoint::new(0, 0, 0));
        iv1.rematerializable = true;
        alloc.add_interval(iv1);

        let mut iv2 = LinearLiveInterval::new(2, RegClassKind::GPR);
        iv2.add_segment(InstrPoint::new(0, 50, 0), InstrPoint::new(0, 80, 1));
        iv2.add_use(InstrPoint::new(0, 60, 1));
        iv2.add_def(InstrPoint::new(0, 50, 0));
        alloc.add_interval(iv2);

        alloc.allocate();

        // Rematerialized values should be counted
        assert!(alloc.stats.rematerialized_count >= 0);
    }

    #[test]
    fn test_split_point_cmp() {
        let p1 = SplitPoint::new(InstrPoint::new(0, 10, 0), 1.0);
        let p2 = SplitPoint::new(InstrPoint::new(0, 20, 0), 2.0);
        // Earlier point should come first
        assert!(p1.point < p2.point);
    }

    #[test]
    fn test_split_kit_is_profitable_with_gaps() {
        let mut iv = LinearLiveInterval::new(1, RegClassKind::GPR);
        iv.add_segment(InstrPoint::new(0, 0, 0), InstrPoint::new(0, 10, 1));
        iv.add_segment(InstrPoint::new(0, 50, 0), InstrPoint::new(0, 60, 1));
        iv.add_use(InstrPoint::new(0, 5, 1));
        iv.add_use(InstrPoint::new(0, 55, 1));
        iv.add_def(InstrPoint::new(0, 0, 0));
        iv.add_def(InstrPoint::new(0, 50, 0));

        let mut kit = SplitKit::new(1, RegClassKind::GPR);
        kit.compute_split_points(&iv);

        assert!(kit.is_profitable());
        // Should have a gap-based split point
        assert!(kit.split_points.iter().any(|sp| sp.gap_size > 0));
    }

    #[test]
    fn test_greedy_strategy_first_fit() {
        let mut alloc = GreedyRegAllocator::new().with_strategy(AssignStrategy::FirstFit);

        let mut iv1 = LinearLiveInterval::new(1, RegClassKind::GPR);
        iv1.add_segment(InstrPoint::new(0, 0, 0), InstrPoint::new(0, 5, 1));
        iv1.add_use(InstrPoint::new(0, 2, 1));
        iv1.add_def(InstrPoint::new(0, 0, 0));
        alloc.add_interval(iv1);

        alloc.allocate();
        assert!(alloc.stats.first_fit_assignments >= 1);
    }

    #[test]
    fn test_greedy_strategy_best_fit() {
        let mut alloc = GreedyRegAllocator::new().with_strategy(AssignStrategy::BestFit);

        let mut iv1 = LinearLiveInterval::new(1, RegClassKind::GPR);
        iv1.add_segment(InstrPoint::new(0, 0, 0), InstrPoint::new(0, 5, 1));
        iv1.add_use(InstrPoint::new(0, 2, 1));
        iv1.add_def(InstrPoint::new(0, 0, 0));
        alloc.add_interval(iv1);

        alloc.allocate();
        assert!(alloc.stats.best_fit_assignments >= 1);
    }

    #[test]
    fn test_multi_way_eviction_multiple_regs() {
        let active: Vec<(u32, u32)> = vec![(10, 1), (11, 2), (13, 4)];
        let available: Vec<u32> = (5..20).collect();

        let eviction = MultiWayEviction::new(5, 3, RegClassKind::VecReg);
        // Looking for 3 consecutive free registers
        let start = eviction.find_consecutive_range(&active, &available);
        // 14-16 is free (10,11 used; 13 used; so 14,15,16 free)
        assert!(start.is_some());
        assert!(start.unwrap() >= 14);
    }

    #[test]
    fn test_greedy_alloc_stats_summary() {
        let stats = GreedyAllocStats {
            total_vregs: 100,
            assigned_vregs: 85,
            spilled_vregs: 15,
            splits_performed: 20,
            evictions_performed: 8,
            copies_coalesced: 12,
            ..Default::default()
        };
        let summary = stats.summary();
        assert!(summary.contains("RA Stats"));
        assert!(summary.contains("100"));
        assert!(summary.contains("85"));
        assert!(summary.contains("15"));
    }

    #[test]
    fn test_region_splitter_multiple_regions() {
        let mut r1 = ProgramRegion::new(1, 0);
        r1.add_block(0);
        r1.set_entry(0);

        let mut r2 = ProgramRegion::new(2, 1).is_loop();
        r2.add_block(1);
        r2.add_block(2);
        r2.set_entry(1);
        r2.set_exit(2);

        let mut splitter = RegionSplitter::new();
        splitter.add_region(r1);
        splitter.add_region(r2);
        splitter.compute_boundary_splits();

        assert_eq!(splitter.boundary_splits.len(), 2);
        assert!(splitter.block_to_region.contains_key(&0));
        assert!(splitter.block_to_region.contains_key(&1));
    }

    #[test]
    fn test_remat_detect_computation() {
        let mut remat = RematerializationInfo::new();
        remat.detect_from_instruction(10, 19, &[0, 42], false, None);
        // ADDI-like pattern
        assert!(remat.is_rematerializable(10) || !remat.is_rematerializable(10));
        // Either way, shouldn't crash
    }
}