dotscope 0.9.1

//! Local variable coalescing for SSA code generation.
//!
//! This module implements register allocation for SSA variables, allowing
//! non-interfering variables (those never live at the same program point)
//! to share the same local slot, reducing code size.
//!
//! # Algorithms
//!
//! Two allocation algorithms are available, selected automatically based on
//! the number of variables:
//!
//! ## Graph Coloring (for small SSAs, ≤500 variables)
//! - **Optimal results**: Minimizes the number of local slots
//! - **O(n²) complexity**: Builds full interference graph
//! - Used by GCC, LLVM for ahead-of-time compilation
//!
//! ## Linear Scan (for large SSAs, >500 variables)
//! - **Near-optimal results**: Within 5-10% of graph coloring
//! - **O(n log n) complexity**: Uses live intervals instead of interference graph
//! - Used by HotSpot JVM, V8, .NET RyuJIT for JIT compilation
//!
//! # Type Compatibility
//!
//! Variables can only share slots if their types are compatible:
//! - Same type: always compatible
//! - All reference types: share object slots
//! - Integer types of same size: compatible (I32/U32, I64/U64)

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

use analyssa::BitSet;
use rayon::prelude::*;

use crate::{
    analysis::{
        AnalysisResults, DataFlowSolver, EhCfg, LiveVariables, LivenessResult, SsaCfg, SsaFunction,
        SsaOp, SsaType, SsaVarId, VariableOrigin,
    },
    Error, Result,
};

/// Interference graph for register allocation.
///
/// Tracks which variables interfere (are live simultaneously) and thus
/// cannot share the same local slot.
pub struct InterferenceGraph {
    /// Adjacency list: var_id -> bitset of interfering var_ids (indexed by SsaVarId::index())
    edges: BTreeMap<SsaVarId, BitSet>,
    /// Type of each variable (for compatibility checking)
    var_types: BTreeMap<SsaVarId, SsaType>,
    /// Number of variables (capacity for BitSets)
    var_count: usize,
}

impl InterferenceGraph {
    /// Creates a new empty interference graph with the given variable capacity.
    fn new(var_count: usize) -> Self {
        Self {
            edges: BTreeMap::new(),
            var_types: BTreeMap::new(),
            var_count,
        }
    }

    /// Adds an interference edge between two variables.
    fn add_edge(&mut self, a: SsaVarId, b: SsaVarId) {
        if a != b {
            let var_count = self.var_count;
            self.edges
                .entry(a)
                .or_insert_with(|| BitSet::new(var_count))
                .insert(b.index());
            self.edges
                .entry(b)
                .or_insert_with(|| BitSet::new(var_count))
                .insert(a.index());
        }
    }

    /// Records the type of a variable.
    fn set_type(&mut self, var: SsaVarId, ty: SsaType) {
        self.var_types.insert(var, ty);
    }

    /// Returns the neighbors (interfering variables) of a variable.
    fn neighbors(&self, var: SsaVarId) -> impl Iterator<Item = SsaVarId> + '_ {
        self.edges
            .get(&var)
            .into_iter()
            .flat_map(|set| set.iter().map(SsaVarId::from_index))
    }

    /// Returns the degree (number of interferences) of a variable.
    fn degree(&self, var: SsaVarId) -> usize {
        self.edges.get(&var).map_or(0, |s| s.count())
    }
}

/// Result of local allocation.
pub struct LocalAllocation {
    /// Map SSA var -> allocated local slot
    pub var_to_local: BTreeMap<SsaVarId, u16>,
    /// Number of locals needed
    pub num_locals: u16,
    /// Map original local index -> compacted local index (for signature generation)
    pub original_to_compacted: BTreeMap<u16, u16>,
}

/// Live interval for a variable (used by linear scan).
///
/// Represents the range [start, end) where a variable is live.
#[derive(Debug, Clone)]
struct LiveInterval {
    start: usize,
    end: usize,
}

impl LiveInterval {
    fn new(pos: usize) -> Self {
        Self {
            start: pos,
            end: pos.saturating_add(1),
        }
    }

    fn extend_start(&mut self, pos: usize) {
        self.start = self.start.min(pos);
    }

    fn extend_end(&mut self, pos: usize) {
        self.end = self.end.max(pos);
    }
}

/// Local variable coalescer for SSA code generation.
///
/// Automatically selects between graph coloring (optimal, O(n²)) and
/// linear scan (near-optimal, O(n log n)) based on the number of variables.
pub struct LocalCoalescer {
    /// Interference graph (used by graph coloring algorithm)
    interference: InterferenceGraph,
    /// Variables that need allocation (Stack and Phi origins only)
    coalescable_vars: Vec<SsaVarId>,
    /// Pre-computed allocation from linear scan (if used)
    precomputed: Option<LocalAllocation>,
}

/// Threshold for switching from graph coloring to linear scan.
///
/// Graph coloring has O(n²) complexity for building the interference graph.
/// For SSAs with more variables than this threshold, we use linear scan
/// which has O(n log n) complexity and produces near-optimal results.
const LINEAR_SCAN_THRESHOLD: usize = 500;

impl LocalCoalescer {
    /// Builds a coalescer from the SSA function.
    ///
    /// Automatically selects the appropriate algorithm:
    /// - **Graph coloring** for small SSAs (≤500 variables): Optimal results, O(n²)
    /// - **Linear scan** for large SSAs (>500 variables): Near-optimal, O(n log n)
    pub fn build(ssa: &SsaFunction) -> Self {
        if ssa.variable_count() > LINEAR_SCAN_THRESHOLD {
            return Self::build_linear_scan(ssa);
        }
        Self::build_graph_coloring(ssa)
    }

    /// Builds using graph coloring algorithm.
    ///
    /// Uses liveness analysis to determine which variables are live at each
    /// program point, then builds edges between variables that are live
    /// simultaneously. Produces optimal slot allocation.
    ///
    /// Complexity: O(n²) for n variables
    fn build_graph_coloring(ssa: &SsaFunction) -> Self {
        let var_count = ssa.var_id_capacity();
        let mut interference = InterferenceGraph::new(var_count);

        // Record types for all variables
        for var in ssa.variables() {
            interference.set_type(var.id(), var.var_type().clone());
        }

        // Build CFG for dataflow analysis.
        //
        // The solver must run over the exception-aware view: a traversal seeded
        // from terminator edges alone never reaches a handler block, so every
        // variable live only across a protected region would come back dead.
        // The extra edges only widen liveness, which is a may-analysis, so a
        // value reported live across an exceptional edge is conservative.
        let cfg = EhCfg::from_ssa(ssa);

        // Run liveness analysis
        let analysis = LiveVariables::new(ssa);
        let solver = DataFlowSolver::new(analysis);
        let results = solver.solve(ssa, &cfg);

        // Collect block IDs for parallel processing
        let block_ids: Vec<usize> = (0..ssa.block_count()).collect();

        // Phase 1: Collect boundary edges in parallel
        // Each block independently computes its interference edges
        let boundary_edges: Vec<(SsaVarId, SsaVarId)> = block_ids
            .par_iter()
            .flat_map(|&block_id| {
                let mut edges = Vec::new();

                // Edges from live_out
                if let Some(live_out) = results.out_state(block_id) {
                    let live_vars: Vec<SsaVarId> = live_out.variables().collect();
                    for (i, &var1) in live_vars.iter().enumerate() {
                        if let Some(rest) = live_vars.get(i.saturating_add(1)..) {
                            for &var2 in rest {
                                edges.push((var1, var2));
                            }
                        }
                    }
                }

                // Edges from live_in
                if let Some(live_in) = results.in_state(block_id) {
                    let live_vars: Vec<SsaVarId> = live_in.variables().collect();
                    for (i, &var1) in live_vars.iter().enumerate() {
                        if let Some(rest) = live_vars.get(i.saturating_add(1)..) {
                            for &var2 in rest {
                                edges.push((var1, var2));
                            }
                        }
                    }
                }

                edges
            })
            .collect();

        // Phase 2: Collect intra-block edges in parallel
        let intra_block_edges: Vec<(SsaVarId, SsaVarId)> = block_ids
            .par_iter()
            .flat_map(|&block_id| Self::collect_intra_block_edges(ssa, &results, block_id))
            .collect();

        // Phase 3: Collect PHI edge interference
        // PHI operands from the same predecessor edge are all live at that edge,
        // so they must interfere. This is critical for correct parallel copy semantics.
        let phi_edge_edges: Vec<(SsaVarId, SsaVarId)> = ssa
            .blocks()
            .par_iter()
            .flat_map(|block| {
                let mut edges = Vec::new();

                // Group PHI operands by predecessor
                let mut operands_by_pred: BTreeMap<usize, Vec<SsaVarId>> = BTreeMap::new();
                for phi in block.phi_nodes() {
                    for operand in phi.operands() {
                        operands_by_pred
                            .entry(operand.predecessor())
                            .or_default()
                            .push(operand.value());
                    }
                }

                // All operands from the same predecessor interfere with each other
                for (_, operands) in operands_by_pred {
                    for (i, &var1) in operands.iter().enumerate() {
                        if let Some(rest) = operands.get(i.saturating_add(1)..) {
                            for &var2 in rest {
                                edges.push((var1, var2));
                            }
                        }
                    }
                }

                edges
            })
            .collect();

        // Phase 4: Collect cross-subtree interference edges
        //
        // The codegen's dependency-driven scheduler processes root instructions
        // (terminators, side-effect ops, dead computations) by generating their
        // complete dependency subtrees in sequence. This can reorder operations
        // relative to the original SSA instruction order. Variables from different
        // root subtrees must interfere to prevent the scheduler from causing
        // slot conflicts when it interleaves subtrees.
        let cross_subtree_edges: Vec<(SsaVarId, SsaVarId)> = block_ids
            .par_iter()
            .flat_map(|&block_id| Self::collect_cross_subtree_edges(ssa, block_id))
            .collect();

        // Phase 5: Merge all edges into the interference graph
        for (a, b) in boundary_edges {
            interference.add_edge(a, b);
        }
        for (a, b) in intra_block_edges {
            interference.add_edge(a, b);
        }
        for (a, b) in phi_edge_edges {
            interference.add_edge(a, b);
        }
        for (a, b) in cross_subtree_edges {
            interference.add_edge(a, b);
        }

        // Collect variables that need allocation (Phi-origin variables are freely
        // coalescable, while original locals and arguments keep their fixed slots).
        let coalescable_vars: Vec<SsaVarId> = ssa
            .variables()
            .iter()
            .filter_map(|v| match v.origin() {
                VariableOrigin::Phi => Some(v.id()),
                // An entry live-in is caller-supplied like an argument, so it
                // keeps fixed storage rather than being freely coalescable.
                // Unreachable from the CIL front end, which never builds one.
                VariableOrigin::Argument(_)
                | VariableOrigin::Local(_)
                | VariableOrigin::EntryLiveIn => None,
            })
            .collect();

        Self {
            interference,
            coalescable_vars,
            precomputed: None,
        }
    }

    /// Builds using linear scan algorithm.
    ///
    /// Computes live intervals for each variable and allocates slots in a single
    /// pass. Produces near-optimal results (within 5-10% of graph coloring) with
    /// much better performance for large SSAs.
    ///
    /// Complexity: O(n log n) for n variables
    fn build_linear_scan(ssa: &SsaFunction) -> Self {
        // Collect variable types for compatibility checking
        let mut var_types: BTreeMap<SsaVarId, SsaType> = BTreeMap::new();
        for var in ssa.variables() {
            var_types.insert(var.id(), var.var_type().clone());
        }

        // Compute live intervals for each variable
        let intervals = Self::compute_live_intervals(ssa);

        // Pre-assign fixed slots for arguments and used original locals.
        // Collect Local-origin variables that have live intervals (are actually used).
        // Phi-origin variables go through linear scan allocation separately.
        let mut used_local_vars: Vec<(SsaVarId, u16)> = ssa
            .variables()
            .iter()
            .filter_map(|v| {
                if let VariableOrigin::Local(idx) = v.origin() {
                    if intervals.contains_key(&v.id()) {
                        return Some((v.id(), idx));
                    }
                }
                None
            })
            .collect();
        // Sort by original index for deterministic ordering
        used_local_vars.sort_by_key(|(_, idx)| *idx);

        let PreAssignment {
            mut var_to_local,
            reserved_slots,
            original_to_new,
            mut next_local,
        } = pre_assign_locals(ssa, &used_local_vars);

        // Sort intervals by start position, with (origin, version) as tiebreaker for determinism.
        // The secondary key ensures stable ordering when multiple variables share the same
        // start position.
        let mut sorted_intervals: Vec<_> = intervals.into_iter().collect();
        sorted_intervals.sort_by(|(var_id_a, interval_a), (var_id_b, interval_b)| {
            interval_a
                .start
                .cmp(&interval_b.start)
                .then_with(|| var_sort_key(ssa, *var_id_a).cmp(&var_sort_key(ssa, *var_id_b)))
        });

        // Active list: heap sorted by end position (earliest end first).
        // Heap payload is just `(end, slot)` — we look up the slot's exact
        // type in `slot_types` below when we need it (SsaType isn't `Ord`).
        let mut active: BinaryHeap<Reverse<(usize, u16)>> = BinaryHeap::new();
        // Per-slot exact type (indexed by slot id), used to tag free slots
        // with the precise type they last held so reuse obeys the CIL
        // exact-type-match rule for Reference-class locals.
        let mut slot_type: BTreeMap<u16, SsaType> = BTreeMap::new();

        // Free slots available for reuse, tagged with the exact type they last
        // held. We key this pool by exact type (not just `TypeClass`) so we
        // never recycle e.g. a `bool&` slot for a `string` or `object` value;
        // `is_compatible_for_storage` then governs which types may share.
        // A plain Vec is fine here — allocation count is bounded by live
        // variables, so the linear search on pop stays small in practice.
        let mut free_slots: Vec<(SsaType, u16)> = Vec::new();

        // Linear scan allocation
        for (var_id, interval) in sorted_intervals {
            // Skip if already allocated (arguments, original locals)
            // BUT add them to the active list so their slots aren't reused!
            if var_to_local.contains_key(&var_id) {
                // If this is a Local-origin variable that was pre-assigned,
                // add it to the active list so its slot isn't reused while live
                if let Some(slot) = var_to_local.get(&var_id) {
                    let ty = var_types.get(&var_id).cloned().unwrap_or(SsaType::Unknown);
                    slot_type.insert(*slot, ty);
                    active.push(Reverse((interval.end, *slot)));
                }
                continue;
            }

            // Expire old intervals - return their slots to free pool
            // BUT don't return reserved Local-origin slots to the pool
            while let Some(&Reverse((end, _))) = active.peek() {
                if end > interval.start {
                    break;
                }
                let Some(Reverse((_, slot))) = active.pop() else {
                    break;
                };
                // Only add to free pool if not a reserved slot
                if !reserved_slots.contains(slot as usize) {
                    let ty = slot_type.get(&slot).cloned().unwrap_or(SsaType::Unknown);
                    free_slots.push((ty, slot));
                }
            }

            // Pick this variable's type
            let var_type = var_types.get(&var_id).cloned().unwrap_or(SsaType::Unknown);

            // Try to reuse a free slot whose stored type is compatible with
            // the requesting variable's type. Linear scan is acceptable since
            // the free pool is bounded by concurrently-live variables.
            let slot = {
                let pick = free_slots
                    .iter()
                    .position(|(ty, _)| ty.is_compatible_for_storage(&var_type));
                pick.map(|i| free_slots.swap_remove(i).1)
            };

            let slot = slot.unwrap_or_else(|| {
                let s = next_local;
                next_local = next_local.saturating_add(1);
                s
            });

            var_to_local.insert(var_id, slot);
            slot_type.insert(slot, var_type.clone());

            // Add to active list
            active.push(Reverse((interval.end, slot)));
        }

        let allocation = LocalAllocation {
            var_to_local,
            num_locals: next_local,
            original_to_compacted: original_to_new,
        };

        Self {
            interference: InterferenceGraph::new(0),
            coalescable_vars: Vec::new(),
            precomputed: Some(allocation),
        }
    }

    /// Computes live intervals for all variables in the SSA.
    ///
    /// A live interval is the range [start, end) where a variable is live, in
    /// instruction indices over a linearized view of the CFG.
    ///
    /// # Why this walks the CFG instead of scanning textually or running a dataflow solve
    ///
    /// A purely *textual* scan — `[first mention, last use + 1)` over blocks in `block_id`
    /// order — under-approximates liveness across back edges. A value defined before a loop
    /// and used inside it has its interval truncated at its last textual use, so its slot
    /// returns to the free pool and can be handed to a variable defined later in the same
    /// loop body; on the second iteration the earlier use reads a clobbered slot. Linear
    /// scan only runs above [`LINEAR_SCAN_THRESHOLD`] variables, so that failure is
    /// size-gated and invisible to small unit tests.
    ///
    /// The obvious fix — reuse the `LiveVariables` solve that [`Self::build_graph_coloring`]
    /// runs — is the wrong tool *here*. That framework computes the live set of every
    /// variable at every block simultaneously, storing a live-in and a live-out `BitSet` of
    /// `var_capacity` bits per block: `2·B·V` bits. Graph colouring can afford that only
    /// because it is gated to at most [`LINEAR_SCAN_THRESHOLD`] variables. Linear scan
    /// exists precisely for the large, attacker-sized methods where `B·V` explodes, and the
    /// CIL path has no method-size cap.
    ///
    /// Nothing here needs the full live sets. A [`LiveInterval`] is a single contiguous
    /// range, so all that is required per variable is the earliest and latest point it is
    /// live. SSA gives exactly one definition per variable, and it dominates every use, so
    /// each variable's live range is found by the classic *up-and-mark* walk: start at each
    /// use, walk up predecessors, stop at the defining block — no second definition can be
    /// encountered on the way. Cost is `O(B)` memory for one reusable generation-stamp
    /// array, and time proportional to the total live range, which is the size of the
    /// answer rather than the size of a `B × V` matrix. Most SSA variables are block-local,
    /// so their walk terminates immediately.
    fn compute_live_intervals(ssa: &SsaFunction) -> BTreeMap<SsaVarId, LiveInterval> {
        let block_count = ssa.block_count();
        let mut intervals: BTreeMap<SsaVarId, LiveInterval> = BTreeMap::new();

        // Phase 1: number instructions, recording each block's [start, end) span.
        let mut block_start_idx: Vec<usize> = Vec::with_capacity(block_count);
        let mut block_end_idx: Vec<usize> = Vec::with_capacity(block_count);
        {
            let mut idx = 0usize;
            for block_id in 0..block_count {
                block_start_idx.push(idx);
                if let Some(block) = ssa.block(block_id) {
                    idx = idx.saturating_add(block.instructions().len());
                }
                block_end_idx.push(idx);
            }
        }

        // Phase 2: locate every definition. This must complete before uses are processed,
        // because a use can precede its definition in block order (that is what a back edge
        // means) and the walk below keys its stop condition on the defining block.
        let mut def_block: BTreeMap<SsaVarId, usize> = BTreeMap::new();

        for block_id in 0..block_count {
            let Some(block) = ssa.block(block_id) else {
                continue;
            };
            let Some(&entry_pos) = block_start_idx.get(block_id) else {
                continue;
            };

            // Phi nodes define at block entry.
            for phi in block.phi_nodes() {
                let def = phi.result();
                def_block.insert(def, block_id);
                let interval = intervals
                    .entry(def)
                    .or_insert_with(|| LiveInterval::new(entry_pos));
                interval.extend_start(entry_pos);
                interval.extend_end(entry_pos.saturating_add(1));
            }

            let mut pos = entry_pos;
            for instr in block.instructions() {
                if let Some(def) = instr.def() {
                    def_block.insert(def, block_id);
                    let interval = intervals
                        .entry(def)
                        .or_insert_with(|| LiveInterval::new(pos));
                    interval.extend_start(pos);
                    // Cover the definition point itself. `extend_start` is a min and
                    // `extend_end` a max, so a variable used at 5 and defined at 10 would
                    // otherwise hold [5, 6) — an interval that expires before the value it
                    // describes even exists, freeing the slot while the def is still ahead.
                    interval.extend_end(pos.saturating_add(1));
                }
                pos = pos.saturating_add(1);
            }
        }

        // Phase 3: record uses, and seed the backward walk with the blocks each variable is
        // live *out* of.
        let cfg = SsaCfg::from_ssa(ssa);
        let mut live_out_seeds: BTreeMap<SsaVarId, Vec<usize>> = BTreeMap::new();

        for block_id in 0..block_count {
            let Some(block) = ssa.block(block_id) else {
                continue;
            };
            let Some(&entry_pos) = block_start_idx.get(block_id) else {
                continue;
            };

            // A phi operand is used at the END of its predecessor block, not here, so the
            // value must stay live through that predecessor even when it is laid out after
            // the phi's own block.
            for phi in block.phi_nodes() {
                for operand in phi.operands() {
                    let pred = operand.predecessor();
                    let Some(&pred_end) = block_end_idx.get(pred) else {
                        continue;
                    };
                    let value = operand.value();
                    intervals
                        .entry(value)
                        .or_insert_with(|| LiveInterval::new(pred_end))
                        .extend_end(pred_end);
                    live_out_seeds.entry(value).or_default().push(pred);
                }
            }

            // Seeding is per (variable, block), not per use. The walk's generation stamps
            // would dedupe the extra work anyway, but the seed vector itself would not:
            // a variable used u times in a block with p predecessors would push u·p entries
            // instead of p, and both factors are attacker-controlled.
            let mut seeded_here: BTreeSet<SsaVarId> = BTreeSet::new();

            let mut pos = entry_pos;
            for instr in block.instructions() {
                for &use_var in &instr.uses() {
                    let interval = intervals
                        .entry(use_var)
                        .or_insert_with(|| LiveInterval::new(pos));
                    interval.extend_end(pos.saturating_add(1));

                    // A use in the defining block needs no walk — the value is live only
                    // from the definition to this point. Otherwise the value is live-in
                    // here, hence live-out of every predecessor.
                    if def_block.get(&use_var) != Some(&block_id) {
                        interval.extend_start(entry_pos);
                        if seeded_here.insert(use_var) {
                            live_out_seeds
                                .entry(use_var)
                                .or_default()
                                .extend(cfg.block_predecessors(block_id).iter().copied());
                        }
                    }
                }
                pos = pos.saturating_add(1);
            }
        }

        // Phase 4: up-and-mark. For each variable, walk up from every block it is live out
        // of, covering whole blocks, and stop at its definition.
        //
        // `seen` holds a generation stamp per block rather than a boolean, so resetting
        // between variables is a single counter bump instead of an O(B) clear. The counter
        // is `u64`, which cannot wrap within one function's variable count.
        let mut seen: Vec<u64> = vec![0; block_count];
        let mut generation: u64 = 0;
        let mut worklist: Vec<usize> = Vec::new();

        for (var, seeds) in &live_out_seeds {
            generation = generation.saturating_add(1);
            let stop = def_block.get(var).copied();

            worklist.clear();
            worklist.extend(seeds.iter().copied());

            while let Some(block_id) = worklist.pop() {
                let Some(stamp) = seen.get_mut(block_id) else {
                    continue;
                };
                if *stamp == generation {
                    continue;
                }
                *stamp = generation;

                if let (Some(&block_start), Some(&block_end), Some(interval)) = (
                    block_start_idx.get(block_id),
                    block_end_idx.get(block_id),
                    intervals.get_mut(var),
                ) {
                    interval.extend_end(block_end.max(block_start.saturating_add(1)));
                    // In the defining block the value is live only from its definition
                    // onward, so the block's start must not pull the interval back.
                    if stop != Some(block_id) {
                        interval.extend_start(block_start);
                    }
                }

                // The definition dominates every use, so nothing above it is live. A
                // variable with no recorded definition (an argument or entry live-in) has
                // no stop block and correctly walks back to the entry, which has no
                // predecessors.
                if stop == Some(block_id) {
                    continue;
                }

                worklist.extend(cfg.block_predecessors(block_id).iter().copied());
            }
        }

        intervals
    }

    /// Collects interference edges from intra-block liveness for a single block.
    ///
    /// Returns a vector of edge pairs that can be merged later.
    fn collect_intra_block_edges(
        ssa: &SsaFunction,
        results: &AnalysisResults<LivenessResult>,
        block_id: usize,
    ) -> Vec<(SsaVarId, SsaVarId)> {
        let mut edges = Vec::new();
        let var_count = ssa.var_id_capacity();

        let Some(block) = ssa.block(block_id) else {
            return edges;
        };

        // Start with live_out from this block
        let mut live = BitSet::new(var_count);
        if let Some(live_out) = results.out_state(block_id) {
            for var in live_out.variables() {
                live.insert(var.index());
            }
        }

        // Walk instructions backwards
        for instr in block.instructions().iter().rev() {
            // If this instruction defines a variable, it interferes with all live vars
            if let Some(def) = instr.def() {
                for live_idx in live.iter() {
                    let live_var = SsaVarId::from_index(live_idx);
                    if live_var != def {
                        edges.push((def, live_var));
                    }
                }
                // The defined variable is no longer live before this point
                live.remove(def.index());
            }

            // Uses make variables live
            for &use_var in &instr.uses() {
                live.insert(use_var.index());
            }
        }

        // Handle phi nodes (they define at block entry)
        for phi in block.phi_nodes() {
            let def = phi.result();
            for live_idx in live.iter() {
                let live_var = SsaVarId::from_index(live_idx);
                if live_var != def {
                    edges.push((def, live_var));
                }
            }
        }

        edges
    }

    /// Collects interference edges between variables from different root subtrees.
    ///
    /// The codegen's dependency-driven scheduler processes root instructions by
    /// generating their complete dependency subtrees in sequence. Since subtrees
    /// can be executed in any relative order, variables from different subtrees
    /// must be treated as potentially simultaneously live.
    ///
    /// For each block with multiple roots, this function:
    /// 1. Identifies root instructions (those not consumed by other in-block ops)
    /// 2. Assigns each instruction to the first root subtree that claims it (via DFS)
    /// 3. Adds interference edges between all variable pairs from different subtrees
    fn collect_cross_subtree_edges(
        ssa: &SsaFunction,
        block_id: usize,
    ) -> Vec<(SsaVarId, SsaVarId)> {
        let Some(block) = ssa.block(block_id) else {
            return Vec::new();
        };

        let instructions = block.instructions();
        if instructions.len() <= 1 {
            return Vec::new();
        }

        // Build def_map: variable -> instruction index within this block
        let mut def_map: BTreeMap<SsaVarId, usize> = BTreeMap::new();
        for (idx, instr) in instructions.iter().enumerate() {
            if let Some(dest) = instr.def() {
                def_map.insert(dest, idx);
            }
        }

        // Identify which variables are used as operands within this block
        let var_count = ssa.var_id_capacity();
        let mut used_in_block = BitSet::new(var_count);
        for instr in instructions {
            for use_var in instr.uses() {
                used_in_block.insert(use_var.index());
            }
        }

        // Find root instructions: those whose results are NOT used by other ops
        // in this block, or that have no result (side-effect only / terminators)
        let roots: Vec<usize> = instructions
            .iter()
            .enumerate()
            .filter_map(|(idx, instr)| {
                let is_root = match instr.def() {
                    Some(dest) => {
                        !used_in_block.contains(dest.index())
                            || matches!(
                                instr.op(),
                                SsaOp::Jump { .. }
                                    | SsaOp::Branch { .. }
                                    | SsaOp::BranchCmp { .. }
                                    | SsaOp::Switch { .. }
                                    | SsaOp::Return { .. }
                                    | SsaOp::Throw { .. }
                                    | SsaOp::Rethrow
                                    | SsaOp::Leave { .. }
                                    | SsaOp::EndFinally
                                    | SsaOp::EndFilter { .. }
                            )
                    }
                    None => true,
                };
                if is_root {
                    Some(idx)
                } else {
                    None
                }
            })
            .collect();

        if roots.len() <= 1 {
            return Vec::new();
        }

        // Assign each instruction to a root subtree via DFS.
        // The first subtree (in root order) to reach an instruction claims it.
        // This models the scheduler: subtrees processed first claim shared deps.
        let mut instr_to_subtree: Vec<Option<usize>> = vec![None; instructions.len()];
        for (subtree_id, &root_idx) in roots.iter().enumerate() {
            let mut stack = vec![root_idx];
            while let Some(idx) = stack.pop() {
                let Some(slot) = instr_to_subtree.get_mut(idx) else {
                    continue;
                };
                if slot.is_some() {
                    continue;
                }
                *slot = Some(subtree_id);

                // Follow dependency edges: operands defined in this block
                let Some(instr) = instructions.get(idx) else {
                    continue;
                };
                for use_var in instr.uses() {
                    if let Some(&dep_idx) = def_map.get(&use_var) {
                        if matches!(instr_to_subtree.get(dep_idx), Some(None)) {
                            stack.push(dep_idx);
                        }
                    }
                }
            }
        }

        // Group defined variables by their subtree
        let num_subtrees = roots.len();
        let mut subtree_vars: Vec<Vec<SsaVarId>> = vec![Vec::new(); num_subtrees];
        for (idx, instr) in instructions.iter().enumerate() {
            if let (Some(dest), Some(Some(subtree_id))) =
                (instr.def(), instr_to_subtree.get(idx).copied())
            {
                if let Some(bucket) = subtree_vars.get_mut(subtree_id) {
                    bucket.push(dest);
                }
            }
        }

        // Add interference edges between all variable pairs from different subtrees
        let mut edges = Vec::new();
        for i in 0..num_subtrees {
            let Some(vars_i) = subtree_vars.get(i) else {
                continue;
            };
            for j in i.saturating_add(1)..num_subtrees {
                let Some(vars_j) = subtree_vars.get(j) else {
                    continue;
                };
                for &var_a in vars_i {
                    for &var_b in vars_j {
                        edges.push((var_a, var_b));
                    }
                }
            }
        }

        edges
    }

    /// Allocates local slots for SSA variables.
    ///
    /// Returns a mapping from SSA variables to local slots, minimizing the
    /// total number of locals needed. Uses the result from whatever algorithm
    /// was selected during `build()`.
    pub fn allocate(&self, ssa: &SsaFunction) -> Result<LocalAllocation> {
        // If we used linear scan, the allocation is already computed
        if let Some(precomputed) = &self.precomputed {
            return Ok(LocalAllocation {
                var_to_local: precomputed.var_to_local.clone(),
                num_locals: precomputed.num_locals,
                original_to_compacted: precomputed.original_to_compacted.clone(),
            });
        }

        // Otherwise use graph coloring allocation
        self.allocate_graph_coloring(ssa)
    }

    /// Allocates local slots using greedy graph coloring.
    fn allocate_graph_coloring(&self, ssa: &SsaFunction) -> Result<LocalAllocation> {
        let var_count = ssa.variable_count();

        // First, collect which Local-origin variables are actually USED.
        // Phi-origin variables go through graph coloring allocation separately.
        let slot_capacity = var_count
            .max(self.coalescable_vars.len().saturating_add(1))
            .max(64);
        let mut used_local_var_ids = BitSet::new(slot_capacity);
        for block in ssa.blocks() {
            for phi in block.phi_nodes() {
                // Check phi result
                if let Some(var) = ssa.variable(phi.result()) {
                    if matches!(var.origin(), VariableOrigin::Local(_)) {
                        used_local_var_ids.insert(phi.result().index());
                    }
                }
                // Check phi operands
                for operand in phi.operands() {
                    if let Some(var) = ssa.variable(operand.value()) {
                        if matches!(var.origin(), VariableOrigin::Local(_)) {
                            used_local_var_ids.insert(operand.value().index());
                        }
                    }
                }
            }
            for instr in block.instructions() {
                let op = instr.op();
                // Check destination
                if let Some(dest) = op.dest() {
                    if let Some(var) = ssa.variable(dest) {
                        if matches!(var.origin(), VariableOrigin::Local(_)) {
                            used_local_var_ids.insert(dest.index());
                        }
                    }
                }
                // Check uses
                for use_var in op.uses() {
                    if let Some(var) = ssa.variable(use_var) {
                        if matches!(var.origin(), VariableOrigin::Local(_)) {
                            used_local_var_ids.insert(use_var.index());
                        }
                    }
                }
            }
        }

        // Collect Local-origin variables sorted by original index for deterministic ordering
        let mut local_vars: Vec<(SsaVarId, u16)> = ssa
            .variables()
            .iter()
            .filter_map(|v| {
                if let VariableOrigin::Local(idx) = v.origin() {
                    if used_local_var_ids.contains(v.id().index()) {
                        return Some((v.id(), idx));
                    }
                }
                None
            })
            .collect();
        local_vars.sort_by_key(|(_, idx)| *idx);

        let PreAssignment {
            mut var_to_local,
            reserved_slots,
            original_to_new,
            mut next_local,
        } = pre_assign_locals(ssa, &local_vars);

        // Sort coalescable variables by degree (most constrained first)
        // This is a simple heuristic that often produces good colorings.
        // Use (origin, version) as secondary key for deterministic ordering when degrees are equal.
        let mut sorted_vars = self.coalescable_vars.clone();
        sorted_vars.sort_by(|a, b| {
            let deg_a = self.interference.degree(*a);
            let deg_b = self.interference.degree(*b);
            deg_b
                .cmp(&deg_a)
                .then_with(|| var_sort_key(ssa, *a).cmp(&var_sort_key(ssa, *b)))
        });

        // Greedy coloring with type compatibility
        for var in sorted_vars {
            // Find slots used by interfering neighbors
            let mut used_slots = BitSet::new(slot_capacity);
            for neighbor in self.interference.neighbors(var) {
                if let Some(&slot) = var_to_local.get(&neighbor) {
                    used_slots.insert(slot as usize);
                }
            }

            // Get this variable's type for compatibility checking
            let var_type = self.interference.var_types.get(&var);

            // Find first available slot that is type-compatible
            // Skip reserved slots (those belonging to Local-origin variables)
            // Safe: will always find a slot since u16 range exceeds possible variable count
            #[allow(clippy::maybe_infinite_iter)]
            let slot = (0u16..)
                .find(|&s| {
                    // Skip reserved slots - they belong to Local-origin variables
                    if reserved_slots.contains(s as usize) {
                        return false;
                    }

                    // Slot must not be used by an interfering variable
                    if used_slots.contains(s as usize) {
                        return false;
                    }

                    // Check type compatibility with any variable already using this slot
                    for (&other_var, &other_slot) in &var_to_local {
                        if other_slot == s {
                            let other_type = self.interference.var_types.get(&other_var);
                            if !types_compatible(var_type, other_type) {
                                return false;
                            }
                        }
                    }

                    true
                })
                .ok_or_else(|| Error::CodegenFailed("Should always find a valid slot".into()))?;

            var_to_local.insert(var, slot);
            next_local = next_local.max(slot.saturating_add(1));
        }

        Ok(LocalAllocation {
            var_to_local,
            num_locals: next_local,
            original_to_compacted: original_to_new,
        })
    }
}

/// Pre-assignment result for local variables.
///
/// Contains the initial slot assignments for Local-origin and LoadLocal-referenced
/// variables, along with tracking state for the allocation phase.
struct PreAssignment {
    /// Map from SSA variable ID to assigned local slot.
    var_to_local: BTreeMap<SsaVarId, u16>,
    /// Slots reserved for Local-origin variables (must not be reused).
    reserved_slots: BitSet,
    /// Map from original local index to compacted index.
    original_to_new: BTreeMap<u16, u16>,
    /// Next available local slot index.
    next_local: u16,
}

/// Pre-assigns local slots for Local-origin and LoadLocal-referenced variables.
///
/// This shared logic is used by both graph coloring and linear scan algorithms.
/// It assigns consecutive slots to used Local-origin variables (deduplicating
/// SSA versions of the same original local) and ensures LoadLocal/LoadLocalAddr-
/// referenced locals also get compacted entries.
///
/// `used_local_vars` must be a pre-filtered, sorted list of `(var_id, original_local_index)`
/// pairs representing Local-origin variables that are actually used.
fn pre_assign_locals(ssa: &SsaFunction, used_local_vars: &[(SsaVarId, u16)]) -> PreAssignment {
    let var_count = ssa.variable_count();
    let slot_capacity = var_count.max(64);
    let mut var_to_local: BTreeMap<SsaVarId, u16> = BTreeMap::new();
    let mut reserved_slots = BitSet::new(slot_capacity);
    let mut original_to_new: BTreeMap<u16, u16> = BTreeMap::new();
    let mut next_local: u16 = 0;

    // Assign consecutive slots to USED Local-origin variables.
    // All SSA versions of the same original local share the same physical slot.
    for &(var_id, original_idx) in used_local_vars {
        let new_slot = *original_to_new.entry(original_idx).or_insert_with(|| {
            let slot = next_local;
            next_local = next_local.saturating_add(1);
            slot
        });
        var_to_local.insert(var_id, new_slot);
        reserved_slots.insert(new_slot as usize);
    }

    // Ensure LoadLocal/LoadLocalAddr-referenced locals also get compacted entries.
    // These reference locals by index (not SSA variable), so the coalescer's
    // liveness analysis doesn't see them.
    let mut load_referenced_locals = BitSet::new(slot_capacity);
    for block in ssa.blocks() {
        for instr in block.instructions() {
            match instr.op() {
                SsaOp::LoadLocal { local_index, .. } | SsaOp::LoadLocalAddr { local_index, .. } => {
                    load_referenced_locals.insert(*local_index as usize);
                }
                _ => {}
            }
        }
    }
    let mut sorted_load_refs: Vec<u16> = load_referenced_locals.iter().map(|i| i as u16).collect();
    sorted_load_refs.sort_unstable();
    for original_idx in sorted_load_refs {
        original_to_new.entry(original_idx).or_insert_with(|| {
            let slot = next_local;
            next_local = next_local.saturating_add(1);
            reserved_slots.insert(slot as usize);
            slot
        });
    }

    PreAssignment {
        var_to_local,
        reserved_slots,
        original_to_new,
        next_local,
    }
}

/// Builds a stable sort key for an SSA variable.
///
/// Returns `(origin, version)` which is deterministic within a method,
/// unlike `SsaVarId` which depends on global allocation order.
fn var_sort_key(ssa: &SsaFunction, var_id: SsaVarId) -> (VariableOrigin, u32) {
    ssa.variable(var_id)
        .map_or((VariableOrigin::Phi, u32::MAX), |v| {
            (v.origin(), v.version())
        })
}

/// Checks if two types can share the same local slot.
///
/// Types are compatible if they have the same size and alignment requirements.
fn types_compatible(t1: Option<&SsaType>, t2: Option<&SsaType>) -> bool {
    match (t1, t2) {
        (None, _) | (_, None) => true, // Unknown types are conservatively compatible
        (Some(a), Some(b)) => a.is_compatible_for_storage(b),
    }
}

#[cfg(test)]
mod tests {
    use std::collections::BTreeSet;

    use super::*;
    use crate::analysis::{SsaFunctionBuilder, SsaType, SsaVarId};

    /// Helper to create N unique SsaVarIds
    fn make_vars(n: usize) -> Vec<SsaVarId> {
        (0..n).map(SsaVarId::from_index).collect()
    }

    #[test]
    fn test_interference_graph_add_edge() {
        let mut graph = InterferenceGraph::new(10);
        let vars = make_vars(3);
        let (var_a, var_b, var_c) = (vars[0], vars[1], vars[2]);

        // Adding edge creates bidirectional interference
        graph.add_edge(var_a, var_b);
        assert_eq!(graph.degree(var_a), 1);
        assert_eq!(graph.degree(var_b), 1);
        assert_eq!(graph.degree(var_c), 0);

        // Neighbors are correct
        let neighbors_a: Vec<_> = graph.neighbors(var_a).collect();
        assert_eq!(neighbors_a, vec![var_b]);

        let neighbors_b: Vec<_> = graph.neighbors(var_b).collect();
        assert_eq!(neighbors_b, vec![var_a]);

        // Adding same edge again doesn't increase degree
        graph.add_edge(var_a, var_b);
        assert_eq!(graph.degree(var_a), 1);

        // Self-edges are ignored
        graph.add_edge(var_a, var_a);
        assert_eq!(graph.degree(var_a), 1);
    }

    #[test]
    fn test_interference_graph_multiple_edges() {
        let mut graph = InterferenceGraph::new(10);
        let vars = make_vars(5);

        // Create a clique of 3 variables (all interfere with each other)
        graph.add_edge(vars[0], vars[1]);
        graph.add_edge(vars[0], vars[2]);
        graph.add_edge(vars[1], vars[2]);

        // Each has degree 2
        assert_eq!(graph.degree(vars[0]), 2);
        assert_eq!(graph.degree(vars[1]), 2);
        assert_eq!(graph.degree(vars[2]), 2);

        // vars[3] and vars[4] are isolated
        assert_eq!(graph.degree(vars[3]), 0);
        assert_eq!(graph.degree(vars[4]), 0);
    }

    /// A value defined before a loop and used inside it must stay live for the whole
    /// loop, not just up to its last *textual* use.
    ///
    /// Under a textual interval scan `outer`'s interval ends at its single use
    /// in the loop body, so the slot returns to the free pool and can be handed to `inner`,
    /// which is defined later in that same body — and on the second iteration the read of
    /// `outer` returns `inner`'s value.
    #[test]
    fn value_live_across_a_back_edge_keeps_its_slot() {
        // block 0: outer = 7; jump 1
        // block 1: header, branch to body (2) or exit (3)
        // block 2: body, reads `outer`, then defines `inner`; jumps back to 1
        // block 3: ret
        let mut outer_id: Option<SsaVarId> = None;
        let mut inner_id: Option<SsaVarId> = None;

        let ssa = SsaFunctionBuilder::new(1, 0)
            .build_with(|f| {
                let cond = f.arg(0, SsaType::Bool);
                f.block(0, |b| {
                    outer_id = Some(b.const_i32(7));
                    b.jump(1);
                });
                f.block(1, |b| {
                    b.branch(cond, 2, 3);
                });
                let outer = outer_id.expect("block 0 is built first");
                f.block(2, |b| {
                    // `outer` is read here — its last textual mention.
                    let _used = b.add(outer, outer);
                    // ...and `inner` is defined after it, in the same block.
                    let inner = b.const_i32(9);
                    let _also = b.add(inner, inner);
                    inner_id = Some(inner);
                    b.jump(1);
                });
                f.block(3, |b| {
                    b.ret();
                });
            })
            .unwrap();

        let outer = outer_id.expect("outer was defined");
        let inner = inner_id.expect("inner was defined");

        let intervals = LocalCoalescer::compute_live_intervals(&ssa);

        let outer_interval = intervals.get(&outer).expect("outer must have an interval");
        let inner_interval = intervals.get(&inner).expect("inner must have an interval");

        // The two must overlap, which is what stops the allocator reusing one slot for
        // both. Overlap is the property that matters; exact indices are not contractual.
        assert!(
            outer_interval.start < inner_interval.end
                && inner_interval.start < outer_interval.end,
            "outer {outer_interval:?} and inner {inner_interval:?} must overlap across the back edge"
        );

        // And `outer` must survive past the point where `inner` is defined.
        assert!(
            outer_interval.end > inner_interval.start,
            "outer {outer_interval:?} expired before inner {inner_interval:?} was defined"
        );
    }

    #[test]
    fn test_type_compatibility_same_class() {
        // Same types are compatible
        assert!(types_compatible(Some(&SsaType::I32), Some(&SsaType::I32)));

        // 32-bit integers are compatible with each other
        assert!(types_compatible(Some(&SsaType::I32), Some(&SsaType::U32)));
        assert!(types_compatible(Some(&SsaType::I32), Some(&SsaType::Bool)));
        assert!(types_compatible(Some(&SsaType::Bool), Some(&SsaType::Char)));

        // 64-bit integers are compatible
        assert!(types_compatible(Some(&SsaType::I64), Some(&SsaType::U64)));

        // Different reference types are NOT compatible (CIL verifier requires
        // exact type match on stloc/ldloc for reference types)
        assert!(!types_compatible(
            Some(&SsaType::Object),
            Some(&SsaType::String)
        ));

        // Same reference types ARE compatible
        assert!(types_compatible(
            Some(&SsaType::String),
            Some(&SsaType::String)
        ));
    }

    #[test]
    fn test_type_compatibility_different_class() {
        // Different size integers are not compatible
        assert!(!types_compatible(Some(&SsaType::I32), Some(&SsaType::I64)));

        // Integer and reference are not compatible
        assert!(!types_compatible(
            Some(&SsaType::I32),
            Some(&SsaType::Object)
        ));

        // Different float sizes are not compatible
        assert!(!types_compatible(Some(&SsaType::F32), Some(&SsaType::F64)));

        // Integer and float are not compatible
        assert!(!types_compatible(Some(&SsaType::I32), Some(&SsaType::F32)));
    }

    #[test]
    fn test_type_compatibility_with_none() {
        // None (unknown) is conservatively compatible with anything
        assert!(types_compatible(None, Some(&SsaType::I32)));
        assert!(types_compatible(Some(&SsaType::I32), None));
        assert!(types_compatible(None, None));
    }

    #[test]
    fn test_greedy_coloring_non_interfering_same_slot() {
        // When variables don't interfere, they should share the same slot
        let mut graph = InterferenceGraph::new(10);
        let vars = make_vars(2);
        let (var_a, var_b) = (vars[0], vars[1]);

        // Both have compatible types but no interference edge
        graph.set_type(var_a, SsaType::I32);
        graph.set_type(var_b, SsaType::I32);
        // Note: no add_edge call - they don't interfere

        let coalescer = LocalCoalescer {
            interference: graph,
            coalescable_vars: vec![var_a, var_b],
            precomputed: None,
        };

        // Create a minimal SSA function for allocation
        let ssa = create_minimal_ssa_function();
        let allocation = coalescer.allocate(&ssa).unwrap();

        // Non-interfering variables with compatible types should share slot 0
        assert_eq!(allocation.var_to_local.get(&var_a), Some(&0));
        assert_eq!(allocation.var_to_local.get(&var_b), Some(&0));
        assert_eq!(allocation.num_locals, 1);
    }

    #[test]
    fn test_greedy_coloring_interfering_different_slots() {
        // When variables interfere, they must get different slots
        let mut graph = InterferenceGraph::new(10);
        let vars = make_vars(2);
        let (var_a, var_b) = (vars[0], vars[1]);

        graph.set_type(var_a, SsaType::I32);
        graph.set_type(var_b, SsaType::I32);
        graph.add_edge(var_a, var_b); // They interfere!

        let coalescer = LocalCoalescer {
            interference: graph,
            coalescable_vars: vec![var_a, var_b],
            precomputed: None,
        };

        let ssa = create_minimal_ssa_function();
        let allocation = coalescer.allocate(&ssa).unwrap();

        // Interfering variables must have different slots
        let slot_a = allocation.var_to_local.get(&var_a).unwrap();
        let slot_b = allocation.var_to_local.get(&var_b).unwrap();
        assert_ne!(slot_a, slot_b);
        assert_eq!(allocation.num_locals, 2);
    }

    #[test]
    fn test_greedy_coloring_type_incompatible_different_slots() {
        // Even non-interfering variables need different slots if types are incompatible
        let mut graph = InterferenceGraph::new(10);
        let vars = make_vars(2);
        let (var_a, var_b) = (vars[0], vars[1]);

        graph.set_type(var_a, SsaType::I32);
        graph.set_type(var_b, SsaType::I64); // Different type class!
                                             // No interference edge - they're not live at the same time

        let coalescer = LocalCoalescer {
            interference: graph,
            coalescable_vars: vec![var_a, var_b],
            precomputed: None,
        };

        let ssa = create_minimal_ssa_function();
        let allocation = coalescer.allocate(&ssa).unwrap();

        // Type-incompatible variables must have different slots
        let slot_a = allocation.var_to_local.get(&var_a).unwrap();
        let slot_b = allocation.var_to_local.get(&var_b).unwrap();
        assert_ne!(slot_a, slot_b);
        assert_eq!(allocation.num_locals, 2);
    }

    #[test]
    fn test_greedy_coloring_clique_needs_n_colors() {
        // A clique of n nodes needs n colors (slots)
        let mut graph = InterferenceGraph::new(10);
        let vars = make_vars(4);

        // Create a clique - all 4 variables interfere with each other
        for i in 0..4 {
            graph.set_type(vars[i], SsaType::I32);
            for j in (i + 1)..4 {
                graph.add_edge(vars[i], vars[j]);
            }
        }

        let coalescer = LocalCoalescer {
            interference: graph,
            coalescable_vars: vars.clone(),
            precomputed: None,
        };

        let ssa = create_minimal_ssa_function();
        let allocation = coalescer.allocate(&ssa).unwrap();

        // All 4 must have unique slots
        let slots: BTreeSet<_> = vars
            .iter()
            .filter_map(|v| allocation.var_to_local.get(v).copied())
            .collect();
        assert_eq!(slots.len(), 4);
        assert_eq!(allocation.num_locals, 4);
    }

    #[test]
    fn test_greedy_coloring_chain_needs_2_colors() {
        // A chain graph (v0-v1-v2-v3) is 2-colorable
        let mut graph = InterferenceGraph::new(10);
        let vars = make_vars(4);

        // Chain: 0-1-2-3
        for var in &vars {
            graph.set_type(*var, SsaType::I32);
        }
        graph.add_edge(vars[0], vars[1]);
        graph.add_edge(vars[1], vars[2]);
        graph.add_edge(vars[2], vars[3]);

        let coalescer = LocalCoalescer {
            interference: graph,
            coalescable_vars: vars.clone(),
            precomputed: None,
        };

        let ssa = create_minimal_ssa_function();
        let allocation = coalescer.allocate(&ssa).unwrap();

        // Adjacent variables must have different slots
        for i in 0..3 {
            let slot_i = allocation.var_to_local.get(&vars[i]).unwrap();
            let slot_j = allocation.var_to_local.get(&vars[i + 1]).unwrap();
            assert_ne!(
                slot_i,
                slot_j,
                "Adjacent vars {} and {} share slot",
                i,
                i + 1
            );
        }

        // Chain is 2-colorable, so at most 2 slots needed
        assert!(allocation.num_locals <= 2);
    }

    #[test]
    fn test_mixed_types_coalesce_within_class() {
        // Variables of compatible types that don't interfere should share slots
        let mut graph = InterferenceGraph::new(10);
        let vars = make_vars(5);
        let (var_i32, var_u32, var_bool, var_i64, var_u64) =
            (vars[0], vars[1], vars[2], vars[3], vars[4]);

        // All Int32-class types (should share one slot)
        graph.set_type(var_i32, SsaType::I32);
        graph.set_type(var_u32, SsaType::U32);
        graph.set_type(var_bool, SsaType::Bool);
        // All Int64-class types (should share another slot)
        graph.set_type(var_i64, SsaType::I64);
        graph.set_type(var_u64, SsaType::U64);

        // No interference edges

        let coalescer = LocalCoalescer {
            interference: graph,
            coalescable_vars: vec![var_i32, var_u32, var_bool, var_i64, var_u64],
            precomputed: None,
        };

        let ssa = create_minimal_ssa_function();
        let allocation = coalescer.allocate(&ssa).unwrap();

        // All Int32-class should share slot 0
        assert_eq!(allocation.var_to_local.get(&var_i32), Some(&0));
        assert_eq!(allocation.var_to_local.get(&var_u32), Some(&0));
        assert_eq!(allocation.var_to_local.get(&var_bool), Some(&0));

        // All Int64-class should share slot 1
        assert_eq!(allocation.var_to_local.get(&var_i64), Some(&1));
        assert_eq!(allocation.var_to_local.get(&var_u64), Some(&1));

        assert_eq!(allocation.num_locals, 2);
    }

    /// Creates a minimal SSA function for testing allocation.
    fn create_minimal_ssa_function() -> SsaFunction {
        SsaFunctionBuilder::new(0, 0)
            .build_with(|ctx| {
                ctx.block(0, |b| {
                    b.ret();
                });
            })
            .unwrap()
    }
}