llvm-native-core-ext 0.1.0

//! Machine CSE — Common Subexpression Elimination on machine instructions.
//! Clean-room behavioral reconstruction.
//!
//! @llvm_behavior: MachineCSE eliminates redundant machine instructions
//! that compute the same value. It operates within a single basic block
//! (local CSE) or across basic blocks (global CSE using available
//! expressions analysis).
//!
//! Two machine instructions are congruent if they:
//! 1. Have the same opcode
//! 2. Have the same operands (after renaming to canonical virtual registers)
//! 3. Produce the same result for all possible inputs (pure computation)
//!
//! Algorithm (local CSE per basic block):
//! 1. Walk instructions in order
//! 2. Hash each instruction based on opcode + operands
//! 3. Look up the hash in a table of available expressions
//! 4. If a congruent instruction is found and still dominates the current
//!    instruction (its definition is still live), replace the current
//!    instruction's uses with the earlier definition
//! 5. Remove the redundant instruction
//! 6. Otherwise, add the instruction to the available-expression table
//!
//! Important constraints for machine CSE:
//! - Must respect physical register liveness
//! - Must not eliminate instructions with side effects
//! - Must handle memory operations conservatively (may alias)
//! - Must preserve debug information

use llvm_native_core::codegen::{MachineBasicBlock, MachineFunction, MachineInstr, MachineOperand, VirtReg};
use std::collections::{HashMap, HashSet};

// ============================================================================
// Expression Key — Unique identifier for a computation
// ============================================================================

/// A hashable key representing a machine instruction expression.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
struct ExprKey {
    /// The instruction opcode.
    opcode: u32,
    /// The canonical operands (register numbers, immediate values).
    operands: Vec<ExprOperand>,
}

/// A canonical operand representation for hashing.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
enum ExprOperand {
    /// A register (virtual or physical).
    Reg(u64),
    /// An immediate value.
    Imm(i64),
    /// A label target.
    Label(String),
    /// A global symbol.
    Global(String),
}

impl From<&MachineOperand> for ExprOperand {
    fn from(op: &MachineOperand) -> Self {
        match op {
            MachineOperand::Reg(vr) => ExprOperand::Reg(*vr as u64),
            MachineOperand::PhysReg(pr) => ExprOperand::Reg((*pr as u64) + (1 << 32)),
            MachineOperand::Imm(v) => ExprOperand::Imm(*v),
            MachineOperand::Label(s) => ExprOperand::Label(s.clone()),
            MachineOperand::Global(s) => ExprOperand::Global(s.clone()),
        }
    }
}

// ============================================================================
// Available Expression Entry
// ============================================================================

/// An entry in the available-expression table.
#[derive(Debug, Clone)]
struct AvailExpr {
    /// The virtual register defined by this expression.
    def_vreg: Option<VirtReg>,
    /// Index of the instruction in the basic block.
    instr_idx: usize,
    /// Whether the expression's definition is still live.
    is_live: bool,
}

// ============================================================================
// Machine CSE Pass
// ============================================================================

/// MachineCSE — Common Subexpression Elimination at the machine level.
pub struct MachineCSE {
    /// Number of redundant instructions eliminated.
    pub eliminated: usize,
    /// Whether CSE is enabled.
    pub enabled: bool,
}

impl MachineCSE {
    /// Create a new MachineCSE pass.
    pub fn new() -> Self {
        Self {
            eliminated: 0,
            enabled: true,
        }
    }

    /// Run MachineCSE on a machine function.
    ///
    /// Returns the number of common subexpressions eliminated.
    pub fn run_on_function(&mut self, mf: &mut MachineFunction) -> usize {
        if !self.enabled {
            return 0;
        }

        self.eliminated = 0;

        let block_count = mf.blocks.len();
        for bb_idx in 0..block_count {
            // We need to process each block independently
            self.eliminate_common_subexprs(mf, bb_idx);
        }

        self.eliminated
    }

    // ========================================================================
    // Block-level CSE
    // ========================================================================

    /// Eliminate common subexpressions within a single basic block.
    pub fn eliminate_common_subexprs(&mut self, mf: &mut MachineFunction, bb: usize) {
        if bb >= mf.blocks.len() {
            return;
        }

        let mut avail: HashMap<ExprKey, AvailExpr> = HashMap::new();
        let block = &mf.blocks[bb];
        let num_insts = block.instructions.len();

        // We need to know which instructions to remove.
        // Strategy: first pass — identify redundancies.
        let mut redundant: Vec<usize> = Vec::new();
        let mut replacements: Vec<(usize, usize)> = Vec::new(); // (redundant_idx, canonical_idx)

        for (idx, mi) in block.instructions.iter().enumerate() {
            // Skip instructions with side effects
            if self.has_side_effects(mi) {
                // Side-effecting instructions must not be CSE'd.
                // However, they may kill available expressions if they
                // modify memory that expressions depend on.
                self.kill_memory_expressions(&mut avail);
                continue;
            }

            // Skip instructions that don't define a value
            if mi.def.is_none() {
                // Instructions that don't define a register can still kill
                // memory expressions (e.g., stores)
                if self.is_memory_store(mi) {
                    self.kill_memory_expressions(&mut avail);
                }
                continue;
            }

            let key = self.build_expr_key(mi);

            if key.operands.is_empty() && !self.is_constant_like(mi) {
                // Skip instructions with no operands that aren't constants
                // (they may have implicit operands we can't track)
                self.add_to_avail(&mut avail, key, mi.def, idx);
                continue;
            }

            if let Some(avail_expr) = avail.get(&key) {
                if avail_expr.is_live {
                    // Found a congruent expression! This instruction is redundant.
                    redundant.push(idx);
                    replacements.push((idx, avail_expr.instr_idx));
                } else {
                    // Congruent expression exists but its definition is dead.
                    // Reuse this slot.
                    self.add_to_avail(&mut avail, key, mi.def, idx);
                }
            } else {
                // New expression — add it
                self.add_to_avail(&mut avail, key, mi.def, idx);
            }
        }

        // Second pass: actually remove redundant instructions
        // Process in reverse order to preserve indices
        redundant.sort_by(|a, b| b.cmp(a)); // descending
        for &redundant_idx in &redundant {
            if redundant_idx < mf.blocks[bb].instructions.len() {
                mf.blocks[bb].instructions.remove(redundant_idx);
                self.eliminated += 1;
            }
        }
    }

    // ========================================================================
    // Expression key construction
    // ========================================================================

    /// Build a hashable expression key from a machine instruction.
    fn build_expr_key(&self, mi: &MachineInstr) -> ExprKey {
        let operands: Vec<ExprOperand> = mi.operands.iter().map(ExprOperand::from).collect();

        ExprKey {
            opcode: mi.opcode,
            operands,
        }
    }

    /// Hash a machine instruction into a u64 for quick lookup.
    pub fn hash_instruction(&self, mi: &MachineInstr) -> u64 {
        let mut hash: u64 = (mi.opcode as u64).wrapping_mul(0x9E3779B97F4A7C15);
        hash = hash.rotate_left(17);

        for op in &mi.operands {
            let op_hash = match op {
                MachineOperand::Reg(vr) => (*vr as u64).wrapping_mul(0xBF58476D1CE4E5B9),
                MachineOperand::PhysReg(pr) => {
                    ((*pr as u64) + (1u64 << 32)).wrapping_mul(0x94D049BB133111EB)
                }
                MachineOperand::Imm(v) => (*v as u64).wrapping_mul(0xC6A4A7935BD1E995),
                MachineOperand::Label(s) => Self::hash_string(s).wrapping_mul(0x9E3779B97F4A7C15),
                MachineOperand::Global(s) => Self::hash_string(s).wrapping_mul(0xBF58476D1CE4E5B9),
            };
            hash = hash.wrapping_add(op_hash);
            hash = hash.rotate_left(13);
        }

        if let Some(def) = mi.def {
            // The def register itself doesn't affect the expression identity,
            // but we include it for uniqueness of the key.
            hash = hash.wrapping_add((def as u64).wrapping_mul(0x94D049BB133111EB));
        }

        hash
    }

    /// Simple string hash.
    fn hash_string(s: &str) -> u64 {
        let mut hash: u64 = 0;
        for byte in s.bytes() {
            hash = hash.wrapping_mul(31).wrapping_add(byte as u64);
        }
        hash
    }

    // ========================================================================
    // Congruence checking
    // ========================================================================

    /// Check if two machine instructions are congruent.
    ///
    /// Two instructions are congruent if they have the same opcode and
    /// the same operands (after canonicalization).
    pub fn are_congruent(&self, a: &MachineInstr, b: &MachineInstr) -> bool {
        if a.opcode != b.opcode {
            return false;
        }

        if a.operands.len() != b.operands.len() {
            return false;
        }

        for (op_a, op_b) in a.operands.iter().zip(b.operands.iter()) {
            if !Self::operand_eq(op_a, op_b) {
                return false;
            }
        }

        true
    }

    /// Check if two machine operands are equal.
    fn operand_eq(a: &MachineOperand, b: &MachineOperand) -> bool {
        match (a, b) {
            (MachineOperand::Reg(va), MachineOperand::Reg(vb)) => va == vb,
            (MachineOperand::PhysReg(pa), MachineOperand::PhysReg(pb)) => pa == pb,
            (MachineOperand::Imm(ia), MachineOperand::Imm(ib)) => ia == ib,
            (MachineOperand::Label(la), MachineOperand::Label(lb)) => la == lb,
            (MachineOperand::Global(ga), MachineOperand::Global(gb)) => ga == gb,
            _ => false,
        }
    }

    // ========================================================================
    // Side effect analysis
    // ========================================================================

    /// Check if a machine instruction has side effects that prevent CSE.
    fn has_side_effects(&self, mi: &MachineInstr) -> bool {
        // Branches, calls, and memory stores have side effects
        self.is_branch_or_call(mi) || self.is_memory_store(mi)
    }

    /// Check if an instruction is a branch or call.
    fn is_branch_or_call(&self, mi: &MachineInstr) -> bool {
        // Approximate opcode classification:
        // 7-12: typically branches and calls
        let opcode = mi.opcode;
        (7..=12).contains(&opcode) || opcode == 0 // call often maps to 0 in some models
    }

    /// Check if an instruction is a memory store.
    fn is_memory_store(&self, mi: &MachineInstr) -> bool {
        // Approximate: opcodes 20-21 typically are store operations
        let opcode = mi.opcode;
        opcode == 20 || opcode == 21
    }

    /// Check if an instruction is a memory load.
    fn is_memory_load(&self, mi: &MachineInstr) -> bool {
        let opcode = mi.opcode;
        opcode == 2 || opcode == 3
    }

    /// Check if an instruction is a constant-like value.
    fn is_constant_like(&self, mi: &MachineInstr) -> bool {
        // Instructions with no register operands (only immediates)
        // that define a value are constant-like
        mi.operands
            .iter()
            .all(|op| matches!(op, MachineOperand::Imm(_)))
            && mi.def.is_some()
    }

    // ========================================================================
    // Available expression management
    // ========================================================================

    /// Add an expression to the available-expression table.
    fn add_to_avail(
        &self,
        avail: &mut HashMap<ExprKey, AvailExpr>,
        key: ExprKey,
        def: Option<VirtReg>,
        idx: usize,
    ) {
        avail.insert(
            key,
            AvailExpr {
                def_vreg: def,
                instr_idx: idx,
                is_live: true,
            },
        );
    }

    /// Kill memory-dependent expressions from the available set.
    ///
    /// When a store occurs, expressions that depend on memory (loads)
    /// may no longer be valid. We conservatively kill all memory loads
    /// from the available set.
    fn kill_memory_expressions(&self, avail: &mut HashMap<ExprKey, AvailExpr>) {
        // In a full implementation, we'd use alias analysis to only
        // kill expressions that may alias with the store.
        // Here we conservatively mark all as potentially invalid.
        // For a more precise implementation, we'd only kill loads.
        for entry in avail.values_mut() {
            // Keep the entry but mark it as potentially invalid
            // In practice, store kills loads but not pure computations
        }

        // More precise: remove only load expressions
        avail.retain(|key, _| {
            // Keep non-load expressions
            !self.is_load_opcode(key.opcode)
        });
    }

    /// Check if an opcode corresponds to a memory load.
    fn is_load_opcode(&self, opcode: u32) -> bool {
        opcode == 2 || opcode == 3
    }
}

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

// ============================================================================
// Value Numbering Info (VNInfo)
// ============================================================================

/// VNInfo associates a value number with a machine instruction.
/// Used by global CSE to track congruence classes across blocks.
#[derive(Debug, Clone)]
pub struct VNInfo {
    /// The value number (congruence class ID).
    pub value_number: u32,
    /// Whether this value is available at this point.
    pub available: bool,
    /// The virtual register that holds this value.
    pub vreg: VirtReg,
    /// The block where this value was defined.
    pub defining_block: usize,
    /// The instruction index within the defining block.
    pub defining_instr: usize,
    /// Whether this value is a constant.
    pub is_constant: bool,
    /// The constant value (if applicable).
    pub constant_value: Option<i64>,
}

impl VNInfo {
    /// Create a new VNInfo entry.
    pub fn new(value_number: u32, vreg: VirtReg, block: usize, instr: usize) -> Self {
        Self {
            value_number,
            available: true,
            vreg,
            defining_block: block,
            defining_instr: instr,
            is_constant: false,
            constant_value: None,
        }
    }

    /// Create a VNInfo for a constant value.
    pub fn constant(value_number: u32, value: i64) -> Self {
        Self {
            value_number,
            available: true,
            vreg: 0,
            defining_block: 0,
            defining_instr: 0,
            is_constant: true,
            constant_value: Some(value),
        }
    }

    /// Invalidate this value (e.g., due to a memory clobber).
    pub fn invalidate(&mut self) {
        self.available = false;
    }
}

/// ValueNumberTable maps expression keys to value numbers for global CSE.
#[derive(Debug, Clone)]
pub struct ValueNumberTable {
    /// Mapping from expression key to value number.
    pub key_to_vn: HashMap<ExprKey, u32>,
    /// Mapping from value number to VNInfo.
    pub vn_to_info: HashMap<u32, VNInfo>,
    /// Next value number to assign.
    pub next_vn: u32,
}

impl ValueNumberTable {
    /// Create a new value number table.
    pub fn new() -> Self {
        Self {
            key_to_vn: HashMap::new(),
            vn_to_info: HashMap::new(),
            next_vn: 1,
        }
    }

    /// Look up or create a value number for an expression.
    pub fn lookup_or_create(
        &mut self,
        key: &ExprKey,
        vreg: VirtReg,
        block: usize,
        instr: usize,
    ) -> (u32, bool) {
        if let Some(&vn) = self.key_to_vn.get(key) {
            // Existing value number found
            if let Some(info) = self.vn_to_info.get(&vn) {
                if info.available {
                    return (vn, true);
                }
            }
        }

        // Create new value number
        let vn = self.next_vn;
        self.next_vn += 1;
        self.key_to_vn.insert(key.clone(), vn);
        self.vn_to_info
            .insert(vn, VNInfo::new(vn, vreg, block, instr));
        (vn, false)
    }

    /// Look up a value number for an expression key.
    pub fn lookup(&self, key: &ExprKey) -> Option<u32> {
        self.key_to_vn.get(key).copied()
    }

    /// Get VNInfo for a value number.
    pub fn get_info(&self, vn: u32) -> Option<&VNInfo> {
        self.vn_to_info.get(&vn)
    }

    /// Invalidate all values (e.g., after a memory write).
    pub fn invalidate_all(&mut self) {
        for info in self.vn_to_info.values_mut() {
            info.available = false;
        }
    }
}

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

// ============================================================================
// Def-Use Chain Analysis for Machine Instructions
// ============================================================================

/// DefUseInfo tracks definition-use relationships for virtual registers.
#[derive(Debug, Clone)]
pub struct DefUseInfo {
    /// The virtual register.
    pub vreg: VirtReg,
    /// Position of the single definition.
    pub def_pos: Option<(usize, usize)>,
    /// All use positions.
    pub uses: Vec<(usize, usize)>,
    /// Whether this register is live across blocks.
    pub is_global: bool,
    /// Whether this is a single-def, single-use temp.
    pub is_single_use: bool,
}

impl DefUseInfo {
    /// Create a new def-use info entry.
    pub fn new(vreg: VirtReg) -> Self {
        Self {
            vreg,
            def_pos: None,
            uses: Vec::new(),
            is_global: false,
            is_single_use: false,
        }
    }

    /// Check if this value can be trivially forwarded.
    pub fn is_trivially_replaceable(&self) -> bool {
        self.def_pos.is_some() && self.uses.len() == 1 && !self.is_global
    }

    /// Check if this is a dead definition (no uses).
    pub fn is_dead(&self) -> bool {
        self.def_pos.is_some() && self.uses.is_empty() && !self.is_global
    }
}

/// DefUseChain computes def-use chains for all virtual registers
/// within a machine function.
pub struct DefUseChain {
    /// Per-register def-use info.
    pub chains: HashMap<VirtReg, DefUseInfo>,
}

impl DefUseChain {
    /// Create a new def-use chain analyzer.
    pub fn new() -> Self {
        Self {
            chains: HashMap::new(),
        }
    }

    /// Compute def-use chains for all registers in the function.
    pub fn compute(&mut self, mf: &MachineFunction) {
        self.chains.clear();

        for (block_idx, block) in mf.blocks.iter().enumerate() {
            for (instr_idx, instr) in block.instructions.iter().enumerate() {
                // Record definition
                if let Some(def_vreg) = instr.def {
                    let entry = self
                        .chains
                        .entry(def_vreg)
                        .or_insert_with(|| DefUseInfo::new(def_vreg));
                    if entry.def_pos.is_none() {
                        entry.def_pos = Some((block_idx, instr_idx));
                    } else {
                        // Multiple definitions — mark as global
                        entry.is_global = true;
                    }
                }

                // Record uses
                for op in &instr.operands {
                    if let MachineOperand::Reg(vreg) = *op {
                        let entry = self
                            .chains
                            .entry(vreg)
                            .or_insert_with(|| DefUseInfo::new(vreg));
                        entry.uses.push((block_idx, instr_idx));
                    }
                }
            }

            // Check if any register is live-out (used in successor blocks)
            for &succ_idx in &block.successors {
                {
                    if succ_idx != block_idx {
                        // Mark all defs in this block that reach uses in
                        // successor blocks as global
                        for instr in &block.instructions {
                            if let Some(def_vreg) = instr.def {
                                if let Some(entry) = self.chains.get_mut(&def_vreg) {
                                    if entry.uses.iter().any(|(bi, _)| *bi == succ_idx) {
                                        entry.is_global = true;
                                    }
                                }
                            }
                        }
                    }
                }
            }
        }

        // Set single-use flag
        for entry in self.chains.values_mut() {
            entry.is_single_use = entry.uses.len() == 1;
        }
    }

    /// Get the def-use info for a virtual register.
    pub fn get_info(&self, vreg: VirtReg) -> Option<&DefUseInfo> {
        self.chains.get(&vreg)
    }

    /// Check if a definition dominates all uses.
    pub fn dominates_all_uses(&self, vreg: VirtReg) -> bool {
        if let Some(info) = self.chains.get(&vreg) {
            if let Some((def_block, _def_instr)) = info.def_pos {
                info.uses
                    .iter()
                    .all(|(use_block, _)| *use_block >= def_block)
            } else {
                false
            }
        } else {
            false
        }
    }

    /// Get all registers that are dead (defined but never used).
    pub fn dead_registers(&self) -> Vec<VirtReg> {
        self.chains
            .iter()
            .filter(|(_, info)| info.is_dead())
            .map(|(&vreg, _)| vreg)
            .collect()
    }

    /// Get all registers that are trivially replaceable.
    pub fn trivially_replaceable(&self) -> Vec<VirtReg> {
        self.chains
            .iter()
            .filter(|(_, info)| info.is_trivially_replaceable())
            .map(|(&vreg, _)| vreg)
            .collect()
    }
}

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

// ============================================================================
// Register Liveness-Aware CSE
// ============================================================================

/// LivenessAwareCSE extends MachineCSE with register liveness tracking.
/// It ensures that CSE replacements only occur when the original
/// definition is still live at the point of the redundant computation.
pub struct LivenessAwareCSE {
    /// Core CSE pass.
    pub cse: MachineCSE,
    /// Def-use chains for the function.
    pub def_use: DefUseChain,
    /// Per-block live register sets at each instruction.
    pub live_sets: Vec<Vec<HashSet<VirtReg>>>,
    /// Whether liveness has been computed.
    pub liveness_computed: bool,
}

impl LivenessAwareCSE {
    /// Create a new liveness-aware CSE pass.
    pub fn new() -> Self {
        Self {
            cse: MachineCSE::new(),
            def_use: DefUseChain::new(),
            live_sets: Vec::new(),
            liveness_computed: false,
        }
    }

    /// Run CSE with liveness awareness.
    pub fn run_on_function(&mut self, mf: &mut MachineFunction) -> usize {
        // Step 1: Compute def-use chains
        self.def_use.compute(mf);

        // Step 2: Compute liveness per instruction
        self.compute_liveness(mf);

        // Step 3: Run CSE with liveness checks
        let mut eliminated = 0;

        for (block_idx, block) in mf.blocks.iter_mut().enumerate() {
            let mut avail: HashMap<ExprKey, (VirtReg, usize)> = HashMap::new();

            // We need to track which instructions to remove
            let mut to_remove: Vec<usize> = Vec::new();
            let mut replacements: Vec<(usize, VirtReg)> = Vec::new();

            for (instr_idx, instr) in block.instructions.iter().enumerate() {
                if self.cse.has_side_effects(instr) {
                    // Side-effecting instructions kill available expressions
                    // that may be affected (memory ops kill memory expressions)
                    if self.cse.is_memory_store(instr) {
                        // Kill memory-dependent expressions (opcodes 2,3 = loads)
                        avail.retain(|key, _| key.opcode != 2 && key.opcode != 3);
                    }
                    continue;
                }

                let key = self.cse.build_expr_key(instr);

                if let Some(&(def_vreg, _def_instr_idx)) = avail.get(&key) {
                    // Found congruent expression — check if the original
                    // definition is still live at this point
                    if self.is_live_at(def_vreg, block_idx, instr_idx) {
                        // Replace use: forward the earlier definition
                        replacements.push((instr_idx, def_vreg));
                        to_remove.push(instr_idx);
                        eliminated += 1;
                        continue;
                    }
                }

                // Add to available expressions
                if let Some(def_vreg) = instr.def {
                    avail.insert(key, (def_vreg, instr_idx));
                }
            }

            // Apply replacements: update uses of the removed definitions
            // First, collect the def_vreg for each removed instruction
            let def_map: HashMap<usize, VirtReg> = replacements
                .iter()
                .map(|&(removed_idx, _)| {
                    let def = block.instructions[removed_idx].def;
                    (removed_idx, def)
                })
                .filter(|(_, def)| def.is_some())
                .map(|(idx, def)| (idx, def.unwrap()))
                .collect();

            for &(removed_idx, replacement_vreg) in &replacements {
                let removed_def = def_map.get(&removed_idx).copied();
                for instr in block.instructions.iter_mut().skip(removed_idx + 1) {
                    for op in &mut instr.operands {
                        if let MachineOperand::Reg(vreg) = *op {
                            if Some(vreg) == removed_def {
                                *op = MachineOperand::Reg(replacement_vreg);
                            }
                        }
                    }
                }
            }

            // Remove redundant instructions (in reverse order)
            to_remove.sort_unstable();
            to_remove.dedup();
            for &idx in to_remove.iter().rev() {
                if idx < block.instructions.len() {
                    block.instructions.remove(idx);
                }
            }
        }

        self.cse.eliminated += eliminated;
        eliminated
    }

    /// Compute live sets for each instruction in each block.
    fn compute_liveness(&mut self, mf: &MachineFunction) {
        self.live_sets.clear();

        for block in &mf.blocks {
            let mut block_live: Vec<HashSet<VirtReg>> = Vec::new();
            let mut live: HashSet<VirtReg> = HashSet::new();

            // Process in reverse within block to compute liveness
            for instr in block.instructions.iter().rev() {
                // Record current live set
                block_live.push(live.clone());

                // A definition kills the register (remove from live set)
                if let Some(def_vreg) = instr.def {
                    live.remove(&def_vreg);
                }

                // Uses add registers to the live set
                for op in &instr.operands {
                    if let MachineOperand::Reg(vreg) = *op {
                        live.insert(vreg);
                    }
                }
            }

            // Reverse back so index 0 = before first instruction
            block_live.reverse();
            self.live_sets.push(block_live);
        }

        self.liveness_computed = true;
    }

    /// Check if a virtual register is live at a given block/instruction point.
    fn is_live_at(&self, vreg: VirtReg, block_idx: usize, instr_idx: usize) -> bool {
        if let Some(block_live) = self.live_sets.get(block_idx) {
            if let Some(live_set) = block_live.get(instr_idx) {
                return live_set.contains(&vreg);
            }
        }
        false
    }

    /// Print liveness debug info.
    pub fn print_liveness(&self) {
        for (block_idx, block_live) in self.live_sets.iter().enumerate() {
            eprintln!("Block {}:", block_idx);
            for (instr_idx, live_set) in block_live.iter().enumerate() {
                if !live_set.is_empty() {
                    eprintln!(
                        "  instr {}: live = {:?}",
                        instr_idx,
                        live_set.iter().collect::<Vec<_>>()
                    );
                }
            }
        }
    }
}

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

// ============================================================================
// Global Machine CSE with Value Numbering
// ============================================================================

/// GlobalMachineCSE extends local CSE to operate across basic blocks
/// using value numbering and availability analysis.
///
/// Algorithm:
///   1. Compute value numbers for all expressions in all blocks.
///   2. Propagate availability of expressions across CFG edges.
///   3. For each instruction, check if a congruent expression is available
///      from a dominating definition.
///   4. Replace redundant instructions with copies from the dominating def.
pub struct GlobalMachineCSE {
    /// Core CSE pass.
    pub local_cse: MachineCSE,
    /// Value number table.
    pub vn_table: ValueNumberTable,
    /// Set of expressions available at the start of each block.
    pub avail_in: Vec<HashSet<u32>>,
    /// Set of expressions available at the end of each block.
    pub avail_out: Vec<HashSet<u32>>,
    /// Number of global redundancies eliminated.
    pub global_eliminated: usize,
}

impl GlobalMachineCSE {
    /// Create a new global CSE pass.
    pub fn new() -> Self {
        Self {
            local_cse: MachineCSE::new(),
            vn_table: ValueNumberTable::new(),
            avail_in: Vec::new(),
            avail_out: Vec::new(),
            global_eliminated: 0,
        }
    }

    /// Run global CSE on a machine function.
    pub fn run_on_function(&mut self, mf: &mut MachineFunction) -> usize {
        self.global_eliminated = 0;

        if mf.blocks.is_empty() {
            return 0;
        }

        // Step 1: Build value numbers for all expressions
        self.build_value_numbers(mf);

        // Step 2: Compute availability via iterative dataflow
        self.compute_availability(mf);

        // Step 3: Eliminate globally redundant expressions
        self.eliminate_global_redundancies(mf);

        self.global_eliminated
    }

    /// Assign value numbers to all expressions in all blocks.
    fn build_value_numbers(&mut self, mf: &mut MachineFunction) {
        self.vn_table = ValueNumberTable::new();

        for (block_idx, block) in mf.blocks.iter_mut().enumerate() {
            for (instr_idx, instr) in block.instructions.iter().enumerate() {
                if self.local_cse.has_side_effects(instr) {
                    // Invalidate memory-related value numbers on stores
                    if self.local_cse.is_memory_store(instr) {
                        self.vn_table.invalidate_all();
                    }
                    continue;
                }

                let key = self.local_cse.build_expr_key(instr);
                let def_vreg = instr.def.unwrap_or(0);

                // Look up or create value number
                self.vn_table
                    .lookup_or_create(&key, def_vreg, block_idx, instr_idx);
            }
        }
    }

    /// Compute available expressions at block boundaries using iterative
    /// dataflow analysis.
    fn compute_availability(&mut self, mf: &MachineFunction) {
        let n = mf.blocks.len();

        // Initialize: avail_in[0] = empty, others = all VNs (top)
        self.avail_in = vec![HashSet::new(); n];
        self.avail_out = vec![HashSet::new(); n];

        // Collect all value numbers
        let all_vns: HashSet<u32> = self.vn_table.vn_to_info.keys().copied().collect();

        // Initialize all non-entry blocks with all VNs (optimistic)
        for i in 1..n {
            self.avail_in[i] = all_vns.clone();
        }

        // Iterative dataflow
        let mut changed = true;
        let mut iterations = 0;

        while changed && iterations < 100 {
            changed = false;
            iterations += 1;

            for block_idx in 0..n {
                // Compute avail_in from predecessors' avail_out
                let mut new_in: HashSet<u32> = if block_idx == 0 {
                    HashSet::new()
                } else {
                    // Intersection of predecessors' avail_out
                    let pred_indices: Vec<usize> = (0..n)
                        .filter(|&i| mf.blocks[i].successors.contains(&block_idx))
                        .collect();

                    if pred_indices.is_empty() {
                        HashSet::new()
                    } else {
                        let mut inter = self.avail_out[pred_indices[0]].clone();
                        for &pi in &pred_indices[1..] {
                            inter = inter.intersection(&self.avail_out[pi]).copied().collect();
                        }
                        inter
                    }
                };

                // Compute avail_out = GEN ∪ (avail_in - KILL)
                let mut new_out = new_in.clone();

                for instr in &mf.blocks[block_idx].instructions {
                    if self.local_cse.has_side_effects(instr) {
                        if self.local_cse.is_memory_store(instr) {
                            // Kill all memory-dependent expressions
                            new_out.clear();
                        }
                        continue;
                    }

                    let key = self.local_cse.build_expr_key(instr);
                    if let Some(vn) = self.vn_table.lookup(&key) {
                        new_out.insert(vn); // GEN
                    }
                }

                if new_in != self.avail_in[block_idx] {
                    self.avail_in[block_idx] = new_in;
                    changed = true;
                }
                if new_out != self.avail_out[block_idx] {
                    self.avail_out[block_idx] = new_out;
                    changed = true;
                }
            }
        }
    }

    /// Eliminate globally redundant expressions.
    fn eliminate_global_redundancies(&mut self, mf: &mut MachineFunction) {
        for (block_idx, block) in mf.blocks.iter_mut().enumerate() {
            let mut avail = self.avail_in[block_idx].clone();
            let mut to_remove: Vec<usize> = Vec::new();

            for (instr_idx, instr) in block.instructions.iter().enumerate() {
                if self.local_cse.has_side_effects(instr) {
                    if self.local_cse.is_memory_store(instr) {
                        avail.clear();
                    }
                    continue;
                }

                let key = self.local_cse.build_expr_key(instr);
                if let Some(vn) = self.vn_table.lookup(&key) {
                    if avail.contains(&vn) {
                        // Redundant! Mark for removal
                        to_remove.push(instr_idx);
                        self.global_eliminated += 1;
                    } else {
                        // Not redundant — add to available set
                        avail.insert(vn);
                    }
                }
            }

            // Remove redundant instructions
            to_remove.sort_unstable();
            to_remove.dedup();
            for &idx in to_remove.iter().rev() {
                if idx < block.instructions.len() {
                    block.instructions.remove(idx);
                }
            }
        }
    }

    /// Get statistics for the global CSE pass.
    pub fn print_stats(&self) {
        eprintln!(
            "GlobalMachineCSE: {} global redundancies eliminated",
            self.global_eliminated
        );
        eprintln!("  Value numbers: {}", self.vn_table.vn_to_info.len());
        eprintln!("  Expression keys: {}", self.vn_table.key_to_vn.len());
    }
}

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

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

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

    fn make_mi(opcode: u32, def: Option<u32>, operands: Vec<MachineOperand>) -> MachineInstr {
        MachineInstr {
            opcode,
            operands,
            def,
        }
    }

    fn make_alu(def: u32, op1: u32, op2: u32) -> MachineInstr {
        MachineInstr {
            opcode: 1,
            operands: vec![MachineOperand::Reg(op1), MachineOperand::Reg(op2)],
            def: Some(def),
        }
    }

    #[test]
    fn test_machine_cse_new() {
        let cse = MachineCSE::new();
        assert_eq!(cse.eliminated, 0);
        assert!(cse.enabled);
    }

    #[test]
    fn test_default() {
        let cse = MachineCSE::default();
        assert!(cse.enabled);
    }

    #[test]
    fn test_hash_instruction_different() {
        let cse = MachineCSE::new();
        let mi1 = make_mi(1, Some(1), vec![MachineOperand::Reg(2)]);
        let mi2 = make_mi(2, Some(1), vec![MachineOperand::Reg(2)]);
        assert_ne!(cse.hash_instruction(&mi1), cse.hash_instruction(&mi2));
    }

    #[test]
    fn test_hash_instruction_same() {
        let cse = MachineCSE::new();
        let mi1 = make_mi(1, Some(1), vec![MachineOperand::Reg(2)]);
        let mi2 = make_mi(1, Some(1), vec![MachineOperand::Reg(2)]);
        // Same opcode, same operands, same def → same hash
        assert_eq!(cse.hash_instruction(&mi1), cse.hash_instruction(&mi2));
    }

    #[test]
    fn test_are_congruent_true() {
        let cse = MachineCSE::new();
        let mi1 = make_mi(
            1,
            Some(1),
            vec![MachineOperand::Reg(2), MachineOperand::Imm(5)],
        );
        let mi2 = make_mi(
            1,
            Some(3),
            vec![MachineOperand::Reg(2), MachineOperand::Imm(5)],
        );
        assert!(cse.are_congruent(&mi1, &mi2));
    }

    #[test]
    fn test_are_congruent_false_opcode() {
        let cse = MachineCSE::new();
        let mi1 = make_mi(1, Some(1), vec![MachineOperand::Reg(2)]);
        let mi2 = make_mi(2, Some(1), vec![MachineOperand::Reg(2)]);
        assert!(!cse.are_congruent(&mi1, &mi2));
    }

    #[test]
    fn test_are_congruent_false_operands() {
        let cse = MachineCSE::new();
        let mi1 = make_mi(1, Some(1), vec![MachineOperand::Reg(2)]);
        let mi2 = make_mi(1, Some(1), vec![MachineOperand::Reg(3)]);
        assert!(!cse.are_congruent(&mi1, &mi2));
    }

    #[test]
    fn test_are_congruent_different_len() {
        let cse = MachineCSE::new();
        let mi1 = make_mi(1, Some(1), vec![MachineOperand::Reg(2)]);
        let mi2 = make_mi(
            1,
            Some(1),
            vec![MachineOperand::Reg(2), MachineOperand::Reg(3)],
        );
        assert!(!cse.are_congruent(&mi1, &mi2));
    }

    #[test]
    fn test_has_side_effects_branch() {
        let cse = MachineCSE::new();
        let mi = make_mi(7, None, vec![]);
        assert!(cse.has_side_effects(&mi));
    }

    #[test]
    fn test_has_side_effects_store() {
        let cse = MachineCSE::new();
        let mi = make_mi(20, None, vec![MachineOperand::Reg(1)]);
        assert!(cse.has_side_effects(&mi));
    }

    #[test]
    fn test_has_side_effects_alu() {
        let cse = MachineCSE::new();
        let mi = make_mi(1, Some(1), vec![MachineOperand::Reg(2)]);
        assert!(!cse.has_side_effects(&mi));
    }

    #[test]
    fn test_is_memory_load() {
        let cse = MachineCSE::new();
        assert!(cse.is_memory_load(&make_mi(2, None, vec![])));
        assert!(!cse.is_memory_load(&make_mi(1, None, vec![])));
    }

    #[test]
    fn test_is_memory_store() {
        let cse = MachineCSE::new();
        assert!(cse.is_memory_store(&make_mi(20, None, vec![])));
        assert!(!cse.is_memory_store(&make_mi(2, None, vec![])));
    }

    #[test]
    fn test_operand_eq() {
        assert!(MachineCSE::operand_eq(
            &MachineOperand::Reg(5),
            &MachineOperand::Reg(5)
        ));
        assert!(!MachineCSE::operand_eq(
            &MachineOperand::Reg(5),
            &MachineOperand::Reg(3)
        ));
        assert!(MachineCSE::operand_eq(
            &MachineOperand::Imm(42),
            &MachineOperand::Imm(42)
        ));
        assert!(!MachineCSE::operand_eq(
            &MachineOperand::Imm(1),
            &MachineOperand::Imm(2)
        ));
        assert!(!MachineCSE::operand_eq(
            &MachineOperand::Reg(1),
            &MachineOperand::Imm(1)
        ));
    }

    #[test]
    fn test_eliminate_block_no_redundancy() {
        let mut cse = MachineCSE::new();
        let mut mf = MachineFunction::new("func");
        let mbb = MachineBasicBlock {
            name: "entry".to_string(),
            instructions: vec![make_alu(1, 10, 11), make_alu(2, 12, 13)],
            successors: Vec::new(),
        };
        mf.push_block(mbb);

        cse.eliminate_common_subexprs(&mut mf, 0);
        assert_eq!(cse.eliminated, 0);
        assert_eq!(mf.blocks[0].instructions.len(), 2);
    }

    #[test]
    fn test_eliminate_block_with_redundancy() {
        let mut cse = MachineCSE::new();
        let mut mf = MachineFunction::new("func");
        let mbb = MachineBasicBlock {
            name: "entry".to_string(),
            instructions: vec![
                make_alu(1, 10, 11),
                make_alu(2, 10, 11), // Redundant: same as first
            ],
            successors: Vec::new(),
        };
        mf.push_block(mbb);

        cse.eliminate_common_subexprs(&mut mf, 0);
        assert_eq!(cse.eliminated, 1);
        assert_eq!(mf.blocks[0].instructions.len(), 1);
    }

    #[test]
    fn test_run_on_function_disabled() {
        let mut cse = MachineCSE::new();
        cse.enabled = false;
        let mut mf = MachineFunction::new("func");
        let mbb = MachineBasicBlock {
            name: "entry".to_string(),
            instructions: vec![make_alu(1, 10, 11), make_alu(2, 10, 11)],
            successors: Vec::new(),
        };
        mf.push_block(mbb);

        let count = cse.run_on_function(&mut mf);
        assert_eq!(count, 0);
        assert_eq!(mf.blocks[0].instructions.len(), 2);
    }

    #[test]
    fn test_run_on_empty_function() {
        let mut cse = MachineCSE::new();
        let mut mf = MachineFunction::new("empty");
        let count = cse.run_on_function(&mut mf);
        assert_eq!(count, 0);
    }

    #[test]
    fn test_hash_string() {
        let h1 = MachineCSE::hash_string("hello");
        let h2 = MachineCSE::hash_string("hello");
        let h3 = MachineCSE::hash_string("world");
        assert_eq!(h1, h2);
        assert_ne!(h1, h3);
    }
}