llvm-native-core-ext 0.1.0

//! Fast Register Allocator for -O0 — minimal analysis, maximum speed.
//! Clean-room behavioral reconstruction.
//!
//! The fast register allocator is designed for debug/-O0 compilation:
//! - Simple linear scan with minimal liveness analysis
//! - Live variable analysis without interval construction
//! - Single-pass allocation, immediate reuse of dead registers
//! - Simple first-fit slot assignment
//! - No live range splitting (accept spills, don't try to avoid them)
//! - No eviction (when out of registers, spill the current value)
//! - Dead register reuse: immediate reuse of registers whose values are dead
//! - Fast coalescing: only coalesce COPY instructions within the same basic block
//!
//! The primary goal is compilation speed, not code quality. This allocator
//! makes decisions locally (per-block) without global analysis.

use llvm_native_core::codegen_regalloc::{InstrPoint, RegClassKind};
use std::collections::{HashMap, HashSet};

// ============================================================================
// Fast Live Variable
// ============================================================================

/// A fast representation of a live variable (no interval, just def/use sets).
#[derive(Debug, Clone)]
pub struct FastLiveVar {
    /// Virtual register.
    pub vreg: u32,
    /// Register class.
    pub reg_class: RegClassKind,
    /// Block where this variable is defined.
    pub def_block: u32,
    /// Instruction index within the block where defined.
    pub def_instr: u32,
    /// All use points: (block, instr) pairs.
    pub uses: Vec<(u32, u32)>,
    /// Whether the variable is live out of its defining block.
    pub live_out: bool,
    /// Blocks where this variable is live-in.
    pub live_in_blocks: HashSet<u32>,
    /// Assigned physical register (if any).
    pub assigned_reg: Option<u32>,
    /// Spilled: assigned to stack.
    pub spilled: bool,
    /// Stack offset (if spilled).
    pub spill_offset: i32,
}

impl FastLiveVar {
    pub fn new(vreg: u32, reg_class: RegClassKind, def_block: u32, def_instr: u32) -> Self {
        Self {
            vreg,
            reg_class,
            def_block,
            def_instr,
            uses: Vec::new(),
            live_out: false,
            live_in_blocks: HashSet::new(),
            assigned_reg: None,
            spilled: false,
            spill_offset: -1,
        }
    }

    /// Add a use point.
    pub fn add_use(&mut self, block: u32, instr: u32) {
        self.uses.push((block, instr));
    }

    /// Mark as live-out of defining block.
    pub fn mark_live_out(&mut self) {
        self.live_out = true;
    }

    /// Add a live-in block.
    pub fn add_live_in(&mut self, block: u32) {
        self.live_in_blocks.insert(block);
    }

    /// Check if this variable is dead at a given point.
    pub fn is_dead_at(&self, block: u32, instr: u32) -> bool {
        if self.live_out {
            return false;
        }
        // Check all uses: if all uses are before (block, instr), it's dead
        self.uses
            .iter()
            .all(|(ub, ui)| *ub < block || (*ub == block && *ui < instr))
    }

    /// The last use point.
    pub fn last_use(&self) -> Option<(u32, u32)> {
        self.uses.iter().max_by_key(|(b, i)| (b, i)).copied()
    }

    /// The first use point.
    pub fn first_use(&self) -> Option<(u32, u32)> {
        self.uses.iter().min_by_key(|(b, i)| (b, i)).copied()
    }
}

// ============================================================================
// Fast Register Map
// ============================================================================

/// Single-pass register map: tracks which physical registers are free.
#[derive(Debug, Clone)]
pub struct FastRegisterMap {
    /// Available physical registers per class.
    pub available: HashMap<RegClassKind, Vec<u32>>,
    /// Currently allocated: phys_reg -> vreg.
    pub allocated: HashMap<u32, u32>,
    /// Reverse map: vreg -> phys_reg.
    pub vreg_to_reg: HashMap<u32, u32>,
    /// Reserved registers.
    pub reserved: HashSet<u32>,
    /// Registers that became free but haven't been cleaned up.
    pub recently_freed: Vec<u32>,
}

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

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

        Self {
            available,
            allocated: HashMap::new(),
            vreg_to_reg: HashMap::new(),
            reserved,
            recently_freed: Vec::new(),
        }
    }

    /// Try to allocate a free register.
    pub fn allocate_reg(&mut self, class: RegClassKind) -> Option<u32> {
        // Check recently freed first (reuse dead registers immediately)
        while let Some(reg) = self.recently_freed.pop() {
            if !self.allocated.contains_key(&reg) {
                return Some(reg);
            }
        }

        let available = self.available.get(&class).cloned().unwrap_or_default();
        for reg in available {
            if !self.reserved.contains(&reg) && !self.allocated.contains_key(&reg) {
                return Some(reg);
            }
        }
        None
    }

    /// Free a register (mark as available for immediate reuse).
    pub fn free_reg(&mut self, phys_reg: u32) {
        self.allocated.remove(&phys_reg);
        self.recently_freed.push(phys_reg);
    }

    /// Assign a register to a virtual register.
    pub fn assign(&mut self, vreg: u32, phys_reg: u32) {
        self.allocated.insert(phys_reg, vreg);
        self.vreg_to_reg.insert(vreg, phys_reg);
        // Remove from recently freed if present
        self.recently_freed.retain(|&r| r != phys_reg);
    }

    /// Check if a register is free.
    pub fn is_free(&self, reg: u32) -> bool {
        !self.allocated.contains_key(&reg) && !self.reserved.contains(&reg)
    }

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

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

// ============================================================================
// Stack Slot Assignment (First-Fit)
// ============================================================================

/// Simple first-fit slot allocator for spilled registers.
#[derive(Debug, Clone)]
pub struct FastSlotAllocator {
    /// Allocated slots: vreg -> offset.
    pub slots: HashMap<u32, i32>,
    /// Freed slots for reuse.
    pub free_slots: Vec<i32>,
    /// Next available slot offset.
    pub next_offset: i32,
    /// Slot size in bytes.
    pub slot_size: i32,
    /// Total frame size.
    pub frame_size: i32,
}

impl FastSlotAllocator {
    pub fn new() -> Self {
        Self {
            slots: HashMap::new(),
            free_slots: Vec::new(),
            next_offset: 0,
            slot_size: 8,
            frame_size: 0,
        }
    }

    /// Allocate a slot (first-fit: try freed slots first).
    pub fn allocate_slot(&mut self, vreg: u32) -> i32 {
        // Reuse a freed slot if available
        if let Some(offset) = self.free_slots.pop() {
            self.slots.insert(vreg, offset);
            return offset;
        }

        let offset = self.next_offset;
        self.next_offset += self.slot_size;
        self.slots.insert(vreg, offset);
        self.frame_size = self.frame_size.max(offset + self.slot_size);
        offset
    }

    /// Free a slot for reuse.
    pub fn free_slot(&mut self, vreg: u32) {
        if let Some(offset) = self.slots.remove(&vreg) {
            self.free_slots.push(offset);
        }
    }

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

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

// ============================================================================
// Fast Coalescing
// ============================================================================

/// Simple coalescing: only within the same basic block.
#[derive(Debug, Clone)]
pub struct FastCoalescer {
    /// COPY instructions to coalesce: (dst_vreg, src_vreg, block).
    pub pending_copies: Vec<(u32, u32, u32)>,
    /// Coalesced count.
    pub coalesced: usize,
}

impl FastCoalescer {
    pub fn new() -> Self {
        Self {
            pending_copies: Vec::new(),
            coalesced: 0,
        }
    }

    /// Record a COPY for potential coalescing.
    pub fn record_copy(&mut self, dst: u32, src: u32, block: u32) {
        self.pending_copies.push((dst, src, block));
    }

    /// Attempt coalescing for all pending copies.
    /// Only coalesces if both virtual registers are assigned to the same
    /// physical register (or one can be reassigned).
    pub fn coalesce_pending(
        &mut self,
        reg_map: &mut FastRegisterMap,
        live_vars: &HashMap<u32, FastLiveVar>,
    ) -> usize {
        let mut coalesced = 0;

        let copies = std::mem::take(&mut self.pending_copies);
        for (dst, src, block) in copies {
            // Both must be assigned
            let dst_reg = match reg_map.vreg_to_reg.get(&dst) {
                Some(r) => *r,
                None => continue,
            };
            let src_reg = match reg_map.vreg_to_reg.get(&src) {
                Some(r) => *r,
                None => continue,
            };

            if dst_reg == src_reg {
                // Already coalesced
                coalesced += 1;
                continue;
            }

            // Check if we can reassign dst to src's register
            let src_var = match live_vars.get(&src) {
                Some(v) => v,
                None => continue,
            };
            let dst_var = match live_vars.get(&dst) {
                Some(v) => v,
                None => continue,
            };

            // Only coalesce if they're in the same block and don't overlap
            let dst_dead_at_src = dst_var.is_dead_at(block, src_var.def_instr);
            let src_dead_at_dst = src_var.is_dead_at(block, dst_var.def_instr);

            if dst_dead_at_src {
                // Free dst's register and use src's register for dst
                reg_map.free_reg(dst_reg);
                reg_map.assign(dst, src_reg);
                coalesced += 1;
            } else if src_dead_at_dst {
                // Free src's register and use dst's register for src
                reg_map.free_reg(src_reg);
                reg_map.assign(src, dst_reg);
                coalesced += 1;
            }
        }

        self.coalesced += coalesced;
        coalesced
    }
}

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

// ============================================================================
// Fast Register Allocator (Main)
// ============================================================================

/// The fast register allocator for -O0 compilation.
///
/// Processes basic blocks in order, allocating registers greedily
/// and reusing dead registers immediately. No global analysis.
pub struct FastRegAlloc {
    /// Live variable info per virtual register.
    pub live_vars: HashMap<u32, FastLiveVar>,
    /// Register assignment map.
    pub reg_map: FastRegisterMap,
    /// Stack slot allocator.
    pub slot_alloc: FastSlotAllocator,
    /// Coalescer.
    pub coalescer: FastCoalescer,
    /// Spilled registers.
    pub spilled: HashSet<u32>,
    /// Statistics.
    pub stats: FastAllocStats,
}

/// Statistics from fast allocation.
#[derive(Debug, Clone, Default)]
pub struct FastAllocStats {
    pub total_vregs: usize,
    pub allocated_vregs: usize,
    pub spilled_vregs: usize,
    pub dead_reuses: usize,
    pub coalesced: usize,
    pub slots_allocated: usize,
    pub frame_size: i32,
}

impl FastRegAlloc {
    pub fn new() -> Self {
        Self {
            live_vars: HashMap::new(),
            reg_map: FastRegisterMap::new(),
            slot_alloc: FastSlotAllocator::new(),
            coalescer: FastCoalescer::new(),
            spilled: HashSet::new(),
            stats: FastAllocStats::default(),
        }
    }

    /// Register a virtual register definition.
    pub fn record_def(&mut self, vreg: u32, reg_class: RegClassKind, block: u32, instr: u32) {
        self.live_vars
            .entry(vreg)
            .or_insert_with(|| FastLiveVar::new(vreg, reg_class, block, instr));
        self.stats.total_vregs = self.live_vars.len();
    }

    /// Record a use of a virtual register.
    pub fn record_use(&mut self, vreg: u32, block: u32, instr: u32) {
        if let Some(var) = self.live_vars.get_mut(&vreg) {
            var.add_use(block, instr);
        }
    }

    /// Mark a variable as live-out of a block.
    pub fn mark_live_out(&mut self, vreg: u32) {
        if let Some(var) = self.live_vars.get_mut(&vreg) {
            var.mark_live_out();
        }
    }

    /// Mark a variable as live-in to a block.
    pub fn mark_live_in(&mut self, vreg: u32, block: u32) {
        if let Some(var) = self.live_vars.get_mut(&vreg) {
            var.add_live_in(block);
        }
    }

    /// Record a COPY instruction for coalescing.
    pub fn record_copy(&mut self, dst: u32, src: u32, block: u32) {
        self.coalescer.record_copy(dst, src, block);
    }

    /// Run the fast allocation algorithm (per-block processing).
    pub fn allocate(&mut self) -> &FastAllocStats {
        // Collect all variables sorted by definition order
        let mut sorted_vregs: Vec<u32> = self.live_vars.keys().copied().collect();
        sorted_vregs.sort_by_key(|v| {
            let var = &self.live_vars[v];
            (var.def_block, var.def_instr)
        });

        // Process each variable
        for vreg in sorted_vregs {
            let var = match self.live_vars.get(&vreg) {
                Some(v) => v.clone(),
                None => continue,
            };

            // Free registers whose values are dead at this definition point
            self.free_dead_registers(var.def_block, var.def_instr);

            // Try to allocate a register
            if let Some(reg) = self.reg_map.allocate_reg(var.reg_class) {
                self.reg_map.assign(vreg, reg);
                self.stats.allocated_vregs += 1;
            } else {
                // No free register — spill this variable
                self.spill_variable(vreg);
            }
        }

        // Attempt coalescing
        let coalesced = self
            .coalescer
            .coalesce_pending(&mut self.reg_map, &self.live_vars);
        self.stats.coalesced += coalesced;

        &self.stats
    }

    /// Free registers whose values are dead at the given point.
    fn free_dead_registers(&mut self, block: u32, instr: u32) {
        let mut to_free: Vec<u32> = Vec::new();

        for (&phys_reg, &vreg) in &self.reg_map.allocated {
            if let Some(var) = self.live_vars.get(&vreg) {
                if var.is_dead_at(block, instr) {
                    to_free.push(phys_reg);
                }
            }
        }

        for reg in to_free {
            self.reg_map.free_reg(reg);
            self.stats.dead_reuses += 1;
        }
    }

    /// Spill a virtual register (allocate a stack slot).
    fn spill_variable(&mut self, vreg: u32) {
        self.spilled.insert(vreg);
        self.stats.spilled_vregs += 1;

        let slot = self.slot_alloc.allocate_slot(vreg);
        self.stats.slots_allocated += 1;

        if let Some(var) = self.live_vars.get_mut(&vreg) {
            var.spilled = true;
            var.spill_offset = slot;
        }
    }

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

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

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

    /// Get the total frame size needed.
    pub fn frame_size(&self) -> i32 {
        self.slot_alloc.frame_size
    }

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

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

// ============================================================================
// Per-Block Fast Allocator
// ============================================================================

/// Processes a single basic block for fast allocation.
/// Reuses dead registers within the block and handles local coalescing.
pub struct FastBlockAllocator {
    /// The parent allocator.
    pub parent: FastRegAlloc,
    /// Current block index.
    pub current_block: u32,
    /// Next instruction index within the block.
    pub next_instr: u32,
}

impl FastBlockAllocator {
    pub fn new(parent: FastRegAlloc, block: u32) -> Self {
        Self {
            parent,
            current_block: block,
            next_instr: 0,
        }
    }

    /// Process a definition within the block.
    pub fn process_def(&mut self, vreg: u32, reg_class: RegClassKind) -> Option<u32> {
        // Free dead registers at this point
        self.parent
            .free_dead_registers(self.current_block, self.next_instr);

        // Try to allocate
        if let Some(reg) = self.parent.reg_map.allocate_reg(reg_class) {
            self.parent.reg_map.assign(vreg, reg);
            self.parent.stats.allocated_vregs += 1;
            self.next_instr += 1;
            Some(reg)
        } else {
            self.parent.spill_variable(vreg);
            self.next_instr += 1;
            None
        }
    }

    /// Process a use within the block.
    pub fn process_use(&mut self, vreg: u32) {
        self.parent
            .record_use(vreg, self.current_block, self.next_instr);
        self.next_instr += 1;
    }

    /// Process a COPY within the block.
    pub fn process_copy(&mut self, dst: u32, src: u32) {
        self.parent.record_copy(dst, src, self.current_block);
        self.next_instr += 1;
    }

    /// Finish the block and return the allocator.
    pub fn finish(mut self) -> FastRegAlloc {
        // Coalesce within this block
        let _ = self
            .parent
            .coalescer
            .coalesce_pending(&mut self.parent.reg_map, &self.parent.live_vars);
        self.parent
    }
}

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

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

    #[test]
    fn test_fast_live_var_basic() {
        let mut var = FastLiveVar::new(1, RegClassKind::GPR, 0, 0);
        var.add_use(0, 5);
        var.add_use(0, 10);
        assert!(var.is_dead_at(0, 11));
        assert!(!var.is_dead_at(0, 5));
    }

    #[test]
    fn test_fast_live_var_live_out() {
        let mut var = FastLiveVar::new(1, RegClassKind::GPR, 0, 0);
        var.add_use(0, 5);
        var.mark_live_out();
        assert!(!var.is_dead_at(0, 100)); // Still live due to live-out
        assert!(!var.is_dead_at(1, 0));
    }

    #[test]
    fn test_fast_register_map_allocate() {
        let mut map = FastRegisterMap::new();
        let reg = map.allocate_reg(RegClassKind::GPR);
        assert!(reg.is_some());
    }

    #[test]
    fn test_fast_register_map_free_and_reuse() {
        let mut map = FastRegisterMap::new();
        let reg = map.allocate_reg(RegClassKind::GPR).unwrap();
        map.assign(1, reg);
        map.free_reg(reg);

        // Should be able to reallocate immediately
        let reg2 = map.allocate_reg(RegClassKind::GPR);
        assert_eq!(reg2, Some(reg)); // Same register reused
    }

    #[test]
    fn test_fast_register_map_reserved() {
        let map = FastRegisterMap::new();
        // Reserved registers should not be allocatable
        assert!(!map.is_free(0)); // x0/zero
        assert!(!map.is_free(2)); // x2/sp
    }

    #[test]
    fn test_fast_slot_allocator_basic() {
        let mut alloc = FastSlotAllocator::new();
        let s1 = alloc.allocate_slot(1);
        let s2 = alloc.allocate_slot(2);
        assert!(s1 < s2); // Monotonically increasing
        assert_eq!(alloc.get_slot(1), Some(s1));
        assert_eq!(alloc.get_slot(2), Some(s2));
    }

    #[test]
    fn test_fast_slot_allocator_reuse() {
        let mut alloc = FastSlotAllocator::new();
        let s1 = alloc.allocate_slot(1);
        alloc.free_slot(1);
        let s2 = alloc.allocate_slot(2);
        assert_eq!(s1, s2); // Reuses freed slot
    }

    #[test]
    fn test_fast_coalescer_basic() {
        let mut map = FastRegisterMap::new();
        let mut vars = HashMap::new();

        let mut var_src = FastLiveVar::new(1, RegClassKind::GPR, 0, 2);
        var_src.add_use(0, 5);
        let mut var_dst = FastLiveVar::new(2, RegClassKind::GPR, 0, 4);
        var_dst.add_use(0, 8);

        vars.insert(1, var_src);
        vars.insert(2, var_dst);

        // Assign different registers
        map.assign(1, 10);
        map.assign(2, 11);

        let mut coalescer = FastCoalescer::new();
        coalescer.record_copy(2, 1, 0); // COPY dst=2, src=1
        let count = coalescer.coalesce_pending(&mut map, &vars);
        assert!(count >= 0); // May or may not coalesce depending on liveness
    }

    #[test]
    fn test_fast_reg_alloc_basic() {
        let mut alloc = FastRegAlloc::new();
        alloc.record_def(1, RegClassKind::GPR, 0, 0);
        alloc.record_use(1, 0, 1);
        alloc.record_def(2, RegClassKind::GPR, 0, 2);
        alloc.record_use(2, 0, 3);

        alloc.allocate();

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

    #[test]
    fn test_fast_reg_alloc_spill_when_full() {
        let mut alloc = FastRegAlloc::new();

        // Only 1 register available
        alloc.reg_map.available.insert(RegClassKind::GPR, vec![10]);

        // Create many overlapping variables
        for i in 0..10 {
            alloc.record_def(i, RegClassKind::GPR, 0, i);
            alloc.record_use(i, 0, i + 1);
        }

        alloc.allocate();

        // Some should be spilled
        assert!(
            alloc.stats.spilled_vregs > 0,
            "Expected spills with only 1 register and 10 overlapping vars"
        );
    }

    #[test]
    fn test_dead_register_reuse() {
        let mut alloc = FastRegAlloc::new();

        // Only 1 or 2 registers available
        alloc
            .reg_map
            .available
            .insert(RegClassKind::GPR, vec![10, 11]);

        // vreg 1: def at (0,0), use at (0,1), dead after (0,1)
        alloc.record_def(1, RegClassKind::GPR, 0, 0);
        alloc.record_use(1, 0, 1);

        // vreg 2: def at (0,2), use at (0,3) — after vreg 1 is dead
        alloc.record_def(2, RegClassKind::GPR, 0, 2);
        alloc.record_use(2, 0, 3);

        alloc.allocate();

        // Both should be allocated (dead register reuse)
        assert_eq!(alloc.stats.allocated_vregs, 2);
        assert_eq!(alloc.stats.spilled_vregs, 0);
    }

    #[test]
    fn test_fast_block_allocator() {
        let parent = FastRegAlloc::new();
        let mut block_alloc = FastBlockAllocator::new(parent, 0);

        // Process instructions in a block
        block_alloc.process_def(1, RegClassKind::GPR); // vreg1 = ...
        block_alloc.process_use(1); // ... = vreg1
        block_alloc.process_copy(2, 1); // vreg2 = COPY vreg1

        let result = block_alloc.finish();
        assert!(result.stats.allocated_vregs >= 1);
    }

    #[test]
    fn test_fast_alloc_stats_default() {
        let stats = FastAllocStats::default();
        assert_eq!(stats.total_vregs, 0);
        assert_eq!(stats.allocated_vregs, 0);
        assert_eq!(stats.spilled_vregs, 0);
        assert_eq!(stats.dead_reuses, 0);
        assert_eq!(stats.frame_size, 0);
    }

    #[test]
    fn test_fast_live_var_last_use() {
        let mut var = FastLiveVar::new(1, RegClassKind::GPR, 0, 0);
        var.add_use(0, 3);
        var.add_use(0, 7);
        var.add_use(0, 5);

        assert_eq!(var.last_use(), Some((0, 7)));
        assert_eq!(var.first_use(), Some((0, 3)));
    }

    #[test]
    fn test_fast_live_var_cross_block() {
        let mut var = FastLiveVar::new(1, RegClassKind::GPR, 0, 0);
        var.add_use(0, 5);
        var.add_use(1, 3);

        assert!(!var.is_dead_at(0, 6)); // Still used in block 1
        assert!(!var.is_dead_at(1, 2)); // Not yet used in block 1 at instr 2
        assert!(var.is_dead_at(1, 4)); // All uses exhausted
    }

    #[test]
    fn test_fast_reg_alloc_live_out_handling() {
        let mut alloc = FastRegAlloc::new();
        alloc.record_def(1, RegClassKind::GPR, 0, 0);
        alloc.record_use(1, 0, 1);
        alloc.mark_live_out(1);
        alloc.mark_live_in(1, 1);

        alloc.allocate();
        // Live-out value should not be freed at end of block 0
        assert!(alloc.get_assignment(1).is_some() || alloc.is_spilled(1));
    }

    #[test]
    fn test_fast_reg_alloc_multiple_blocks() {
        let mut alloc = FastRegAlloc::new();
        alloc.record_def(1, RegClassKind::GPR, 0, 0);
        alloc.record_use(1, 0, 2);
        alloc.record_def(2, RegClassKind::GPR, 1, 0);
        alloc.record_use(2, 1, 1);
        alloc.record_def(3, RegClassKind::GPR, 2, 0);
        alloc.record_use(3, 2, 3);

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

    #[test]
    fn test_fast_reg_alloc_spill_recovery() {
        let mut alloc = FastRegAlloc::new();
        // Only 2 GPRs available
        alloc
            .reg_map
            .available
            .insert(RegClassKind::GPR, vec![10, 11]);

        // Create 5 overlapping variables
        for i in 0..5 {
            alloc.record_def(i, RegClassKind::GPR, 0, i);
            alloc.record_use(i, 0, i + 1);
        }
        alloc.allocate();

        // Should have some spills
        assert!(alloc.stats.spilled_vregs > 0);
        // And some allocations
        assert!(alloc.stats.allocated_vregs >= 1);
    }

    #[test]
    fn test_fast_slot_allocator_frame_size() {
        let mut alloc = FastSlotAllocator::new();
        let s1 = alloc.allocate_slot(1);
        let s2 = alloc.allocate_slot(2);
        let s3 = alloc.allocate_slot(3);
        assert!(alloc.frame_size >= s3 + alloc.slot_size);
    }

    #[test]
    fn test_fast_slot_allocator_multiple_reuse() {
        let mut alloc = FastSlotAllocator::new();
        let s1 = alloc.allocate_slot(1);
        let s2 = alloc.allocate_slot(2);
        alloc.free_slot(1);
        alloc.free_slot(2);
        let s3 = alloc.allocate_slot(3); // Should reuse s2 (last freed)
        let s4 = alloc.allocate_slot(4); // Should reuse s1
        assert_eq!(s3, s2);
        assert_eq!(s4, s1);
    }

    #[test]
    fn test_fast_register_map_available_regs() {
        let map = FastRegisterMap::new();
        let gprs = map.available.get(&RegClassKind::GPR).unwrap();
        assert!(!gprs.is_empty());
        assert!(!gprs.contains(&0)); // x0 not available
    }

    #[test]
    fn test_fast_register_map_allocate_exhaustion() {
        let mut map = FastRegisterMap::new();
        map.available.insert(RegClassKind::GPR, vec![10, 11, 12]);

        let r1 = map.allocate_reg(RegClassKind::GPR).unwrap();
        map.assign(1, r1);
        let r2 = map.allocate_reg(RegClassKind::GPR).unwrap();
        map.assign(2, r2);
        let r3 = map.allocate_reg(RegClassKind::GPR).unwrap();
        map.assign(3, r3);

        // Should be exhausted
        assert!(map.allocate_reg(RegClassKind::GPR).is_none());

        // Free one and reallocate
        map.free_reg(r1);
        let r4 = map.allocate_reg(RegClassKind::GPR);
        assert_eq!(r4, Some(r1));
    }

    #[test]
    fn test_fast_coalescer_same_block() {
        let mut map = FastRegisterMap::new();
        let mut vars = HashMap::new();

        let mut var_src = FastLiveVar::new(1, RegClassKind::GPR, 0, 2);
        var_src.add_use(0, 5);
        let mut var_dst = FastLiveVar::new(2, RegClassKind::GPR, 0, 4);
        var_dst.add_use(0, 8);

        vars.insert(1, var_src);
        vars.insert(2, var_dst);

        map.assign(1, 10);
        map.assign(2, 11);

        let mut coalescer = FastCoalescer::new();
        coalescer.record_copy(2, 1, 0);
        let count = coalescer.coalesce_pending(&mut map, &vars);
        assert!(count <= 1);
    }

    #[test]
    fn test_fast_coalescer_multiple_copies() {
        let mut map = FastRegisterMap::new();
        let mut vars = HashMap::new();

        for i in 1..=4 {
            let mut var = FastLiveVar::new(i, RegClassKind::GPR, 0, i * 2);
            var.add_use(0, i * 2 + 1);
            vars.insert(i, var);
            map.assign(i, 9 + i);
        }

        let mut coalescer = FastCoalescer::new();
        coalescer.record_copy(2, 1, 0);
        coalescer.record_copy(4, 3, 0);
        coalescer.coalesce_pending(&mut map, &vars);
        assert!(coalescer.coalesced <= 2);
    }

    #[test]
    fn test_fast_block_allocator_process_sequence() {
        let parent = FastRegAlloc::new();
        let mut block_alloc = FastBlockAllocator::new(parent, 0);

        let r1 = block_alloc.process_def(1, RegClassKind::GPR);
        assert!(r1.is_some());

        block_alloc.process_use(1);

        let r2 = block_alloc.process_def(2, RegClassKind::GPR);
        assert!(r2.is_some());
        assert_ne!(r1, r2);

        block_alloc.process_use(2);
        block_alloc.process_copy(2, 1);

        let result = block_alloc.finish();
        assert_eq!(result.stats.allocated_vregs, 2);
    }

    #[test]
    fn test_fast_block_allocator_spill_under_pressure() {
        let parent = FastRegAlloc::new();
        let mut block_alloc = FastBlockAllocator::new(parent, 0);

        // Only 2 registers available
        block_alloc
            .parent
            .reg_map
            .available
            .insert(RegClassKind::GPR, vec![10, 11]);

        let r1 = block_alloc.process_def(1, RegClassKind::GPR);
        let r2 = block_alloc.process_def(2, RegClassKind::GPR);
        let r3 = block_alloc.process_def(3, RegClassKind::GPR);

        assert!(r1.is_some());
        assert!(r2.is_some());
        // r3 should have spilled
        assert!(r3.is_none());
    }

    #[test]
    fn test_fast_alloc_stats_field_updates() {
        let mut alloc = FastRegAlloc::new();
        alloc.record_def(1, RegClassKind::GPR, 0, 0);
        alloc.record_use(1, 0, 1);
        alloc.allocate();

        assert_eq!(alloc.stats.total_vregs, 1);
        assert!(alloc.stats.allocated_vregs <= 1);
        assert!(alloc.frame_size() >= 0);
    }

    #[test]
    fn test_fast_reg_alloc_fp_reg_class() {
        let mut alloc = FastRegAlloc::new();
        alloc.record_def(50, RegClassKind::FPR64, 0, 0);
        alloc.record_use(50, 0, 1);
        alloc.allocate();

        // FPR should have been allocated
        let assigned = alloc.get_assignment(50);
        assert!(assigned.is_some() || alloc.is_spilled(50));
    }

    #[test]
    fn test_fast_reg_alloc_vec_reg_class() {
        let mut alloc = FastRegAlloc::new();
        alloc.record_def(70, RegClassKind::VecReg, 0, 0);
        alloc.record_use(70, 0, 1);
        alloc.allocate();

        let assigned = alloc.get_assignment(70);
        assert!(assigned.is_some() || alloc.is_spilled(70));
    }
}