llvm-native-core 0.1.15

//! BPF Instruction Selection
//!
//! Translates LLVM IR (in SelectionDAG form) to BPF MC instructions.
//! This is the core lowering pass for the BPF backend, converting
//! target-independent IR operations into the BPF instruction set.
//!
//! ## Selection Strategy
//!
//! BPF is a simple RISC-like ISA with a small number of operations.
//! Selection maps:
//!
//! - Integer arithmetic → ALU64/ALU32 instructions
//! - Memory operations → LDX/STX instructions
//! - Control flow → conditional jumps (JEQ, JGT, etc.)
//! - Function calls → CALL
//! - Constants → MOV with immediate, or LD_DW for 64-bit values
//!
//! ## Register Allocation
//!
//! This instruction selector uses a simple virtual-register-to-physical
//! mapping strategy. In a full backend, register allocation would be
//! handled by a separate regalloc pass; here we perform a basic
//! greedy assignment.
//!
//! Clean-room implementation from the BPF ISA specification,
//! the LLVM SelectionDAG documentation, and the BPF backend
//! architecture description. No LLVM C++ source code was consulted.

use super::bpf_instr_info::{reg_name, BpfInstrInfo, BpfOpcode, BPF_REG_10};
use super::bpf_mc_encoder::BpfInstr;
use std::collections::HashMap;

// ============================================================================
// Selection DAG node types (simplified)
// ============================================================================

/// A simplified SelectionDAG node representing an LLVM IR operation.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SDNodeKind {
    // Constants
    ConstantI32(i32),
    ConstantI64(i64),
    ConstantF32(u32),
    ConstantF64(u64),
    // Memory
    Load,
    Store,
    // Arithmetic
    Add,
    Sub,
    Mul,
    SDiv,
    UDiv,
    And,
    Or,
    Xor,
    Shl,
    Sra,
    Srl,
    // Comparison
    IcmpEq,
    IcmpNe,
    IcmpSgt,
    IcmpSge,
    IcmpSlt,
    IcmpSle,
    IcmpUgt,
    IcmpUge,
    IcmpUlt,
    IcmpUle,
    // Control flow
    Br,
    BrCond,
    // Other
    CopyToReg,
    CopyFromReg,
    Call,
    Ret,
    FrameIndex,
    Phi,
}

/// A node in the selection DAG (simplified representation).
#[derive(Debug, Clone)]
pub struct SDNode {
    pub kind: SDNodeKind,
    pub operands: Vec<usize>,    // indices into DAG node array
    pub value_type: SdValueType, // result type
    pub node_id: usize,
    // For constant nodes
    pub const_value: i64,
}

/// Value types for SelectionDAG nodes.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SdValueType {
    I1,
    I8,
    I16,
    I32,
    I64,
    F32,
    F64,
}

// ============================================================================
// Selection DAG
// ============================================================================

/// A simplified SelectionDAG.
pub struct SelectionDAG {
    pub nodes: Vec<SDNode>,
    pub entry_node: Option<usize>,
    pub root_node: Option<usize>,
}

impl Default for SelectionDAG {
    fn default() -> Self {
        Self {
            nodes: Vec::new(),
            entry_node: None,
            root_node: None,
        }
    }
}

impl SelectionDAG {
    /// Create a new empty DAG.
    pub fn new() -> Self {
        Self::default()
    }

    /// Add a node and return its index.
    pub fn add_node(&mut self, node: SDNode) -> usize {
        let idx = self.nodes.len();
        self.nodes.push(node);
        idx
    }
}

// ============================================================================
// Virtual register tracking
// ============================================================================

/// Tracks virtual register assignments.
pub struct RegTracker {
    /// Next virtual register number.
    next_vreg: u32,
    /// Virtual-to-physical register mapping.
    vreg_to_phys: HashMap<u32, u8>,
    /// Physical register allocation bitmap.
    phys_regs_used: u64,
    /// Stack slot assignments for spilled regs.
    stack_slots: HashMap<u32, i16>,
    /// Next stack slot offset.
    next_stack_slot: i16,
}

impl Default for RegTracker {
    fn default() -> Self {
        Self {
            next_vreg: 0,
            vreg_to_phys: HashMap::new(),
            phys_regs_used: 0,
            stack_slots: HashMap::new(),
            next_stack_slot: -8, // grow downward
        }
    }
}

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

    /// Allocate a new virtual register.
    pub fn new_vreg(&mut self) -> u32 {
        let vreg = self.next_vreg;
        self.next_vreg += 1;
        vreg
    }

    /// Assign a virtual register to a physical register.
    pub fn assign_phys(&mut self, vreg: u32, phys: u8) {
        self.vreg_to_phys.insert(vreg, phys);
        self.phys_regs_used |= 1u64 << phys;
    }

    /// Get the physical register for a virtual register, if assigned.
    pub fn get_phys(&self, vreg: u32) -> Option<u8> {
        self.vreg_to_phys.get(&vreg).copied()
    }

    /// Find a free physical register (r0–r9).
    pub fn find_free_reg(&self) -> Option<u8> {
        for reg in 0..10u8 {
            if self.phys_regs_used & (1u64 << reg) == 0 {
                return Some(reg);
            }
        }
        None
    }

    /// Spill a virtual register to the stack.
    pub fn spill(&mut self, vreg: u32) -> i16 {
        if let Some(&slot) = self.stack_slots.get(&vreg) {
            return slot;
        }
        let slot = self.next_stack_slot;
        self.next_stack_slot -= 8;
        self.stack_slots.insert(vreg, slot);
        slot
    }
}

// ============================================================================
// BPF Instruction Selector
// ============================================================================

/// The BPF instruction selector — converts SelectionDAG nodes to
/// BPF MC instructions.
pub struct BpfISel {
    /// Emitted BPF instructions.
    pub instrs: Vec<BpfInstr>,
    /// Next available virtual register.
    reg_tracker: RegTracker,
    /// DAG node result → virtual register mapping.
    node_to_vreg: HashMap<usize, u32>,
    /// Whether 32-bit ALU operations are available.
    pub has_alu32: bool,
}

impl Default for BpfISel {
    fn default() -> Self {
        Self {
            instrs: Vec::new(),
            reg_tracker: RegTracker::new(),
            node_to_vreg: HashMap::new(),
            has_alu32: true,
        }
    }
}

impl BpfISel {
    /// Create a new BPF instruction selector.
    pub fn new() -> Self {
        Self::default()
    }

    /// Create with alu32 disabled (for BPF v1 compatibility).
    pub fn without_alu32() -> Self {
        Self {
            has_alu32: false,
            ..Self::default()
        }
    }

    /// Run instruction selection on a SelectionDAG.
    pub fn select(&mut self, dag: &SelectionDAG) -> Result<Vec<BpfInstr>, String> {
        self.instructions.clear();
        self.reg_tracker = RegTracker::new();
        self.node_to_vreg.clear();

        // Walk the DAG and emit BPF instructions.
        // A full implementation would do a topological walk with pattern matching.
        // Here we do a simplified linear walk and map each node kind to BPF.

        for (idx, node) in dag.nodes.iter().enumerate() {
            self.select_node(idx, node, dag)?;
        }

        Ok(self.instructions.clone())
    }

    /// Select a single DAG node.
    fn select_node(&mut self, idx: usize, node: &SDNode, dag: &SelectionDAG) -> Result<(), String> {
        match &node.kind {
            SDNodeKind::ConstantI32(val) => {
                let vreg = self.reg_tracker.new_vreg();
                let phys = self.assign_or_find_free(vreg)?;
                self.node_to_vreg.insert(idx, vreg);

                self.instructions.push(BpfInstr {
                    opcode: BpfOpcode::MOV,
                    dst_reg: phys,
                    src_reg: 0,
                    offset: 0,
                    imm: *val,
                });
            }
            SDNodeKind::ConstantI64(val) => {
                let vreg = self.reg_tracker.new_vreg();
                let phys = self.assign_or_find_free(vreg)?;
                self.node_to_vreg.insert(idx, vreg);

                // Use LD_DW for full 64-bit immediate
                self.instructions.push(BpfInstr {
                    opcode: BpfOpcode::LD_DW,
                    dst_reg: phys,
                    src_reg: 0,
                    offset: 0,
                    imm: *val as i32, // low 32 bits; full 64-bit in actual ld dw
                });
            }
            SDNodeKind::Add => {
                self.select_binary(BpfOpcode::ADD, idx, node, dag)?;
            }
            SDNodeKind::Sub => {
                self.select_binary(BpfOpcode::SUB, idx, node, dag)?;
            }
            SDNodeKind::Mul => {
                self.select_binary(BpfOpcode::MUL, idx, node, dag)?;
            }
            SDNodeKind::SDiv => {
                self.select_binary(BpfOpcode::SDIV, idx, node, dag)?;
            }
            SDNodeKind::UDiv => {
                self.select_binary(BpfOpcode::DIV, idx, node, dag)?;
            }
            SDNodeKind::And => {
                self.select_binary(BpfOpcode::AND, idx, node, dag)?;
            }
            SDNodeKind::Or => {
                self.select_binary(BpfOpcode::OR, idx, node, dag)?;
            }
            SDNodeKind::Xor => {
                self.select_binary(BpfOpcode::XOR, idx, node, dag)?;
            }
            SDNodeKind::Shl => {
                self.select_binary(BpfOpcode::LSH, idx, node, dag)?;
            }
            SDNodeKind::Sra => {
                self.select_binary(BpfOpcode::ARSH, idx, node, dag)?;
            }
            SDNodeKind::Srl => {
                self.select_binary(BpfOpcode::RSH, idx, node, dag)?;
            }
            SDNodeKind::IcmpEq => {
                self.select_compare(BpfOpcode::JEQ, idx, node, dag)?;
            }
            SDNodeKind::IcmpNe => {
                self.select_compare(BpfOpcode::JNE, idx, node, dag)?;
            }
            SDNodeKind::IcmpSgt => {
                self.select_compare(BpfOpcode::JSGT, idx, node, dag)?;
            }
            SDNodeKind::IcmpSge => {
                self.select_compare(BpfOpcode::JSGE, idx, node, dag)?;
            }
            SDNodeKind::IcmpSlt => {
                self.select_compare(BpfOpcode::JSLT, idx, node, dag)?;
            }
            SDNodeKind::IcmpSle => {
                self.select_compare(BpfOpcode::JSLE, idx, node, dag)?;
            }
            SDNodeKind::IcmpUgt => {
                self.select_compare(BpfOpcode::JGT, idx, node, dag)?;
            }
            SDNodeKind::IcmpUge => {
                self.select_compare(BpfOpcode::JGE, idx, node, dag)?;
            }
            SDNodeKind::IcmpUlt => {
                self.select_compare(BpfOpcode::JLT, idx, node, dag)?;
            }
            SDNodeKind::IcmpUle => {
                self.select_compare(BpfOpcode::JLE, idx, node, dag)?;
            }
            SDNodeKind::Load => {
                self.select_load(idx, node, dag)?;
            }
            SDNodeKind::Store => {
                self.select_store(idx, node, dag)?;
            }
            SDNodeKind::Ret => {
                // Map return value (operand 0) to r0, then emit EXIT
                if !node.operands.is_empty() {
                    let ret_val_idx = node.operands[0];
                    if let Some(&vreg) = self.node_to_vreg.get(&ret_val_idx) {
                        if let Some(phys) = self.reg_tracker.get_phys(vreg) {
                            if phys != 0 {
                                self.instructions.push(BpfInstr {
                                    opcode: BpfOpcode::MOV,
                                    dst_reg: 0,
                                    src_reg: phys,
                                    offset: 0,
                                    imm: 0,
                                });
                            }
                        }
                    }
                }
                self.instructions.push(BpfInstr {
                    opcode: BpfOpcode::EXIT,
                    dst_reg: 0,
                    src_reg: 0,
                    offset: 0,
                    imm: 0,
                });
            }
            SDNodeKind::Call => {
                self.instructions.push(BpfInstr {
                    opcode: BpfOpcode::CALL,
                    dst_reg: 0,
                    src_reg: 0,
                    offset: 0,
                    imm: node.const_value as i32,
                });
            }
            SDNodeKind::CopyToReg | SDNodeKind::CopyFromReg => {
                // Register copies handled by register assignment
            }
            SDNodeKind::Br => {
                // Unconditional branch — handled by the basic block layout.
                // Emitted as JA when lowering the CFG.
            }
            SDNodeKind::BrCond => {
                // Conditional branch — handled by CFG lowering.
                // The comparison result from Icmp is used with the branch.
            }
            SDNodeKind::FrameIndex => {
                let vreg = self.reg_tracker.new_vreg();
                let phys = self.assign_or_find_free(vreg)?;
                self.node_to_vreg.insert(idx, vreg);

                // FrameIndex is r10 + offset
                let offset = node.const_value as i16;
                self.instructions.push(BpfInstr {
                    opcode: BpfOpcode::ADD,
                    dst_reg: phys,
                    src_reg: BPF_REG_10,
                    offset: 0,
                    imm: offset as i32,
                });
            }
            SDNodeKind::Phi => {
                // Phi nodes are eliminated during SSA destruction.
                // In a simple lowering, they become MOV instructions.
            }
            _ => {
                // Unhandled node types
            }
        }
        Ok(())
    }

    /// Select a binary operation.
    fn select_binary(
        &mut self,
        opcode: BpfOpcode,
        idx: usize,
        node: &SDNode,
        dag: &SelectionDAG,
    ) -> Result<(), String> {
        if node.operands.len() < 2 {
            return Err(format!("Binary op {:?} needs 2 operands", opcode));
        }

        let lhs_idx = node.operands[0];
        let rhs_idx = node.operands[1];

        // Get the vregs for the operands
        self.ensure_operand_materialized(lhs_idx, dag)?;
        self.ensure_operand_materialized(rhs_idx, dag)?;

        let lhs_vreg = *self.node_to_vreg.get(&lhs_idx).unwrap();
        let rhs_vreg = *self.node_to_vreg.get(&rhs_idx).unwrap();

        let lhs_phys = self.reg_tracker.get_phys(lhs_vreg).unwrap();
        let rhs_phys = self.reg_tracker.get_phys(rhs_vreg).unwrap();

        // Result gets a new vreg
        let result_vreg = self.reg_tracker.new_vreg();
        let result_phys = self.assign_or_find_free(result_vreg)?;

        self.node_to_vreg.insert(idx, result_vreg);

        self.instructions.push(BpfInstr {
            opcode,
            dst_reg: result_phys,
            src_reg: rhs_phys,
            offset: 0,
            imm: 0,
        });

        // If the result is in a different register from the lhs, copy lhs to result first
        if result_phys != lhs_phys {
            // Insert a MOV before the operation
            let last = self.instructions.len() - 1;
            self.instructions.insert(
                last,
                BpfInstr {
                    opcode: BpfOpcode::MOV,
                    dst_reg: result_phys,
                    src_reg: lhs_phys,
                    offset: 0,
                    imm: 0,
                },
            );
        }

        Ok(())
    }

    /// Select a comparison operation.
    fn select_compare(
        &mut self,
        branch_op: BpfOpcode,
        idx: usize,
        node: &SDNode,
        dag: &SelectionDAG,
    ) -> Result<(), String> {
        if node.operands.len() < 2 {
            return Err("Compare needs 2 operands".into());
        }

        // Comparisons in BPF are integrated into branches.
        // We emit a conditional jump with the comparison result,
        // and a MOV to set r0 to 0 or 1 for the boolean result.

        let lhs_idx = node.operands[0];
        let rhs_idx = node.operands[1];

        self.ensure_operand_materialized(lhs_idx, dag)?;
        self.ensure_operand_materialized(rhs_idx, dag)?;

        let lhs_vreg = *self.node_to_vreg.get(&lhs_idx).unwrap();
        let rhs_vreg = *self.node_to_vreg.get(&rhs_idx).unwrap();

        let lhs_phys = self.reg_tracker.get_phys(lhs_vreg).unwrap();
        let rhs_phys = self.reg_tracker.get_phys(rhs_vreg).unwrap();

        // Emit: jump to +2 if condition true; mov r0, 0; ja +1; mov r0, 1
        // (The branch target offset +2 is a placeholder; resolved by the MC layer.)
        let vreg = self.reg_tracker.new_vreg();
        let result_phys = self.assign_or_find_free(vreg)?;
        self.node_to_vreg.insert(idx, vreg);

        // Conditional jump: if condition true, skip to "set_true"
        self.instructions.push(BpfInstr {
            opcode: branch_op,
            dst_reg: lhs_phys,
            src_reg: rhs_phys,
            offset: 2, // skip next 2 instructions
            imm: 0,
        });

        // false path: set result to 0
        self.instructions.push(BpfInstr {
            opcode: BpfOpcode::MOV,
            dst_reg: result_phys,
            src_reg: 0,
            offset: 0,
            imm: 0,
        });

        // jump past true path
        self.instructions.push(BpfInstr {
            opcode: BpfOpcode::JA,
            dst_reg: 0,
            src_reg: 0,
            offset: 1,
            imm: 0,
        });

        // true path: set result to 1
        self.instructions.push(BpfInstr {
            opcode: BpfOpcode::MOV,
            dst_reg: result_phys,
            src_reg: 0,
            offset: 0,
            imm: 1,
        });

        Ok(())
    }

    /// Select a load instruction.
    fn select_load(&mut self, idx: usize, node: &SDNode, dag: &SelectionDAG) -> Result<(), String> {
        if node.operands.is_empty() {
            return Err("Load needs address operand".into());
        }

        let addr_idx = node.operands[0];
        self.ensure_operand_materialized(addr_idx, dag)?;
        let addr_vreg = *self.node_to_vreg.get(&addr_idx).unwrap();
        let addr_phys = self.reg_tracker.get_phys(addr_vreg).unwrap();

        let vreg = self.reg_tracker.new_vreg();
        let result_phys = self.assign_or_find_free(vreg)?;
        self.node_to_vreg.insert(idx, vreg);

        self.instructions.push(BpfInstr {
            opcode: BpfOpcode::LDX_DW,
            dst_reg: result_phys,
            src_reg: addr_phys,
            offset: 0,
            imm: 0,
        });

        Ok(())
    }

    /// Select a store instruction.
    fn select_store(
        &mut self,
        idx: usize,
        node: &SDNode,
        dag: &SelectionDAG,
    ) -> Result<(), String> {
        if node.operands.len() < 2 {
            return Err("Store needs address and value operands".into());
        }

        let val_idx = node.operands[0];
        let addr_idx = node.operands[1];

        self.ensure_operand_materialized(val_idx, dag)?;
        self.ensure_operand_materialized(addr_idx, dag)?;

        let val_vreg = *self.node_to_vreg.get(&val_idx).unwrap();
        let addr_vreg = *self.node_to_vreg.get(&addr_idx).unwrap();

        let val_phys = self.reg_tracker.get_phys(val_vreg).unwrap();
        let addr_phys = self.reg_tracker.get_phys(addr_vreg).unwrap();

        self.instructions.push(BpfInstr {
            opcode: BpfOpcode::STX_DW,
            dst_reg: addr_phys,
            src_reg: val_phys,
            offset: 0,
            imm: 0,
        });

        Ok(())
    }

    /// Ensure an operand's value is materialized in a physical register.
    fn ensure_operand_materialized(
        &mut self,
        node_idx: usize,
        _dag: &SelectionDAG,
    ) -> Result<(), String> {
        // If the node already has a vreg assigned, it was already materialized.
        if self.node_to_vreg.contains_key(&node_idx) {
            return Ok(());
        }

        // Otherwise, this is an un-materialized node — in a full implementation
        // we would recursively select it. For now, create a vreg for it.
        let vreg = self.reg_tracker.new_vreg();
        let phys = self.assign_or_find_free(vreg)?;
        self.node_to_vreg.insert(node_idx, vreg);

        Ok(())
    }

    /// Assign a virtual register to a physical register, finding a free one.
    fn assign_or_find_free(&mut self, vreg: u32) -> Result<u8, String> {
        if let Some(phys) = self.reg_tracker.find_free_reg() {
            self.reg_tracker.assign_phys(vreg, phys);
            Ok(phys)
        } else {
            // No free regs — spill to stack in a full implementation
            Err("Register allocation failed: no free registers".into())
        }
    }

    /// Get the number of instructions emitted.
    pub fn instr_count(&self) -> usize {
        self.instructions.len()
    }

    /// Reset the selector state for a new function.
    pub fn reset(&mut self) {
        self.instructions.clear();
        self.reg_tracker = RegTracker::new();
        self.node_to_vreg.clear();
    }
}

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

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

    fn make_dag_with_add() -> SelectionDAG {
        let mut dag = SelectionDAG::new();
        let c1 = dag.add_node(SDNode {
            kind: SDNodeKind::ConstantI32(10),
            operands: vec![],
            value_type: SdValueType::I32,
            node_id: 0,
            const_value: 10,
        });
        let c2 = dag.add_node(SDNode {
            kind: SDNodeKind::ConstantI32(20),
            operands: vec![],
            value_type: SdValueType::I32,
            node_id: 1,
            const_value: 20,
        });
        let add = dag.add_node(SDNode {
            kind: SDNodeKind::Add,
            operands: vec![c1, c2],
            value_type: SdValueType::I32,
            node_id: 2,
            const_value: 0,
        });
        dag.root_node = Some(add);
        dag
    }

    #[test]
    fn test_isel_create() {
        let isel = BpfISel::new();
        assert_eq!(isel.instructions.len(), 0);
        assert!(isel.has_alu32);
    }

    #[test]
    fn test_isel_without_alu32() {
        let isel = BpfISel::without_alu32();
        assert!(!isel.has_alu32);
    }

    #[test]
    fn test_select_add() {
        let dag = make_dag_with_add();
        let mut isel = BpfISel::new();
        let result = isel.select(&dag);
        assert!(result.is_ok());
        let instrs = result.unwrap();
        assert!(!instrs.is_empty());

        // Should contain ADD instruction
        let has_add = instrs.iter().any(|i| i.opcode == BpfOpcode::ADD);
        assert!(has_add, "Expected ADD instruction in output");
    }

    #[test]
    fn test_select_constants() {
        let dag = SelectionDAG::new();
        let mut isel = BpfISel::new();
        // Just test that the selector runs without panic on an empty DAG
        let result = isel.select(&dag);
        assert!(result.is_ok());
        assert_eq!(result.unwrap().len(), 0);
    }

    #[test]
    fn test_instr_count() {
        let dag = make_dag_with_add();
        let mut isel = BpfISel::new();
        let _ = isel.select(&dag);
        assert!(isel.instr_count() > 0);
    }

    #[test]
    fn test_reset() {
        let dag = make_dag_with_add();
        let mut isel = BpfISel::new();
        let _ = isel.select(&dag);
        assert!(isel.instr_count() > 0);

        isel.reset();
        assert_eq!(isel.instr_count(), 0);
    }

    #[test]
    fn test_select_sub() {
        let mut dag = SelectionDAG::new();
        let c1 = dag.add_node(SDNode {
            kind: SDNodeKind::ConstantI32(100),
            operands: vec![],
            value_type: SdValueType::I32,
            node_id: 0,
            const_value: 100,
        });
        let c2 = dag.add_node(SDNode {
            kind: SDNodeKind::ConstantI32(30),
            operands: vec![],
            value_type: SdValueType::I32,
            node_id: 1,
            const_value: 30,
        });
        let sub = dag.add_node(SDNode {
            kind: SDNodeKind::Sub,
            operands: vec![c1, c2],
            value_type: SdValueType::I32,
            node_id: 2,
            const_value: 0,
        });
        dag.root_node = Some(sub);

        let mut isel = BpfISel::new();
        let result = isel.select(&dag);
        assert!(result.is_ok());
        let instrs = result.unwrap();
        let has_sub = instrs.iter().any(|i| i.opcode == BpfOpcode::SUB);
        assert!(has_sub);
    }

    #[test]
    fn test_select_mul_and_or_xor() {
        let ops = vec![
            (SDNodeKind::Mul, BpfOpcode::MUL),
            (SDNodeKind::And, BpfOpcode::AND),
            (SDNodeKind::Or, BpfOpcode::OR),
            (SDNodeKind::Xor, BpfOpcode::XOR),
        ];

        for (sd_op, bpf_op) in ops {
            let mut dag = SelectionDAG::new();
            let c1 = dag.add_node(SDNode {
                kind: SDNodeKind::ConstantI32(5),
                operands: vec![],
                value_type: SdValueType::I32,
                node_id: 0,
                const_value: 5,
            });
            let c2 = dag.add_node(SDNode {
                kind: SDNodeKind::ConstantI32(3),
                operands: vec![],
                value_type: SdValueType::I32,
                node_id: 1,
                const_value: 3,
            });
            let node = dag.add_node(SDNode {
                kind: sd_op,
                operands: vec![c1, c2],
                value_type: SdValueType::I32,
                node_id: 2,
                const_value: 0,
            });
            dag.root_node = Some(node);

            let mut isel = BpfISel::new();
            let result = isel.select(&dag);
            assert!(result.is_ok(), "Failed for op {:?}", sd_op);
            let instrs = result.unwrap();
            let found = instrs.iter().any(|i| i.opcode == bpf_op);
            assert!(found, "Expected {:?} for {:?}", bpf_op, sd_op);
        }
    }

    #[test]
    fn test_select_icmp_eq() {
        let mut dag = SelectionDAG::new();
        let c1 = dag.add_node(SDNode {
            kind: SDNodeKind::ConstantI32(42),
            operands: vec![],
            value_type: SdValueType::I32,
            node_id: 0,
            const_value: 42,
        });
        let c2 = dag.add_node(SDNode {
            kind: SDNodeKind::ConstantI32(42),
            operands: vec![],
            value_type: SdValueType::I32,
            node_id: 1,
            const_value: 42,
        });
        let cmp = dag.add_node(SDNode {
            kind: SDNodeKind::IcmpEq,
            operands: vec![c1, c2],
            value_type: SdValueType::I1,
            node_id: 2,
            const_value: 0,
        });
        dag.root_node = Some(cmp);

        let mut isel = BpfISel::new();
        let result = isel.select(&dag);
        assert!(result.is_ok());
        let instrs = result.unwrap();
        let has_jeq = instrs.iter().any(|i| i.opcode == BpfOpcode::JEQ);
        assert!(has_jeq, "Expected JEQ instruction");
    }

    #[test]
    fn test_select_ret() {
        let mut dag = SelectionDAG::new();
        let r = dag.add_node(SDNode {
            kind: SDNodeKind::Ret,
            operands: vec![],
            value_type: SdValueType::I1,
            node_id: 0,
            const_value: 0,
        });
        dag.root_node = Some(r);

        let mut isel = BpfISel::new();
        let result = isel.select(&dag);
        assert!(result.is_ok());
        let instrs = result.unwrap();
        let has_exit = instrs.iter().any(|i| i.opcode == BpfOpcode::EXIT);
        assert!(has_exit, "Expected EXIT instruction");
    }

    #[test]
    fn test_select_call() {
        let mut dag = SelectionDAG::new();
        let call_node = dag.add_node(SDNode {
            kind: SDNodeKind::Call,
            operands: vec![],
            value_type: SdValueType::I64,
            node_id: 0,
            const_value: 7,
        });
        dag.root_node = Some(call_node);

        let mut isel = BpfISel::new();
        let result = isel.select(&dag);
        assert!(result.is_ok());
        let instrs = result.unwrap();
        let has_call = instrs.iter().any(|i| i.opcode == BpfOpcode::CALL);
        assert!(has_call, "Expected CALL instruction");
    }

    #[test]
    fn test_select_frame_index() {
        let mut dag = SelectionDAG::new();
        let fi = dag.add_node(SDNode {
            kind: SDNodeKind::FrameIndex,
            operands: vec![],
            value_type: SdValueType::I64,
            node_id: 0,
            const_value: -16, // offset
        });
        dag.root_node = Some(fi);

        let mut isel = BpfISel::new();
        let result = isel.select(&dag);
        assert!(result.is_ok());
        let instrs = result.unwrap();
        let has_add = instrs
            .iter()
            .any(|i| i.opcode == BpfOpcode::ADD && i.src_reg == BPF_REG_10);
        assert!(has_add, "Expected ADD with r10 for FrameIndex");
    }

    #[test]
    fn test_select_load_store() {
        let mut dag = SelectionDAG::new();
        let addr = dag.add_node(SDNode {
            kind: SDNodeKind::FrameIndex,
            operands: vec![],
            value_type: SdValueType::I64,
            node_id: 0,
            const_value: -8,
        });
        let load = dag.add_node(SDNode {
            kind: SDNodeKind::Load,
            operands: vec![addr],
            value_type: SdValueType::I64,
            node_id: 1,
            const_value: 0,
        });
        dag.root_node = Some(load);

        let mut isel = BpfISel::new();
        let result = isel.select(&dag);
        assert!(result.is_ok());
        let instrs = result.unwrap();
        let has_ldx = instrs.iter().any(|i| i.opcode == BpfOpcode::LDX_DW);
        assert!(has_ldx, "Expected LDX_DW for load");
    }

    #[test]
    fn test_select_shift_ops() {
        let shifts = vec![
            (SDNodeKind::Shl, BpfOpcode::LSH),
            (SDNodeKind::Sra, BpfOpcode::ARSH),
            (SDNodeKind::Srl, BpfOpcode::RSH),
        ];

        for (sd_op, bpf_op) in shifts {
            let mut dag = SelectionDAG::new();
            let c1 = dag.add_node(SDNode {
                kind: SDNodeKind::ConstantI32(16),
                operands: vec![],
                value_type: SdValueType::I32,
                node_id: 0,
                const_value: 16,
            });
            let c2 = dag.add_node(SDNode {
                kind: SDNodeKind::ConstantI32(2),
                operands: vec![],
                value_type: SdValueType::I32,
                node_id: 1,
                const_value: 2,
            });
            let node = dag.add_node(SDNode {
                kind: sd_op,
                operands: vec![c1, c2],
                value_type: SdValueType::I32,
                node_id: 2,
                const_value: 0,
            });
            dag.root_node = Some(node);

            let mut isel = BpfISel::new();
            let result = isel.select(&dag);
            assert!(result.is_ok(), "Shift {:?} failed", sd_op);
        }
    }

    #[test]
    fn test_register_tracker() {
        let mut rt = RegTracker::new();
        let v0 = rt.new_vreg();
        let v1 = rt.new_vreg();
        assert_eq!(v0, 0);
        assert_eq!(v1, 1);

        let free = rt.find_free_reg();
        assert!(free.is_some());

        rt.assign_phys(v0, 3);
        assert_eq!(rt.get_phys(v0), Some(3));
        assert_eq!(rt.get_phys(v1), None);

        let slot = rt.spill(v1);
        assert!(slot < 0);
    }

    #[test]
    fn test_register_tracker_exhaust() {
        let mut rt = RegTracker::new();
        // Allocate all 10 GPRs
        for i in 0..10 {
            let v = rt.new_vreg();
            rt.assign_phys(v, i);
        }
        assert!(rt.find_free_reg().is_none());
    }

    #[test]
    fn test_selection_dag_empty() {
        let dag = SelectionDAG::new();
        assert_eq!(dag.nodes.len(), 0);
        assert!(dag.root_node.is_none());
    }

    #[test]
    fn test_multiple_icmp_kinds() {
        let kinds = vec![
            (SDNodeKind::IcmpNe, BpfOpcode::JNE),
            (SDNodeKind::IcmpSgt, BpfOpcode::JSGT),
            (SDNodeKind::IcmpSge, BpfOpcode::JSGE),
            (SDNodeKind::IcmpUgt, BpfOpcode::JGT),
            (SDNodeKind::IcmpUge, BpfOpcode::JGE),
        ];

        for (sd_op, bpf_op) in kinds {
            let mut dag = SelectionDAG::new();
            let c1 = dag.add_node(SDNode {
                kind: SDNodeKind::ConstantI32(10),
                operands: vec![],
                value_type: SdValueType::I32,
                node_id: 0,
                const_value: 10,
            });
            let c2 = dag.add_node(SDNode {
                kind: SDNodeKind::ConstantI32(20),
                operands: vec![],
                value_type: SdValueType::I32,
                node_id: 1,
                const_value: 20,
            });
            let node = dag.add_node(SDNode {
                kind: sd_op,
                operands: vec![c1, c2],
                value_type: SdValueType::I1,
                node_id: 2,
                const_value: 0,
            });
            dag.root_node = Some(node);

            let mut isel = BpfISel::new();
            let result = isel.select(&dag);
            assert!(result.is_ok(), "ICmp {:?} failed", sd_op);
        }
    }

    #[test]
    fn test_constant_i64_uses_lddw() {
        let mut dag = SelectionDAG::new();
        let c = dag.add_node(SDNode {
            kind: SDNodeKind::ConstantI64(0x7EAD_BEEF_CAFE_BABE),
            operands: vec![],
            value_type: SdValueType::I64,
            node_id: 0,
            const_value: 0x7EAD_BEEF_CAFE_BABE,
        });
        dag.root_node = Some(c);

        let mut isel = BpfISel::new();
        let result = isel.select(&dag);
        assert!(result.is_ok());
        let instrs = result.unwrap();
        let has_lddw = instrs.iter().any(|i| i.opcode == BpfOpcode::LD_DW);
        assert!(has_lddw, "Expected LD_DW for i64 constant");
    }
}