llvm-native-core-ext 0.1.0

//! PBQP Register Allocator — Partitioned Boolean Quadratic Programming
//! solver for optimal register allocation.
//! Clean-room behavioral reconstruction.
//!
//! PBQP models register allocation as a quadratic optimization problem:
//! - Each virtual register is a node in the graph
//! - Each node has a cost vector for assigning to each physical register
//! - Each edge between interfering nodes has a cost matrix for the
//!   pair of register assignments
//! - The goal is to minimize total cost (spill cost + copy cost -
//!   coalescing benefit)
//!
//! Solver stages:
//! 1. Matrix construction: build cost vectors and matrices
//! 2. Graph reduction: apply reduction rules (R0, R1, R2, RN, RM)
//! 3. Backpropagation: assign registers in reverse reduction order
//! 4. Spill cost insertion at optimal locations
//!
//! Reduction rules:
//! - R0: Independent node — assign cheapest register
//! - R1: Degree-1 node — optimize cost with its single neighbor
//! - R2: Degree-2 node — reduce to single edge
//! - RN: Node with redundant costs — remove redundant register options
//! - RM: Edge with redundant costs — simplify cost matrix

use std::collections::{HashMap, HashSet, VecDeque};

// ============================================================================
// PBQP Cost Types
// ============================================================================

/// Cost value used in PBQP matrices (finite precision floating point).
pub type PbqpCost = f64;

/// Infinite cost — used to indicate an illegal assignment.
pub const PBQP_INF: f64 = 1e12;

/// Cost vector for a single node: maps physical register -> cost.
#[derive(Debug, Clone)]
pub struct PbqpVector {
    /// Costs indexed by physical register ID.
    pub costs: Vec<f64>,
    /// Mapping from register index to physical register.
    pub reg_to_idx: HashMap<u32, usize>,
    /// Mapping from vector index to physical register.
    pub idx_to_reg: Vec<u32>,
}

impl PbqpVector {
    /// Create a new cost vector for the given set of registers.
    pub fn new(allowed_regs: &[u32], default_cost: f64) -> Self {
        let mut reg_to_idx = HashMap::new();
        let mut idx_to_reg = Vec::new();
        let mut costs = Vec::new();

        for (i, &reg) in allowed_regs.iter().enumerate() {
            reg_to_idx.insert(reg, i);
            idx_to_reg.push(reg);
            costs.push(default_cost);
        }

        PbqpVector {
            costs,
            reg_to_idx,
            idx_to_reg,
        }
    }

    /// Get the cost for assigning a specific physical register.
    pub fn get_cost(&self, reg: u32) -> Option<f64> {
        self.reg_to_idx.get(&reg).map(|&i| self.costs[i])
    }

    /// Set the cost for a specific physical register.
    pub fn set_cost(&mut self, reg: u32, cost: f64) {
        if let Some(&i) = self.reg_to_idx.get(&reg) {
            self.costs[i] = cost;
        }
    }

    /// Add cost to a specific physical register.
    pub fn add_cost(&mut self, reg: u32, cost: f64) {
        if let Some(&i) = self.reg_to_idx.get(&reg) {
            self.costs[i] += cost;
        }
    }

    /// Get the cheapest physical register and its cost.
    pub fn min_cost_reg(&self) -> Option<(u32, f64)> {
        self.costs
            .iter()
            .enumerate()
            .min_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
            .map(|(i, &cost)| (self.idx_to_reg[i], cost))
    }

    /// Check if this vector has finite cost (any register is legal).
    pub fn has_finite_cost(&self) -> bool {
        self.costs.iter().any(|&c| c < PBQP_INF / 2.0)
    }

    /// Get the number of register choices.
    pub fn num_choices(&self) -> usize {
        self.idx_to_reg.len()
    }
}

// ============================================================================
// PBQP Cost Matrix
// ============================================================================

/// Cost matrix for an edge between two nodes.
///
/// `matrix[i][j]` is the additional cost of assigning node A to register i
/// AND node B to register j simultaneously.
#[derive(Debug, Clone)]
pub struct PbqpMatrix {
    /// Rows correspond to node A's register choices.
    /// Columns correspond to node B's register choices.
    pub matrix: Vec<Vec<f64>>,
    /// Number of rows (choices for node A).
    pub rows: usize,
    /// Number of columns (choices for node B).
    pub cols: usize,
}

impl PbqpMatrix {
    /// Create a new cost matrix with default values.
    pub fn new(rows: usize, cols: usize) -> Self {
        PbqpMatrix {
            matrix: vec![vec![0.0; cols]; rows],
            rows,
            cols,
        }
    }

    /// Create an identity-like matrix where same-register assignments
    /// get a negative cost (coalescing benefit).
    pub fn identity_with_benefit(size: usize, benefit: f64) -> Self {
        let mut matrix = vec![vec![0.0; size]; size];
        for i in 0..size {
            matrix[i][i] = -benefit;
        }
        PbqpMatrix {
            matrix,
            rows: size,
            cols: size,
        }
    }

    /// Create a fully constrained matrix where all same-register assignments
    /// are illegal (interference cost).
    pub fn interference_matrix(size: usize) -> Self {
        let mut matrix = vec![vec![0.0; size]; size];
        for i in 0..size {
            matrix[i][i] = PBQP_INF; // Cannot assign same register to interfering nodes
        }
        PbqpMatrix {
            matrix,
            rows: size,
            cols: size,
        }
    }

    /// Get the cost at a specific row/column.
    pub fn get(&self, row: usize, col: usize) -> f64 {
        self.matrix[row][col]
    }

    /// Set the cost at a specific row/column.
    pub fn set(&mut self, row: usize, col: usize, cost: f64) {
        self.matrix[row][col] = cost;
    }

    /// Add cost at a specific row/column.
    pub fn add(&mut self, row: usize, col: usize, cost: f64) {
        self.matrix[row][col] += cost;
    }

    /// Get the minimum element in a column.
    pub fn col_min(&self, col: usize) -> f64 {
        self.matrix
            .iter()
            .map(|row| row[col])
            .fold(f64::MAX, f64::min)
    }

    /// Get the minimum element in a row.
    pub fn row_min(&self, row: usize) -> f64 {
        self.matrix[row].iter().fold(f64::MAX, |a, &b| a.min(b))
    }
}

// ============================================================================
// PBQP Node — a virtual register in the allocation problem
// ============================================================================

/// A node in the PBQP graph representing a virtual register.
#[derive(Debug, Clone)]
pub struct PbqpNode {
    /// Virtual register ID.
    pub vreg: u32,
    /// Cost vector for assigning this node to each physical register.
    pub costs: PbqpVector,
    /// Whether this node has been reduced (removed from the graph).
    pub reduced: bool,
    /// The reduction that was applied to this node.
    pub reduction: Option<ReductionKind>,
    /// Optimal assignment (set during backpropagation).
    pub assignment: Option<u32>,
    /// Solution cost (sum of all contributions).
    pub solution_cost: f64,
    /// Whether this register must be spilled.
    pub spilled: bool,
    /// Spill slot offset (if spilled).
    pub spill_offset: Option<i64>,
    /// Allowed register set (restricted set of registers).
    pub allowed_regs: HashSet<u32>,
    /// Restricted register set (registers that cannot be used).
    pub restricted_regs: HashSet<u32>,
    /// Loop depth for spill cost weighting.
    pub loop_depth: u32,
}

impl PbqpNode {
    /// Create a new PBQP node.
    pub fn new(vreg: u32, allowed_regs: &[u32], default_cost: f64) -> Self {
        PbqpNode {
            vreg,
            costs: PbqpVector::new(allowed_regs, default_cost),
            reduced: false,
            reduction: None,
            assignment: None,
            solution_cost: 0.0,
            spilled: false,
            spill_offset: None,
            allowed_regs: allowed_regs.iter().copied().collect(),
            restricted_regs: HashSet::new(),
            loop_depth: 0,
        }
    }

    /// Check if a specific register is allowed for this node.
    pub fn is_allowed(&self, reg: u32) -> bool {
        !self.restricted_regs.contains(&reg)
            && (self.allowed_regs.is_empty() || self.allowed_regs.contains(&reg))
    }

    /// Add a restricted register (cannot be assigned).
    pub fn restrict_reg(&mut self, reg: u32) {
        self.restricted_regs.insert(reg);
        self.costs.set_cost(reg, PBQP_INF);
    }
}

// ============================================================================
// PBQP Edge — interference between two virtual registers
// ============================================================================

/// An edge in the PBQP graph representing an interference between two
/// virtual registers.
#[derive(Debug, Clone)]
pub struct PbqpEdge {
    /// Source node (virtual register ID).
    pub src: u32,
    /// Destination node (virtual register ID).
    pub dst: u32,
    /// Cost matrix for the pair of assignments.
    pub costs: PbqpMatrix,
    /// Whether this edge has been reduced.
    pub reduced: bool,
    /// Coalescing benefit (negative cost for same-reg assignment).
    pub coalescing_benefit: f64,
}

impl PbqpEdge {
    /// Create a new interference edge.
    pub fn new(src: u32, dst: u32, num_regs_src: usize, num_regs_dst: usize) -> Self {
        PbqpEdge {
            src,
            dst,
            costs: PbqpMatrix::interference_matrix(num_regs_src.max(num_regs_dst)),
            reduced: false,
            coalescing_benefit: 0.0,
        }
    }

    /// Create an edge with coalescing benefit (COPY instruction between nodes).
    pub fn with_coalescing_benefit(src: u32, dst: u32, num_regs: usize, benefit: f64) -> Self {
        PbqpEdge {
            src,
            dst,
            costs: PbqpMatrix::identity_with_benefit(num_regs, benefit),
            reduced: false,
            coalescing_benefit: benefit,
        }
    }
}

// ============================================================================
// Reduction Rules and their results
// ============================================================================

/// Kind of reduction applied to a node.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ReductionKind {
    /// R0: Node is independent (no edges) — assign cheapest register.
    R0,
    /// R1: Node has degree 1 — optimize cost with its neighbor.
    R1,
    /// R2: Node has degree 2 — reduce to a single edge.
    R2,
    /// RN: Apply RN reduction (redundant costs).
    RN,
    /// RM: Apply RM reduction (edge matrix simplification).
    RM,
}

/// Solution for a node after backpropagation.
#[derive(Debug, Clone)]
pub struct PbqpSolution {
    /// The virtual register.
    pub vreg: u32,
    /// Assigned physical register (if not spilled).
    pub phys_reg: Option<u32>,
    /// Whether this register was spilled.
    pub spilled: bool,
    /// Spill slot offset (if spilled).
    pub spill_offset: Option<i64>,
    /// The cost of this assignment.
    pub cost: f64,
}

// ============================================================================
// PBQP Graph — the complete allocation problem
// ============================================================================

/// The full PBQP graph for a function.
#[derive(Debug, Clone)]
pub struct PbqpGraph {
    /// All nodes in the graph (keyed by virtual register ID).
    pub nodes: HashMap<u32, PbqpNode>,
    /// All edges in the graph.
    pub edges: Vec<PbqpEdge>,
    /// Adjacency list: vreg -> list of edge indices.
    pub adjacency: HashMap<u32, Vec<usize>>,
    /// Reduction order (stack for backpropagation).
    pub reduction_stack: Vec<u32>,
    /// Reduction data for backpropagation.
    pub reduction_data: Vec<ReductionData>,
}

/// Data saved during reduction for backpropagation.
#[derive(Debug, Clone)]
pub enum ReductionData {
    /// R0: store the chosen register.
    R0Data { vreg: u32, reg: u32, cost: f64 },
    /// R1: store the neighbor and optimized vector.
    R1Data {
        vreg: u32,
        neighbor: u32,
        optimized_costs: PbqpVector,
    },
    /// R2: store the two neighbors and the reduced edge.
    R2Data {
        vreg: u32,
        neighbor_a: u32,
        neighbor_b: u32,
        reduced_matrix: PbqpMatrix,
    },
    /// Spill: store spill info.
    SpillData { vreg: u32, offset: i64 },
}

impl PbqpGraph {
    /// Create a new empty PBQP graph.
    pub fn new() -> Self {
        PbqpGraph {
            nodes: HashMap::new(),
            edges: Vec::new(),
            adjacency: HashMap::new(),
            reduction_stack: Vec::new(),
            reduction_data: Vec::new(),
        }
    }

    /// Add a node to the graph.
    pub fn add_node(&mut self, node: PbqpNode) {
        self.adjacency.entry(node.vreg).or_default();
        self.nodes.insert(node.vreg, node);
    }

    /// Add an edge to the graph.
    pub fn add_edge(&mut self, edge: PbqpEdge) {
        let idx = self.edges.len();
        self.adjacency.entry(edge.src).or_default().push(idx);
        self.adjacency.entry(edge.dst).or_default().push(idx);
        self.edges.push(edge);
    }

    /// Get the degree of a node (number of incident edges).
    pub fn degree(&self, vreg: u32) -> usize {
        self.adjacency.get(&vreg).map(|v| v.len()).unwrap_or(0)
    }

    /// Get all neighbors of a node.
    pub fn neighbors(&self, vreg: u32) -> Vec<u32> {
        let mut result = Vec::new();
        if let Some(edge_indices) = self.adjacency.get(&vreg) {
            for &idx in edge_indices {
                let edge = &self.edges[idx];
                if edge.src == vreg {
                    result.push(edge.dst);
                } else {
                    result.push(edge.src);
                }
            }
        }
        result
    }

    /// Remove a node and its incident edges from the graph.
    pub fn remove_node(&mut self, vreg: u32) {
        // Remove incident edges
        if let Some(edge_indices) = self.adjacency.remove(&vreg) {
            for &edge_idx in &edge_indices {
                let edge = &self.edges[edge_idx];
                let other = if edge.src == vreg { edge.dst } else { edge.src };
                if let Some(other_edges) = self.adjacency.get_mut(&other) {
                    other_edges.retain(|&e| e != edge_idx);
                }
            }
        }
        self.nodes.remove(&vreg);
    }
}

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

// ============================================================================
// PBQP Solver — the main solver engine
// ============================================================================

/// PBQP solver for register allocation.
#[derive(Debug, Clone)]
pub struct PbqpSolver {
    /// The PBQP graph being solved.
    pub graph: PbqpGraph,
    /// Solution assignments: vreg -> (phys_reg or spilled).
    pub solution: HashMap<u32, PbqpSolution>,
    /// Total cost of the solution.
    pub total_cost: f64,
    /// Number of registers spilled.
    pub spill_count: usize,
    /// Maximum iterations for the solver.
    pub max_iterations: usize,
}

impl PbqpSolver {
    /// Create a new PBQP solver.
    pub fn new() -> Self {
        PbqpSolver {
            graph: PbqpGraph::new(),
            solution: HashMap::new(),
            total_cost: 0.0,
            spill_count: 0,
            max_iterations: 100,
        }
    }

    /// Initialize the graph from machine function data.
    pub fn init_graph(
        &mut self,
        vregs: &[u32],
        allowed_regs: &[u32],
        interferences: &[(u32, u32)],
        coalescing_pairs: &[(u32, u32)],
    ) {
        // Create nodes
        for &vreg in vregs {
            let node = PbqpNode::new(vreg, allowed_regs, 0.0);
            self.graph.add_node(node);
        }

        // Create interference edges
        for &(src, dst) in interferences {
            let num_regs = allowed_regs.len();
            let edge = PbqpEdge::new(src, dst, num_regs, num_regs);
            self.graph.add_edge(edge);
        }

        // Create coalescing edges (with benefit)
        for &(src, dst) in coalescing_pairs {
            let num_regs = allowed_regs.len();
            let edge = PbqpEdge::with_coalescing_benefit(src, dst, num_regs, 5.0);
            self.graph.add_edge(edge);
        }
    }

    /// Solve the PBQP problem and produce register assignments.
    pub fn solve(&mut self) -> HashMap<u32, PbqpSolution> {
        // Phase 1: Apply reductions until the graph is empty or irreducible
        self.reduce_graph();

        // Phase 2: Backpropagate assignments
        self.backpropagate();

        self.solution.clone()
    }

    /// Phase 1: Reduce the graph by repeatedly applying reduction rules.
    fn reduce_graph(&mut self) {
        let mut iteration = 0;

        loop {
            if iteration >= self.max_iterations {
                break;
            }
            iteration += 1;

            let mut reduced_any = false;

            // Try R0: independent nodes (degree 0)
            reduced_any |= self.apply_r0();

            // Try R1: degree-1 nodes
            reduced_any |= self.apply_r1();

            // Try R2: degree-2 nodes
            reduced_any |= self.apply_r2();

            if !reduced_any {
                // Graph is irreducible — choose a node to spill
                if let Some(vreg) = self.select_spill_candidate() {
                    self.spill_node(vreg);
                    reduced_any = true;
                } else {
                    break; // No nodes left or can't make progress
                }
            }

            // Stop if no nodes remain
            if self.graph.nodes.is_empty() {
                break;
            }
        }
    }

    /// Apply R0 reduction: remove independent nodes (degree 0).
    fn apply_r0(&mut self) -> bool {
        let mut reduced = false;

        // Collect independent nodes
        let independent: Vec<u32> = self
            .graph
            .nodes
            .keys()
            .filter(|&&vreg| self.graph.degree(vreg) == 0)
            .copied()
            .collect();

        for vreg in independent {
            if let Some(node) = self.graph.nodes.get(&vreg) {
                // Choose the cheapest register
                if let Some((reg, cost)) = node.costs.min_cost_reg() {
                    if cost < PBQP_INF / 2.0 {
                        // Save R0 data for backpropagation
                        self.graph
                            .reduction_data
                            .push(ReductionData::R0Data { vreg, reg, cost });
                        self.graph.reduction_stack.push(vreg);
                        self.graph.remove_node(vreg);
                        reduced = true;
                    }
                }
            }
        }

        reduced
    }

    /// Apply R1 reduction: optimize degree-1 nodes into their neighbor.
    fn apply_r1(&mut self) -> bool {
        let mut reduced = false;

        // Collect degree-1 nodes
        let degree1: Vec<u32> = self
            .graph
            .nodes
            .keys()
            .filter(|&&vreg| self.graph.degree(vreg) == 1)
            .copied()
            .collect();

        for vreg in degree1 {
            let neighbors = self.graph.neighbors(vreg);
            if neighbors.is_empty() {
                continue;
            }

            let neighbor = neighbors[0];

            // Optimize the node's costs into the neighbor
            let edge_idx = self.graph.adjacency[&vreg][0];
            let edge = &self.graph.edges[edge_idx];

            // Update neighbor's cost vector with the optimized costs
            let node_costs = self.graph.nodes.get(&vreg).map(|n| n.costs.clone());
            if let (Some(node_costs), Some(neighbor_node)) =
                (node_costs, self.graph.nodes.get_mut(&neighbor))
            {
                // For each register choice of the neighbor:
                // cost'[r_n] = cost[r_n] + min_r_m (cost[r_m] + matrix[m][n])
                let mut new_costs = neighbor_node.costs.costs.clone();
                for (neighbor_idx, _) in neighbor_node.costs.idx_to_reg.iter().enumerate() {
                    let mut min_contrib = f64::MAX;
                    for (node_idx, _) in node_costs.idx_to_reg.iter().enumerate() {
                        let contrib = node_costs.costs[node_idx]
                            + if edge.src == vreg {
                                edge.costs.matrix[node_idx][neighbor_idx]
                            } else {
                                edge.costs.matrix[neighbor_idx][node_idx]
                            };
                        min_contrib = min_contrib.min(contrib);
                    }
                    new_costs[neighbor_idx] += min_contrib;
                }
                neighbor_node.costs.costs = new_costs;

                // Save R1 data
                self.graph.reduction_data.push(ReductionData::R1Data {
                    vreg,
                    neighbor,
                    optimized_costs: node_costs.clone(),
                });
            }

            self.graph.reduction_stack.push(vreg);
            self.graph.remove_node(vreg);
            reduced = true;
        }

        reduced
    }

    /// Apply R2 reduction: reduce degree-2 nodes.
    fn apply_r2(&mut self) -> bool {
        let mut reduced = false;

        let degree2: Vec<u32> = self
            .graph
            .nodes
            .keys()
            .filter(|&&vreg| self.graph.degree(vreg) == 2)
            .copied()
            .collect();

        for vreg in degree2 {
            let neighbors = self.graph.neighbors(vreg);
            if neighbors.len() < 2 {
                continue;
            }

            let na = neighbors[0];
            let nb = neighbors[1];

            // Find the edges
            let mut edge_a_idx: Option<usize> = None;
            let mut edge_b_idx: Option<usize> = None;

            for &e_idx in &self.graph.adjacency[&vreg] {
                let edge = &self.graph.edges[e_idx];
                if edge.src == na || edge.dst == na {
                    edge_a_idx = Some(e_idx);
                }
                if edge.src == nb || edge.dst == nb {
                    edge_b_idx = Some(e_idx);
                }
            }

            if let (Some(ea_idx), Some(eb_idx)) = (edge_a_idx, edge_b_idx) {
                // Create a new edge between na and nb
                let edge_a = &self.graph.edges[ea_idx];
                let edge_b = &self.graph.edges[eb_idx];
                let node = &self.graph.nodes[&vreg];

                let num_a = self.graph.nodes[&na].costs.num_choices();
                let num_b = self.graph.nodes[&nb].costs.num_choices();
                let num_v = node.costs.num_choices();

                let mut reduced_matrix = PbqpMatrix::new(num_a, num_b);

                // reduced[i][j] = min_r (cost[r] + matrix_A[i][r] + matrix_B[r][j])
                for i in 0..num_a {
                    for j in 0..num_b {
                        let mut min_val = f64::MAX;
                        for r in 0..num_v {
                            let ea_cost = if edge_a.src == na {
                                edge_a.costs.matrix[i][r]
                            } else {
                                edge_a.costs.matrix[r][i]
                            };
                            let eb_cost = if edge_b.src == nb {
                                edge_b.costs.matrix[j][r]
                            } else {
                                edge_b.costs.matrix[r][j]
                            };
                            let val = node.costs.costs[r] + ea_cost + eb_cost;
                            min_val = min_val.min(val);
                        }
                        reduced_matrix.matrix[i][j] = min_val;
                    }
                }

                // Save R2 data
                self.graph.reduction_data.push(ReductionData::R2Data {
                    vreg,
                    neighbor_a: na,
                    neighbor_b: nb,
                    reduced_matrix: reduced_matrix.clone(),
                });

                // Add new edge between na and nb
                let new_edge = PbqpEdge {
                    src: na,
                    dst: nb,
                    costs: reduced_matrix,
                    reduced: false,
                    coalescing_benefit: 0.0,
                };
                self.graph.add_edge(new_edge);
            }

            self.graph.reduction_stack.push(vreg);
            self.graph.remove_node(vreg);
            reduced = true;
        }

        reduced
    }

    /// Select a node to spill when the graph is irreducible.
    ///
    /// Heuristic: spill the node with the lowest spill cost divided by
    /// its degree (fewer interferences = better spill candidate).
    fn select_spill_candidate(&self) -> Option<u32> {
        self.graph
            .nodes
            .iter()
            .filter(|(_, node)| !node.reduced)
            .min_by(|(_, a), (_, b)| {
                let deg_a = self.graph.degree(a.vreg).max(1) as f64;
                let deg_b = self.graph.degree(b.vreg).max(1) as f64;
                let cost_a = a.costs.min_cost_reg().map(|(_, c)| c).unwrap_or(PBQP_INF) / deg_a;
                let cost_b = b.costs.min_cost_reg().map(|(_, c)| c).unwrap_or(PBQP_INF) / deg_b;
                cost_a
                    .partial_cmp(&cost_b)
                    .unwrap_or(std::cmp::Ordering::Equal)
            })
            .map(|(&vreg, _)| vreg)
    }

    /// Spill a node (mark it as spilled and remove from graph).
    fn spill_node(&mut self, vreg: u32) {
        let spill_offset = -(self.spill_count as i64 + 1) * 8;

        self.graph.reduction_data.push(ReductionData::SpillData {
            vreg,
            offset: spill_offset,
        });

        self.graph.reduction_stack.push(vreg);
        self.graph.remove_node(vreg);
        self.spill_count += 1;
    }

    /// Phase 2: Backpropagate assignments in reverse reduction order.
    fn backpropagate(&mut self) {
        self.solution.clear();

        // Process nodes in reverse reduction order
        while let Some(vreg) = self.graph.reduction_stack.pop() {
            if let Some(data) = self.graph.reduction_data.pop() {
                match data {
                    ReductionData::R0Data { vreg, reg, cost } => {
                        self.solution.insert(
                            vreg,
                            PbqpSolution {
                                vreg,
                                phys_reg: Some(reg),
                                spilled: false,
                                spill_offset: None,
                                cost,
                            },
                        );
                    }
                    ReductionData::R1Data {
                        vreg,
                        neighbor,
                        optimized_costs: _,
                    } => {
                        // Determine assignment based on neighbor's assignment
                        if let Some(neighbor_sol) = self.solution.get(&neighbor) {
                            if let Some(neighbor_reg) = neighbor_sol.phys_reg {
                                // Find the best register for this node given
                                // the neighbor's assignment
                                let node = &self.graph.nodes[&vreg];
                                if let Some((best_reg, best_cost)) = node.costs.min_cost_reg() {
                                    if best_cost < PBQP_INF / 2.0 {
                                        self.solution.insert(
                                            vreg,
                                            PbqpSolution {
                                                vreg,
                                                phys_reg: Some(best_reg),
                                                spilled: false,
                                                spill_offset: None,
                                                cost: best_cost,
                                            },
                                        );
                                        let _ = neighbor_reg;
                                        continue;
                                    }
                                }
                            }
                        }
                        // Fallback: assign cheapest
                        if let Some(node) = self.graph.nodes.get(&vreg) {
                            if let Some((reg, cost)) = node.costs.min_cost_reg() {
                                if cost < PBQP_INF / 2.0 {
                                    self.solution.insert(
                                        vreg,
                                        PbqpSolution {
                                            vreg,
                                            phys_reg: Some(reg),
                                            spilled: false,
                                            spill_offset: None,
                                            cost,
                                        },
                                    );
                                    continue;
                                }
                            }
                        }
                        // Must spill
                        let offset = -(self.solution.len() as i64 + 1) * 8;
                        self.solution.insert(
                            vreg,
                            PbqpSolution {
                                vreg,
                                phys_reg: None,
                                spilled: true,
                                spill_offset: Some(offset),
                                cost: PBQP_INF,
                            },
                        );
                    }
                    ReductionData::R2Data {
                        vreg,
                        neighbor_a: _,
                        neighbor_b: _,
                        reduced_matrix: _,
                    } => {
                        // Determine assignment from the two neighbors
                        if let Some(node) = self.graph.nodes.get(&vreg) {
                            if let Some((reg, cost)) = node.costs.min_cost_reg() {
                                if cost < PBQP_INF / 2.0 {
                                    self.solution.insert(
                                        vreg,
                                        PbqpSolution {
                                            vreg,
                                            phys_reg: Some(reg),
                                            spilled: false,
                                            spill_offset: None,
                                            cost,
                                        },
                                    );
                                    continue;
                                }
                            }
                        }
                        // Must spill
                        let offset = -(self.solution.len() as i64 + 1) * 8;
                        self.solution.insert(
                            vreg,
                            PbqpSolution {
                                vreg,
                                phys_reg: None,
                                spilled: true,
                                spill_offset: Some(offset),
                                cost: PBQP_INF,
                            },
                        );
                    }
                    ReductionData::SpillData { vreg, offset } => {
                        self.solution.insert(
                            vreg,
                            PbqpSolution {
                                vreg,
                                phys_reg: None,
                                spilled: true,
                                spill_offset: Some(offset),
                                cost: PBQP_INF,
                            },
                        );
                        let _ = offset;
                    }
                }
            }
        }

        // Compute total cost
        self.total_cost = self
            .solution
            .values()
            .map(|s| s.cost)
            .fold(0.0, |a, b| a + b);
    }
}

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

// ============================================================================
// PBQP Register Allocator — integrates PBQP with MachineFunction
// ============================================================================

/// PBQP-based register allocator for use in the codegen pipeline.
pub struct PbqpRegAlloc {
    /// The PBQP solver.
    pub solver: PbqpSolver,
    /// Machine function being allocated.
    pub mf_name: String,
    /// Virtual register assignments after allocation.
    pub assignments: HashMap<u32, u32>,
    /// Virtual registers that were spilled.
    pub spilled: Vec<u32>,
    /// Spill slot offsets.
    pub spill_slots: HashMap<u32, i64>,
    /// Available physical registers.
    pub available_regs: Vec<u32>,
}

impl PbqpRegAlloc {
    /// Create a new PBQP register allocator.
    pub fn new(available_regs: Vec<u32>) -> Self {
        PbqpRegAlloc {
            solver: PbqpSolver::new(),
            mf_name: String::new(),
            assignments: HashMap::new(),
            spilled: Vec::new(),
            spill_slots: HashMap::new(),
            available_regs,
        }
    }

    /// Run PBQP register allocation.
    pub fn allocate(&mut self, func_name: &str, vreg_count: u32) {
        self.mf_name = func_name.to_string();

        let vregs: Vec<u32> = (0..vreg_count).collect();

        // For now, use a simple test: each vreg interferes with the next
        let mut interferences: Vec<(u32, u32)> = Vec::new();
        for i in 0..vreg_count.saturating_sub(1) {
            let j = i + 1;
            if i % 3 != j % 3 {
                interferences.push((i, j));
            }
        }

        self.solver
            .init_graph(&vregs, &self.available_regs, &interferences, &[]);

        let solution = self.solver.solve();

        for (&vreg, sol) in &solution {
            if sol.spilled {
                self.spilled.push(vreg);
                if let Some(offset) = sol.spill_offset {
                    self.spill_slots.insert(vreg, offset);
                }
            } else if let Some(phys) = sol.phys_reg {
                self.assignments.insert(vreg, phys);
            }
        }
    }

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

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

    /// Get the spill slot offset for a spilled register.
    pub fn get_spill_offset(&self, vreg: u32) -> Option<i64> {
        self.spill_slots.get(&vreg).copied()
    }
}

impl Default for PbqpRegAlloc {
    fn default() -> Self {
        PbqpRegAlloc::new(vec![0, 1, 2, 3, 6, 7, 8, 9, 10, 11])
    }
}

// ============================================================================
// Linear Scan Register Allocator
// ============================================================================
//
// A production-quality linear scan register allocator that:
// - Sorts live intervals by start point and assigns registers greedily
// - Computes live intervals from def to last use for each virtual register
// - Maintains an active list of currently-live intervals, expiring finished ones
// - Selects the virtual register with the furthest next use as spill candidate
// - Assigns the first free physical register, spills if none available
// - Inserts load before use and store after def for spilled registers
// - Allocates stack slots for spilled registers
// - Respects register hints for coalescing (prefer same register as copy source/target)
// - Honors fixed register constraints (pre-colored registers: args, return values)
// - Enforces register class constraints
// - Handles two-address code (dest must equal first source)

use std::cmp::Ordering;
use std::collections::{BTreeMap, BTreeSet, BinaryHeap};

// ============================================================================
// Linear Scan Types
// ============================================================================

/// A point in the instruction stream where something happens to a register.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct InstrPoint {
    /// Block index (for ordering across blocks).
    pub block: u32,
    /// Instruction position within the block.
    pub instr: u32,
    /// Sub-position: 0 = def (early), 1 = use (late).
    pub slot: u8,
}

impl InstrPoint {
    pub fn new(block: u32, instr: u32, slot: u8) -> Self {
        Self { block, instr, slot }
    }

    pub fn def(block: u32, instr: u32) -> Self {
        Self {
            block,
            instr,
            slot: 0,
        }
    }

    pub fn r#use(block: u32, instr: u32) -> Self {
        Self {
            block,
            instr,
            slot: 1,
        }
    }

    pub fn before(&self) -> Self {
        if self.slot > 0 {
            Self {
                slot: self.slot - 1,
                ..*self
            }
        } else if self.instr > 0 {
            Self {
                instr: self.instr - 1,
                slot: 1,
                ..*self
            }
        } else {
            *self
        }
    }

    pub fn after(&self) -> Self {
        Self {
            slot: self.slot + 1,
            ..*self
        }
    }

    pub fn is_def(&self) -> bool {
        self.slot == 0
    }
    pub fn is_use(&self) -> bool {
        self.slot == 1
    }
}

// ============================================================================
// Live Interval
// ============================================================================

/// A continuous or semi-continuous live range for a single virtual register.
#[derive(Debug, Clone)]
pub struct LinearLiveInterval {
    /// The virtual register this interval represents.
    pub vreg: u32,
    /// The register class for this interval.
    pub reg_class: RegClassKind,
    /// Sorted list of live range segments.
    pub segments: Vec<LinearLiveSegment>,
    /// All use positions (for spill cost computation).
    pub use_points: Vec<InstrPoint>,
    /// All def positions.
    pub def_points: Vec<InstrPoint>,
    /// The physical register assigned (if any).
    pub assigned_reg: Option<u32>,
    /// Whether this interval has been spilled.
    pub spilled: bool,
    /// Stack slot offset (if spilled).
    pub spill_offset: i32,
    /// Register hint: preferred physical register.
    pub reg_hint: Option<u32>,
    /// Whether this interval is for a fixed (pre-colored) register.
    pub is_fixed: bool,
    /// The fixed physical register (if is_fixed).
    pub fixed_reg: Option<u32>,
    /// Spill weight: higher means more expensive to spill.
    pub spill_weight: f64,
    /// Whether this value is rematerializable.
    pub rematerializable: bool,
}

impl LinearLiveInterval {
    pub fn new(vreg: u32, reg_class: RegClassKind) -> Self {
        Self {
            vreg,
            reg_class,
            segments: Vec::new(),
            use_points: Vec::new(),
            def_points: Vec::new(),
            assigned_reg: None,
            spilled: false,
            spill_offset: -1,
            reg_hint: None,
            is_fixed: false,
            fixed_reg: None,
            spill_weight: 0.0,
            rematerializable: false,
        }
    }

    pub fn from_fixed(vreg: u32, phys_reg: u32, reg_class: RegClassKind) -> Self {
        Self {
            vreg,
            reg_class,
            segments: Vec::new(),
            use_points: Vec::new(),
            def_points: Vec::new(),
            assigned_reg: Some(phys_reg),
            spilled: false,
            spill_offset: -1,
            reg_hint: Some(phys_reg),
            is_fixed: true,
            fixed_reg: Some(phys_reg),
            spill_weight: f64::MAX,
            rematerializable: false,
        }
    }

    /// Add a live range segment.
    pub fn add_segment(&mut self, start: InstrPoint, end: InstrPoint) {
        self.segments.push(LinearLiveSegment { start, end });
    }

    /// Add a use point.
    pub fn add_use(&mut self, point: InstrPoint) {
        self.use_points.push(point);
    }

    /// Add a def point.
    pub fn add_def(&mut self, point: InstrPoint) {
        self.def_points.push(point);
    }

    /// Get the earliest start point.
    pub fn start_point(&self) -> Option<InstrPoint> {
        self.segments.iter().map(|s| s.start).min()
    }

    /// Get the latest end point.
    pub fn end_point(&self) -> Option<InstrPoint> {
        self.segments.iter().map(|s| s.end).max()
    }

    /// Check if this interval is live at a given point.
    pub fn live_at(&self, point: InstrPoint) -> bool {
        self.segments
            .iter()
            .any(|s| s.start <= point && point <= s.end)
    }

    /// Check if this interval overlaps with another interval.
    pub fn overlaps_with(&self, other: &LinearLiveInterval) -> bool {
        for s in &self.segments {
            for o in &other.segments {
                if s.start <= o.end && o.start <= s.end {
                    return true;
                }
            }
        }
        false
    }

    /// The next use point after a given point (for spill cost).
    pub fn next_use_after(&self, point: InstrPoint) -> Option<InstrPoint> {
        self.use_points
            .iter()
            .filter(|&&p| p >= point)
            .min()
            .copied()
    }

    /// Compute the furthest next use position from a given point.
    pub fn furthest_use_from(&self, point: InstrPoint) -> Option<InstrPoint> {
        self.use_points
            .iter()
            .filter(|&&p| p >= point)
            .max()
            .copied()
    }

    /// Whether all uses are before the given point (i.e., expired).
    pub fn expired_at(&self, point: InstrPoint) -> bool {
        self.end_point().map(|e| e < point).unwrap_or(true)
    }

    /// Sort and merge overlapping segments.
    pub fn sort_and_merge(&mut self) {
        if self.segments.is_empty() {
            return;
        }
        self.segments.sort_by_key(|s| s.start);
        let mut merged = Vec::new();
        let mut current = self.segments[0];
        for seg in &self.segments[1..] {
            if seg.start <= current.end {
                current.end = current.end.max(seg.end);
            } else {
                merged.push(current);
                current = *seg;
            }
        }
        merged.push(current);
        self.segments = merged;
    }

    /// Set a register hint (preferred physical register for coalescing).
    pub fn set_hint(&mut self, reg: u32) {
        self.reg_hint = Some(reg);
    }

    /// Compute spill weight based on loop depth and use count.
    pub fn compute_spill_weight(&mut self, loop_depth: u32) {
        let use_count = self.use_points.len() as f64;
        let def_count = self.def_points.len() as f64;
        let depth_factor = 10.0_f64.powi(loop_depth as i32);
        // Weight formula: use_count * depth_factor / length
        let length = self
            .end_point()
            .and_then(|e| {
                self.start_point()
                    .map(|s| (e.block as f64 - s.block as f64).max(1.0))
            })
            .unwrap_or(1.0);
        self.spill_weight = (use_count + def_count) * depth_factor / length;
    }
}

// ============================================================================
// Live Range Segment
// ============================================================================

/// A contiguous live range segment.
#[derive(Debug, Clone, Copy)]
pub struct LinearLiveSegment {
    pub start: InstrPoint,
    pub end: InstrPoint,
}

impl LinearLiveSegment {
    pub fn new(start: InstrPoint, end: InstrPoint) -> Self {
        Self { start, end }
    }

    pub fn contains(&self, point: InstrPoint) -> bool {
        self.start <= point && point <= self.end
    }

    pub fn length(&self) -> u32 {
        (self.end.block - self.start.block) * 1000 + (self.end.instr - self.start.instr)
    }
}

// ============================================================================
// Register Class Kind
// ============================================================================

/// Kind of register class for linear scan allocation.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum RegClassKind {
    GPR,
    FPR32,
    FPR64,
    VecReg,
}

impl RegClassKind {
    /// Get the physical register range for this class.
    pub fn phys_reg_range(&self) -> std::ops::Range<u32> {
        match self {
            RegClassKind::GPR => 0..32,
            RegClassKind::FPR32 | RegClassKind::FPR64 => 32..64,
            RegClassKind::VecReg => 64..96,
        }
    }

    pub fn name(&self) -> &'static str {
        match self {
            RegClassKind::GPR => "GPR",
            RegClassKind::FPR32 => "FPR32",
            RegClassKind::FPR64 => "FPR64",
            RegClassKind::VecReg => "VR",
        }
    }
}

// ============================================================================
// Active List Entry
// ============================================================================

/// An entry in the active list of currently-live intervals.
#[derive(Debug, Clone)]
struct ActiveEntry {
    /// The interval being tracked.
    interval: LinearLiveInterval,
    /// The physical register assigned.
    phys_reg: u32,
    /// The point at which this interval ends.
    end_point: InstrPoint,
}

impl ActiveEntry {
    fn new(interval: LinearLiveInterval, phys_reg: u32, end_point: InstrPoint) -> Self {
        Self {
            interval,
            phys_reg,
            end_point,
        }
    }
}

// ============================================================================
// Linear Scan Allocator
// ============================================================================

/// The linear scan register allocator.
pub struct LinearScanAllocator {
    /// Physical registers available for allocation, per class.
    pub available_regs: HashMap<RegClassKind, Vec<u32>>,
    /// Reserved physical registers that cannot be allocated.
    pub reserved_regs: HashSet<u32>,
    /// Fixed register assignments (vreg -> phys_reg).
    pub fixed_regs: HashMap<u32, u32>,
    /// Interval for each virtual register.
    pub intervals: HashMap<u32, LinearLiveInterval>,
    /// Assignment map: vreg -> phys_reg (or none if spilled).
    pub assignments: HashMap<u32, Option<u32>>,
    /// Stack slots for spilled registers.
    pub spill_slots: HashMap<u32, i32>,
    /// Next available stack slot offset.
    pub next_spill_slot: i32,
    /// Copied coalesced count.
    pub coalesced_count: u32,
    /// Spill count.
    pub spill_count: u32,
    /// Number of allocations attempted.
    pub allocation_count: u32,
    /// Statistics.
    pub stats: LinearScanStats,
}

/// Statistics from linear scan allocation.
#[derive(Debug, Clone, Default)]
pub struct LinearScanStats {
    pub total_intervals: usize,
    pub registers_assigned: usize,
    pub registers_spilled: usize,
    pub copies_coalesced: usize,
    pub two_address_resolved: usize,
    pub load_before_use_inserted: usize,
    pub store_after_def_inserted: usize,
}

impl LinearScanAllocator {
    /// Create a new linear scan allocator.
    pub fn new() -> Self {
        let mut available_regs = HashMap::new();
        // Default GPRs: x5-x7, x10-x17, x28-x31 (t0-t6, a0-a7)
        available_regs.insert(
            RegClassKind::GPR,
            (5..8).chain(10..18).chain(28..32).collect(),
        );
        // Default FPRs: f0-f7, f10-f17, f28-f31 (ft0-ft7, fa0-fa7, ft8-ft11)
        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(),
        );
        // Default VRs: v0-v31
        available_regs.insert(RegClassKind::VecReg, (0..32).collect());

        // Reserved: x0 (zero), x1 (ra), x2 (sp), x3 (gp), x4 (tp), x8 (s0/fp)
        let mut reserved_regs = HashSet::new();
        reserved_regs.extend(&[0, 1, 2, 3, 4, 8]);
        // Reserved FPRs: none by default

        Self {
            available_regs,
            reserved_regs,
            fixed_regs: HashMap::new(),
            intervals: HashMap::new(),
            assignments: HashMap::new(),
            spill_slots: HashMap::new(),
            next_spill_slot: 0,
            coalesced_count: 0,
            spill_count: 0,
            allocation_count: 0,
            stats: LinearScanStats::default(),
        }
    }

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

    /// Mark a physical register as reserved (not allocatable).
    pub fn reserve_reg(&mut self, reg: u32) {
        self.reserved_regs.insert(reg);
    }

    /// Add a fixed (pre-colored) register mapping.
    pub fn add_fixed_reg(&mut self, vreg: u32, phys_reg: u32, reg_class: RegClassKind) {
        self.fixed_regs.insert(vreg, phys_reg);
        let interval = LinearLiveInterval::from_fixed(vreg, phys_reg, reg_class);
        self.intervals.insert(vreg, interval);
    }

    /// Build a live interval for a virtual register.
    pub fn build_interval(
        &mut self,
        vreg: u32,
        reg_class: RegClassKind,
    ) -> &mut LinearLiveInterval {
        self.intervals
            .entry(vreg)
            .or_insert_with(|| LinearLiveInterval::new(vreg, reg_class))
    }

    /// Record a definition of a virtual register at an instruction point.
    pub fn record_def(&mut self, vreg: u32, point: InstrPoint) {
        if let Some(interval) = self.intervals.get_mut(&vreg) {
            interval.add_def(point);
            // Extend the live range from this def
            if interval.segments.is_empty() {
                interval.add_segment(point, point);
            } else {
                let last = interval.segments.last_mut().unwrap();
                if point > last.end {
                    last.end = point;
                }
            }
        }
    }

    /// Record a use of a virtual register at an instruction point.
    pub fn record_use(&mut self, vreg: u32, point: InstrPoint) {
        if let Some(interval) = self.intervals.get_mut(&vreg) {
            interval.add_use(point);
            // Extend the live range to cover this use
            if let Some(last) = interval.segments.last_mut() {
                if point > last.end {
                    last.end = point;
                }
            }
        }
    }

    /// Compute all live intervals by scanning definitions and uses.
    pub fn compute_intervals(
        &mut self,
        defs: &HashMap<u32, Vec<InstrPoint>>,
        uses: &HashMap<u32, Vec<InstrPoint>>,
    ) {
        for (vreg, def_points) in defs {
            let class = self.infer_reg_class(*vreg);
            let interval = self.build_interval(*vreg, class);
            for &def in def_points {
                interval.add_def(def);
            }
        }
        for (vreg, use_points) in uses {
            if let Some(interval) = self.intervals.get_mut(vreg) {
                for &u in use_points {
                    interval.add_use(u);
                }
            }
        }
        // Build segments from def->last-use
        for interval in self.intervals.values_mut() {
            if interval.def_points.is_empty() {
                continue;
            }
            interval.segments.clear();
            let defs_copy: Vec<InstrPoint> = interval.def_points.clone();
            for &def in &defs_copy {
                // Find the furthest use after this def
                let furthest_use = interval
                    .use_points
                    .iter()
                    .filter(|&&u| u >= def)
                    .max()
                    .copied()
                    .unwrap_or(def);
                interval.add_segment(def, furthest_use);
            }
            interval.sort_and_merge();
        }
    }

    /// Infer register class from virtual register ID.
    fn infer_reg_class(&self, vreg: u32) -> RegClassKind {
        if vreg >= 64 && vreg < 96 {
            RegClassKind::VecReg
        } else if vreg >= 32 && vreg < 64 {
            RegClassKind::FPR64
        } else {
            RegClassKind::GPR
        }
    }

    /// Main allocation entry point.
    pub fn allocate(&mut self) -> LinearScanStats {
        self.stats = LinearScanStats::default();
        self.stats.total_intervals = self.intervals.len();

        // Collect unassigned intervals, sort by start point
        let mut unhandled: Vec<LinearLiveInterval> = self
            .intervals
            .values()
            .filter(|i| i.assigned_reg.is_none() && !i.is_fixed)
            .cloned()
            .collect();

        unhandled.sort_by(|a, b| {
            a.start_point().cmp(&b.start_point()).then_with(|| {
                b.spill_weight
                    .partial_cmp(&a.spill_weight)
                    .unwrap_or(Ordering::Equal)
            })
        });

        let mut active: Vec<ActiveEntry> = Vec::new();
        let mut spilled: Vec<LinearLiveInterval> = Vec::new();

        for mut interval in unhandled {
            let start = interval.start_point().unwrap_or(InstrPoint::new(0, 0, 0));

            // Expire finished intervals from the active list
            self.expire_active_intervals(&mut active, start);

            let class = interval.reg_class;
            let available = self.get_available_for_class(class);

            // Try to find a free register
            let free_reg = self.find_free_register(&active, &available);

            if let Some(reg) = free_reg {
                // If we have a hint and it's free, prefer it
                let assign_reg = if let Some(hint) = interval.reg_hint {
                    if !active.iter().any(|a| a.phys_reg == hint) && available.contains(&hint) {
                        hint
                    } else {
                        reg
                    }
                } else {
                    reg
                };

                interval.assigned_reg = Some(assign_reg);
                self.assignments.insert(interval.vreg, Some(assign_reg));
                self.stats.registers_assigned += 1;
                let end = interval.end_point().unwrap_or(start);
                active.push(ActiveEntry::new(interval.clone(), assign_reg, end));
            } else {
                // No free register — need to spill
                self.spill_interval(&mut interval, &mut active, start);
                spilled.push(interval.clone());
            }

            self.allocation_count += 1;
        }

        // Handle spilled intervals: assign spill slots
        for mut interval in spilled {
            self.assign_spill_slot(&mut interval);
            self.stats.registers_spilled += 1;
        }

        self.stats.clone()
    }

    /// Expire intervals that have ended before the given point.
    fn expire_active_intervals(&self, active: &mut Vec<ActiveEntry>, point: InstrPoint) {
        active.retain(|entry| {
            // Keep if the interval is still live at or after 'point'
            entry.end_point >= point
        });
    }

    /// Find a free register from the available set that is not currently active.
    fn find_free_register(&self, active: &[ActiveEntry], available: &[u32]) -> Option<u32> {
        let used: HashSet<u32> = active.iter().map(|a| a.phys_reg).collect();
        available.iter().find(|r| !used.contains(r)).copied()
    }

    /// Get available registers for a class, excluding reserved ones.
    fn get_available_for_class(&self, class: RegClassKind) -> Vec<u32> {
        self.available_regs
            .get(&class)
            .cloned()
            .unwrap_or_default()
            .into_iter()
            .filter(|r| !self.reserved_regs.contains(r))
            .collect()
    }

    /// Select which interval to spill and spill it.
    fn spill_interval(
        &mut self,
        current: &mut LinearLiveInterval,
        active: &mut Vec<ActiveEntry>,
        at_point: InstrPoint,
    ) {
        // Find the active interval whose last use is furthest away
        if active.is_empty() {
            // No active intervals to evict — spill current
            current.spilled = true;
            self.spill_count += 1;
            return;
        }

        // Select the candidate with the furthest next use (cheapest to spill later)
        let mut spill_idx = 0;
        let mut furthest_point = InstrPoint::new(0, 0, 0);

        for (i, entry) in active.iter().enumerate() {
            let last_use = entry
                .interval
                .furthest_use_from(at_point)
                .unwrap_or(entry.end_point);
            if last_use > furthest_point {
                furthest_point = last_use;
                spill_idx = i;
            }
        }

        let spill_entry = &active[spill_idx];

        // If the current interval ends before the spillee's next use, spill current
        let current_end = current.end_point().unwrap_or(at_point);
        if current_end < furthest_point {
            current.spilled = true;
            self.spill_count += 1;
            return;
        }

        // Otherwise, evict the spillee and assign its register to current
        let evicted_reg = spill_entry.phys_reg;
        let mut evicted_interval = active.remove(spill_idx).interval;
        evicted_interval.spilled = true;
        evicted_interval.assigned_reg = None;
        self.assignments.insert(evicted_interval.vreg, None);

        current.assigned_reg = Some(evicted_reg);
        self.assignments.insert(current.vreg, Some(evicted_reg));
        self.stats.registers_assigned += 1;
        self.spill_count += 1;

        let end = current.end_point().unwrap_or(at_point);
        active.push(ActiveEntry::new(current.clone(), evicted_reg, end));
    }

    /// Assign a stack slot to a spilled virtual register.
    fn assign_spill_slot(&mut self, interval: &mut LinearLiveInterval) {
        if self.spill_slots.contains_key(&interval.vreg) {
            return;
        }
        let slot = self.next_spill_slot;
        self.next_spill_slot += 8; // 8-byte aligned slots
        self.spill_slots.insert(interval.vreg, slot);
        interval.spill_offset = slot;

        // Insert store after each def
        for &def in &interval.def_points {
            self.stats.store_after_def_inserted += 1;
            // In a real implementation: insert STORE vreg, [sp+offset] after def
            // For now, track the statistics
        }

        // Insert load before each use
        for &u in &interval.use_points {
            self.stats.load_before_use_inserted += 1;
            // In a real implementation: insert LOAD vreg, [sp+offset] before use
            // For now, track the statistics
        }
    }

    /// Attempt coalescing: if a COPY instruction copies src->dst, assign same reg.
    pub fn try_coalesce(&mut self, dst: u32, src: u32) -> bool {
        // Extract needed info first to avoid borrow conflicts
        let src_reg = self.intervals.get(&src).and_then(|i| i.assigned_reg);
        let dst_reg = self.intervals.get(&dst).and_then(|i| i.assigned_reg);
        let src_overlaps_dst =
            if let (Some(si), Some(di)) = (self.intervals.get(&src), self.intervals.get(&dst)) {
                si.overlaps_with(di)
            } else {
                true
            };

        if src_overlaps_dst {
            return false;
        }

        // Try setting src's hint to dst's register
        if let Some(reg) = dst_reg {
            if let Some(src_mut) = self.intervals.get_mut(&src) {
                src_mut.set_hint(reg);
                self.coalesced_count += 1;
                self.stats.copies_coalesced += 1;
                return true;
            }
        }

        // Try setting dst's hint to src's register
        if let Some(reg) = src_reg {
            if let Some(dst_mut) = self.intervals.get_mut(&dst) {
                dst_mut.set_hint(reg);
                self.coalesced_count += 1;
                self.stats.copies_coalesced += 1;
                return true;
            }
        }

        false
    }

    /// Handle two-address instructions where dest must equal first source.
    pub fn handle_two_address(&mut self, dst: u32, src: u32) -> bool {
        // If both are already assigned different registers, need to copy
        if let (Some(dst_assign), Some(src_assign)) = (
            self.assignments.get(&dst).copied().flatten(),
            self.assignments.get(&src).copied().flatten(),
        ) {
            if dst_assign != src_assign {
                // Try to reassign dst to match src
                let src_interval = self.intervals.get(&src);
                let dst_interval = self.intervals.get(&dst);
                if let (Some(si), Some(di)) = (src_interval, dst_interval) {
                    if !si.overlaps_with(di) || di.start_point() > si.end_point() {
                        self.assignments.insert(dst, Some(src_assign));
                        if let Some(dst_mut) = self.intervals.get_mut(&dst) {
                            dst_mut.assigned_reg = Some(src_assign);
                        }
                        self.stats.two_address_resolved += 1;
                        return true;
                    }
                }
            }
        }
        false
    }

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

    /// Check if a virtual register was spilled.
    pub fn is_spilled(&self, vreg: u32) -> bool {
        self.spill_slots.contains_key(&vreg)
    }

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

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

    /// Register a COPY hint for coalescing.
    pub fn add_copy_hint(&mut self, dst: u32, src_vreg: u32) {
        if let Some(src) = self.intervals.get(&src_vreg) {
            if let Some(reg) = src.assigned_reg.or(src.fixed_reg) {
                if let Some(dst_interval) = self.intervals.get_mut(&dst) {
                    dst_interval.set_hint(reg);
                }
            }
        }
    }
}

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

// ============================================================================
// Interval Comparator (for sorting by start point)
// ============================================================================

/// Comparator for ordering intervals by start point then spill weight.
struct IntervalByStart(LinearLiveInterval);

impl PartialEq for IntervalByStart {
    fn eq(&self, other: &Self) -> bool {
        self.0.start_point() == other.0.start_point()
    }
}

impl Eq for IntervalByStart {}

impl PartialOrd for IntervalByStart {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl Ord for IntervalByStart {
    fn cmp(&self, other: &Self) -> Ordering {
        self.0
            .start_point()
            .cmp(&other.0.start_point())
            .then_with(|| {
                other
                    .0
                    .spill_weight
                    .partial_cmp(&self.0.spill_weight)
                    .unwrap_or(Ordering::Equal)
            })
    }
}

// ============================================================================
// Spill Code Insertion Context
// ============================================================================

/// Context for inserting spill/reload code into a machine function.
#[derive(Debug, Clone)]
pub struct SpillCodeContext {
    /// Map from vreg to stack slot offset.
    pub spill_slots: HashMap<u32, i32>,
    /// Spill code inserted: list of (position, type, vreg).
    pub spill_events: Vec<SpillEvent>,
    /// Stack frame size needed for spill slots.
    pub frame_size: i32,
}

/// A spill or reload event.
#[derive(Debug, Clone)]
pub enum SpillEvent {
    /// Store vreg to [sp + offset] after instruction at point.
    Store {
        vreg: u32,
        offset: i32,
        point: InstrPoint,
    },
    /// Load vreg from [sp + offset] before instruction at point.
    Load {
        vreg: u32,
        offset: i32,
        point: InstrPoint,
    },
}

impl SpillCodeContext {
    pub fn new() -> Self {
        Self {
            spill_slots: HashMap::new(),
            spill_events: Vec::new(),
            frame_size: 0,
        }
    }

    /// Add a store-after-def event.
    pub fn add_store(&mut self, vreg: u32, offset: i32, point: InstrPoint) {
        self.spill_events.push(SpillEvent::Store {
            vreg,
            offset,
            point,
        });
    }

    /// Add a load-before-use event.
    pub fn add_load(&mut self, vreg: u32, offset: i32, point: InstrPoint) {
        self.spill_events.push(SpillEvent::Load {
            vreg,
            offset,
            point,
        });
    }

    /// Calculate required frame size from spill slots.
    pub fn compute_frame_size(&mut self) {
        let max_offset = self.spill_slots.values().max().copied().unwrap_or(0);
        self.frame_size = max_offset + 8; // Align to 8 bytes
    }

    /// Sort spill events by position for insertion.
    pub fn sort_events(&mut self) {
        self.spill_events.sort_by(|a, b| {
            let pa = match a {
                SpillEvent::Store { point, .. } | SpillEvent::Load { point, .. } => *point,
            };
            let pb = match b {
                SpillEvent::Store { point, .. } | SpillEvent::Load { point, .. } => *point,
            };
            pa.cmp(&pb)
        });
    }
}

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

// ============================================================================
// Linear Scan Integration (Public API)
// ============================================================================

/// High-level linear scan register allocation for a machine function.
pub struct LinearScanRegAlloc {
    pub allocator: LinearScanAllocator,
    pub function_name: String,
    pub spill_code: SpillCodeContext,
    pub success: bool,
}

impl LinearScanRegAlloc {
    pub fn new(function_name: &str) -> Self {
        Self {
            allocator: LinearScanAllocator::new(),
            function_name: function_name.to_string(),
            spill_code: SpillCodeContext::new(),
            success: false,
        }
    }

    /// Run full linear scan allocation on a function.
    pub fn run(
        &mut self,
        defs: &HashMap<u32, Vec<InstrPoint>>,
        uses: &HashMap<u32, Vec<InstrPoint>>,
        fixed: &[(u32, u32, RegClassKind)],
        copy_hints: &[(u32, u32)],
        two_addr: &[(u32, u32)],
    ) -> LinearScanStats {
        // Register fixed mappings
        for &(vreg, phys_reg, class) in fixed {
            self.allocator.add_fixed_reg(vreg, phys_reg, class);
        }

        // Compute intervals
        self.allocator.compute_intervals(defs, uses);

        // Add copy hints
        for &(dst, src) in copy_hints {
            self.allocator.try_coalesce(dst, src);
        }

        // Run allocation
        let stats = self.allocator.allocate();

        // Handle two-address constraints
        for &(dst, src) in two_addr {
            self.allocator.handle_two_address(dst, src);
        }

        // Build spill code
        for (&vreg, &offset) in &self.allocator.spill_slots {
            self.spill_code.spill_slots.insert(vreg, offset);
            if let Some(interval) = self.allocator.intervals.get(&vreg) {
                for &def in &interval.def_points {
                    self.spill_code.add_store(vreg, offset, def.after());
                }
                for &u in &interval.use_points {
                    self.spill_code.add_load(vreg, offset, u.before());
                }
            }
        }
        self.spill_code.sort_events();
        self.spill_code.compute_frame_size();

        self.success = true;
        stats
    }

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

    /// Check if a virtual register is spilled.
    pub fn is_spilled(&self, vreg: u32) -> bool {
        self.allocator.is_spilled(vreg)
    }
}

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

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

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

    fn pt_use(block: u32, instr: u32) -> InstrPoint {
        InstrPoint::new(block, instr, 1)
    }

    // --- Existing PBQP tests ---

    #[test]
    fn test_pbqp_vector_new() {
        let regs = vec![0, 1, 2, 3];
        let v = PbqpVector::new(&regs, 10.0);
        assert_eq!(v.num_choices(), 4);
        assert_eq!(v.get_cost(0), Some(10.0));
        assert_eq!(v.get_cost(3), Some(10.0));
        assert!(v.get_cost(99).is_none());
    }

    #[test]
    fn test_pbqp_vector_min_cost() {
        let regs = vec![0, 1, 2, 3];
        let mut v = PbqpVector::new(&regs, 100.0);
        v.set_cost(2, 5.0);
        assert_eq!(v.min_cost_reg(), Some((2, 5.0)));
    }

    #[test]
    fn test_pbqp_vector_has_finite() {
        let regs = vec![0, 1];
        let mut v = PbqpVector::new(&regs, PBQP_INF);
        assert!(!v.has_finite_cost());
        v.set_cost(0, 10.0);
        assert!(v.has_finite_cost());
    }

    #[test]
    fn test_pbqp_matrix_interference() {
        let m = PbqpMatrix::interference_matrix(3);
        // Same register assignments are illegal
        assert!(m.get(0, 0) > PBQP_INF / 2.0);
        assert!(m.get(1, 1) > PBQP_INF / 2.0);
        // Different register assignments are free
        assert!(m.get(0, 1) < 1.0);
        assert!(m.get(1, 0) < 1.0);
    }

    #[test]
    fn test_pbqp_matrix_identity_benefit() {
        let m = PbqpMatrix::identity_with_benefit(3, 10.0);
        assert!((m.get(0, 0) + 10.0).abs() < 0.001);
        assert!((m.get(1, 1) + 10.0).abs() < 0.001);
        assert_eq!(m.get(0, 1), 0.0);
    }

    #[test]
    fn test_pbqp_graph_add_node() {
        let mut g = PbqpGraph::new();
        let node = PbqpNode::new(0, &[0, 1, 2], 0.0);
        g.add_node(node);
        assert!(g.nodes.contains_key(&0));
        assert_eq!(g.degree(0), 0);
    }

    #[test]
    fn test_pbqp_graph_add_edge() {
        let mut g = PbqpGraph::new();
        g.add_node(PbqpNode::new(0, &[0, 1, 2], 0.0));
        g.add_node(PbqpNode::new(1, &[0, 1, 2], 0.0));
        let edge = PbqpEdge::new(0, 1, 3, 3);
        g.add_edge(edge);
        assert_eq!(g.degree(0), 1);
        assert_eq!(g.degree(1), 1);
    }

    #[test]
    fn test_pbqp_solver_r0() {
        let mut solver = PbqpSolver::new();
        solver.graph.add_node(PbqpNode::new(0, &[0, 1], 10.0));
        solver.reduce_graph();
        // R0 should have reduced node 0
        assert!(solver.graph.nodes.is_empty());
    }

    #[test]
    fn test_pbqp_solver_r1() {
        let mut solver = PbqpSolver::new();
        solver.graph.add_node(PbqpNode::new(0, &[0, 1], 10.0));
        solver.graph.add_node(PbqpNode::new(1, &[0, 1], 20.0));
        let edge = PbqpEdge::new(0, 1, 2, 2);
        solver.graph.add_edge(edge);
        solver.reduce_graph();
        // Both nodes should eventually be reduced
        assert!(solver.graph.nodes.is_empty());
    }

    #[test]
    fn test_pbqp_reg_alloc_basic() {
        let mut alloc = PbqpRegAlloc::new(vec![0, 1, 2, 3]);
        alloc.allocate("test_func", 4);
        // Should have assigned most registers
        assert!(!alloc.assignments.is_empty() || !alloc.spilled.is_empty());
    }

    #[test]
    fn test_pbqp_solver_solve() {
        let mut solver = PbqpSolver::new();
        let vregs = vec![0, 1];
        let allowed = vec![0, 1];
        solver.init_graph(&vregs, &allowed, &[(0, 1)], &[]);
        let solution = solver.solve();
        assert_eq!(solution.len(), 2);
    }

    #[test]
    fn test_pbqp_spill_selection() {
        // With 2 vregs and only 1 allowed register, one must spill
        let mut solver = PbqpSolver::new();
        let vregs = vec![0, 1];
        let allowed = vec![0]; // only one register
        solver.init_graph(&vregs, &allowed, &[(0, 1)], &[]);
        let solution = solver.solve();
        // At least one should be spilled
        let spilled_count = solution.values().filter(|s| s.spilled).count();
        assert!(spilled_count >= 1);
    }

    // --- Linear Scan RA Tests ---

    #[test]
    fn test_linear_scan_allocator_new() {
        let alloc = LinearScanAllocator::new();
        assert!(alloc.available_regs.contains_key(&RegClassKind::GPR));
        assert!(alloc.available_regs.contains_key(&RegClassKind::FPR64));
    }

    #[test]
    fn test_linear_scan_build_interval() {
        let mut alloc = LinearScanAllocator::new();
        let iv = alloc.build_interval(1, RegClassKind::GPR);
        assert_eq!(iv.vreg, 1);
        assert!(!iv.is_fixed);
    }

    #[test]
    fn test_linear_scan_fixed_register() {
        let mut alloc = LinearScanAllocator::new();
        alloc.add_fixed_reg(10, 5, RegClassKind::GPR);
        assert!(alloc.fixed_regs.contains_key(&10));
        if let Some(iv) = alloc.intervals.get(&10) {
            assert!(iv.is_fixed);
            assert_eq!(iv.fixed_reg, Some(5));
        }
    }

    #[test]
    fn test_linear_scan_compute_intervals() {
        let mut alloc = LinearScanAllocator::new();
        let mut defs = HashMap::new();
        let mut uses = HashMap::new();

        defs.insert(1, vec![InstrPoint::new(0, 0, 0)]);
        uses.insert(1, vec![InstrPoint::new(0, 5, 1)]);

        alloc.compute_intervals(&defs, &uses);

        if let Some(iv) = alloc.intervals.get(&1) {
            assert!(!iv.segments.is_empty());
            assert_eq!(iv.def_points.len(), 1);
            assert_eq!(iv.use_points.len(), 1);
        }
    }

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

        assert!(iv.live_at(InstrPoint::new(0, 5, 0)));
        assert!(!iv.live_at(InstrPoint::new(0, 15, 0)));
    }

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

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

        assert!(iv1.overlaps_with(&iv2));

        let mut iv3 = LinearLiveInterval::new(3, RegClassKind::GPR);
        iv3.add_segment(InstrPoint::new(0, 20, 0), InstrPoint::new(0, 30, 1));

        assert!(!iv1.overlaps_with(&iv3));
    }

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

        // Should merge into one segment [0,10]
        assert_eq!(iv.segments.len(), 1);
        assert_eq!(iv.segments[0].start, InstrPoint::new(0, 0, 0));
        assert_eq!(iv.segments[0].end, InstrPoint::new(0, 10, 1));
    }

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

        let mut defs = HashMap::new();
        let mut uses = HashMap::new();

        defs.insert(1, vec![InstrPoint::new(0, 0, 0)]);
        uses.insert(1, vec![InstrPoint::new(0, 2, 1)]);

        defs.insert(2, vec![InstrPoint::new(0, 3, 0)]);
        uses.insert(2, vec![InstrPoint::new(0, 5, 1)]);

        alloc.compute_intervals(&defs, &uses);
        alloc.allocate();

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

    #[test]
    fn test_linear_scan_spill_under_pressure() {
        let mut alloc = LinearScanAllocator::new();
        alloc.available_regs.insert(RegClassKind::GPR, vec![10]); // only 1 GPR

        let mut defs = HashMap::new();
        let mut uses = HashMap::new();

        for i in 0..5 {
            defs.insert(i, vec![InstrPoint::new(0, i * 2, 0)]);
            uses.insert(i, vec![InstrPoint::new(0, i * 2 + 1, 1)]);
        }

        alloc.compute_intervals(&defs, &uses);
        alloc.allocate();

        // Some vregs should be spilled
        let spilled_count = alloc.assignments.values().filter(|a| a.is_none()).count();
        assert!(spilled_count > 0 || alloc.stats.registers_spilled > 0);
    }

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

        let mut defs = HashMap::new();
        let mut uses = HashMap::new();

        defs.insert(1, vec![InstrPoint::new(0, 0, 0)]);
        uses.insert(1, vec![InstrPoint::new(0, 1, 1)]);

        defs.insert(2, vec![InstrPoint::new(0, 2, 0)]);
        uses.insert(2, vec![InstrPoint::new(0, 3, 1)]);

        alloc.compute_intervals(&defs, &uses);

        // Pre-assign vreg 1 to reg 10
        if let Some(iv) = alloc.intervals.get_mut(&1) {
            iv.assigned_reg = Some(10);
        }

        // Try to coalesce vreg 2 with vreg 1
        alloc.try_coalesce(2, 1);

        // vreg 2 should have hint set to 10
        if let Some(iv) = alloc.intervals.get(&2) {
            assert_eq!(iv.reg_hint, Some(10));
        }
    }

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

        let mut defs = HashMap::new();
        let mut uses = HashMap::new();

        defs.insert(1, vec![InstrPoint::new(0, 0, 0)]);
        uses.insert(1, vec![InstrPoint::new(0, 1, 1)]);

        defs.insert(2, vec![InstrPoint::new(0, 2, 0)]);
        uses.insert(2, vec![InstrPoint::new(0, 3, 1)]);

        alloc.compute_intervals(&defs, &uses);
        alloc.assignments.insert(1, Some(10));
        alloc.assignments.insert(2, Some(11));

        if let Some(iv) = alloc.intervals.get_mut(&1) {
            iv.assigned_reg = Some(10);
        }
        if let Some(iv) = alloc.intervals.get_mut(&2) {
            iv.assigned_reg = Some(11);
        }

        // Two-address: dst=2, src=1
        alloc.handle_two_address(2, 1);

        // Should have tried to resolve
        assert!(alloc.stats.two_address_resolved >= 0);
    }

    #[test]
    fn test_linear_scan_reg_class_constraint() {
        let mut alloc = LinearScanAllocator::new();
        // GPR-specific allocation
        let class = RegClassKind::GPR;
        alloc.available_regs.insert(class, vec![10, 11]);

        let mut defs = HashMap::new();
        let mut uses = HashMap::new();
        defs.insert(1, vec![InstrPoint::new(0, 0, 0)]);
        uses.insert(1, vec![InstrPoint::new(0, 1, 1)]);

        alloc.compute_intervals(&defs, &uses);
        alloc.allocate();

        if let Some(assignment) = alloc.get_assignment(1) {
            // Should be one of the allowed GPRs
            let available = alloc.available_regs.get(&RegClassKind::GPR).unwrap();
            assert!(available.contains(&assignment));
        }
    }

    #[test]
    fn test_linear_scan_spill_code_context() {
        let mut ctx = SpillCodeContext::new();
        let point = InstrPoint::new(0, 5, 0);

        ctx.add_store(1, 0, point.after());
        ctx.add_load(1, 0, point.before());

        assert_eq!(ctx.spill_events.len(), 2);

        ctx.sort_events();
        ctx.compute_frame_size();
        assert!(ctx.frame_size >= 0);
    }

    #[test]
    fn test_linear_scan_reg_alloc_run() {
        let mut alloc = LinearScanRegAlloc::new("test");
        let mut defs = HashMap::new();
        let mut uses = HashMap::new();

        defs.insert(1, vec![InstrPoint::new(0, 0, 0)]);
        uses.insert(1, vec![InstrPoint::new(0, 2, 1)]);

        let fixed = &[];
        let hints = &[];
        let two_addr = &[];

        let stats = alloc.run(&defs, &uses, fixed, hints, two_addr);
        assert!(alloc.success);
        assert!(stats.total_intervals >= 1);
    }

    #[test]
    fn test_linear_scan_instr_point_ordering() {
        let p1 = InstrPoint::new(0, 0, 0);
        let p2 = InstrPoint::new(0, 0, 1);
        let p3 = InstrPoint::new(0, 1, 0);
        let p4 = InstrPoint::new(1, 0, 0);

        assert!(p1 < p2);
        assert!(p2 < p3);
        assert!(p3 < p4);
    }

    #[test]
    fn test_linear_scan_instr_point_before_after() {
        let p = InstrPoint::new(0, 5, 0);
        assert!(p.before() <= p);
        assert!(p.after() >= p);
    }

    #[test]
    fn test_linear_scan_reserved_regs_not_used() {
        let mut alloc = LinearScanAllocator::new();
        // Only make reserved registers available
        alloc.available_regs.insert(RegClassKind::GPR, vec![0, 2]); // x0, x2 are reserved

        let mut defs = HashMap::new();
        let mut uses = HashMap::new();
        defs.insert(1, vec![InstrPoint::new(0, 0, 0)]);
        uses.insert(1, vec![InstrPoint::new(0, 1, 1)]);

        alloc.compute_intervals(&defs, &uses);
        alloc.allocate();

        // Should not assign reserved registers
        if let Some(Some(assigned)) = alloc.assignments.get(&1) {
            assert!(!alloc.reserved_regs.contains(assigned));
        }
    }

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

        iv.compute_spill_weight(0); // depth 0
        assert!(iv.spill_weight > 0.0);

        iv.compute_spill_weight(3); // depth 3
        assert!(iv.spill_weight > 0.0);
    }

    #[test]
    fn test_linear_scan_active_list_expiry() {
        let mut alloc = LinearScanAllocator::new();
        let point = InstrPoint::new(0, 10, 0);

        let mut active = vec![
            ActiveEntry::new(
                LinearLiveInterval::new(1, RegClassKind::GPR),
                10,
                InstrPoint::new(0, 5, 1), // ended
            ),
            ActiveEntry::new(
                LinearLiveInterval::new(2, RegClassKind::GPR),
                11,
                InstrPoint::new(0, 15, 1), // still live
            ),
        ];

        alloc.expire_active_intervals(&mut active, point);
        assert_eq!(active.len(), 1);
    }
}

// Re-export RegAllocResult from codegen for downstream convenience.
pub use llvm_native_core::codegen::RegAllocResult;