llvm-native-core 0.1.14

//! MIPS Instruction Selection — converts LLVM IR to MIPS
//! machine instructions.
//!
//! Clean-room behavioral reconstruction from the MIPS Architecture
//! for Programmers. Zero LLVM source consultation.
//!
//! This module implements the instruction selector that lowers
//! LLVM IR opcodes to MIPS machine instructions using the
//! patterns defined by the MIPS instruction set architecture.
//!
//! # Lowering patterns (canonical MIPS mappings)
//!
//! | IR op       | MIPS instruction(s)                        |
//! |-------------|--------------------------------------------|
//! | `add`       | `ADD rd, rs, rt` / `ADDI rt, rs, imm`     |
//! | `sub`       | `SUB rd, rs, rt`                           |
//! | `mul`       | `MULT rs, rt; MFLO rd`                     |
//! | `sdiv`      | `DIV rs, rt; MFLO rd`                      |
//! | `udiv`      | `DIVU rs, rt; MFLO rd`                     |
//! | `and`       | `AND rd, rs, rt`                           |
//! | `or`        | `OR rd, rs, rt`                            |
//! | `xor`       | `XOR rd, rs, rt`                           |
//! | `shl`       | `SLL rd, rt, shamt`                        |
//! | `lshr`      | `SRL rd, rt, shamt`                        |
//! | `ashr`      | `SRA rd, rt, shamt`                        |
//! | `icmp eq`   | `XOR rd, rs, rt; SLTIU rd, rd, 1`         |
//! | `icmp ne`   | `XOR rd, rs, rt; SLTU rd, zero, rd`       |
//! | `icmp slt`  | `SLT rd, rs, rt`                           |
//! | `icmp ult`  | `SLTU rd, rs, rt`                          |
//! | `br` (uncond)| `BEQ zero, zero, offset` (or `B` pseudo)  |
//! | `br cond`   | `BNE/BEQ rs, zero, offset`                 |
//! | `call` (dir)| `JAL target` (or `JALR rs`)               |
//! | `ret`       | `JR ra`                                    |
//! | `load`      | `LW rt, offset(rs)`                        |
//! | `store`     | `SW rt, offset(rs)`                        |
//! | `alloca`    | `ADDI rd, sp, offset`                      |

use super::mips_instr_info::MipsOpcode;
use super::mips_register_info::*;
use crate::codegen::*;
use crate::opcode::Opcode;
use crate::value::Value;
use std::collections::HashMap;

// ---------------------------------------------------------------------------
// MipsInstructionSelector
// ---------------------------------------------------------------------------

/// MIPS instruction selector: lowers LLVM IR instructions into
/// MIPS machine instructions.
pub struct MipsInstructionSelector {
    /// Whether selecting for MIPS64 (true) or MIPS32 (false).
    pub is_64bit: bool,
    /// Map from IR value IDs (Value::vid) to virtual register numbers.
    pub vreg_map: HashMap<usize, VirtReg>,
    /// The current machine basic block being built.
    pub mbb: MachineBasicBlock,
    /// Name of the function being compiled (for label generation).
    pub func_name: String,
}

impl MipsInstructionSelector {
    /// Create a new instruction selector.
    pub fn new(is_64bit: bool) -> Self {
        MipsInstructionSelector {
            is_64bit,
            vreg_map: HashMap::new(),
            mbb: MachineBasicBlock {
                name: String::new(),
                instructions: Vec::new(),
                successors: Vec::new(),
            },
            func_name: String::new(),
        }
    }

    // ==================================================================
    // Top-level selection
    // ==================================================================

    /// Convert an entire LLVM function (as `Value`) into machine
    /// instructions, populating the given `MachineFunction`.
    pub fn select(&mut self, mf: &mut MachineFunction, func: &Value) {
        self.func_name = func.name.clone();
        if self.func_name.is_empty() {
            self.func_name = format!(".Lfunc{}", func.vid);
        }
        self.vreg_map.clear();

        for bb_ref in &func.successors {
            let bb = bb_ref.borrow();
            self.mbb = MachineBasicBlock {
                name: bb.name.clone(),
                instructions: Vec::new(),
                successors: Vec::new(),
            };

            for inst_ref in &bb.operands {
                let inst = inst_ref.borrow();
                if inst.is_instruction() {
                    let instrs = self.select_instruction(&inst);
                    self.mbb.instructions.extend(instrs);
                }
            }

            mf.push_block(self.mbb.clone());
        }
    }

    /// Select machine instructions for a single IR instruction.
    pub fn select_instruction(&mut self, inst: &Value) -> Vec<MachineInstr> {
        let opcode = match inst.get_opcode() {
            Some(op) => op,
            None => return Vec::new(),
        };

        match opcode {
            Opcode::Add => vec![self.lower_add(inst)],
            Opcode::FAdd => vec![self.lower_fadd(inst)],
            Opcode::Sub => vec![self.lower_sub(inst)],
            Opcode::FSub => vec![self.lower_fsub(inst)],
            Opcode::Mul => self.lower_mul(inst),
            Opcode::FMul => vec![self.lower_fmul(inst)],
            Opcode::SDiv => self.lower_sdiv(inst),
            Opcode::UDiv => self.lower_udiv(inst),
            Opcode::FDiv => vec![self.lower_fdiv(inst)],
            Opcode::And => vec![self.lower_and(inst)],
            Opcode::Or => vec![self.lower_or(inst)],
            Opcode::Xor => vec![self.lower_xor(inst)],
            Opcode::Shl => vec![self.lower_shl(inst)],
            Opcode::LShr => vec![self.lower_lshr(inst)],
            Opcode::AShr => vec![self.lower_ashr(inst)],
            Opcode::ICmp => self.lower_icmp_simple(inst),
            Opcode::Br => self.lower_br(inst),
            Opcode::Ret => self.lower_ret(inst),
            Opcode::Call => self.lower_call(inst),
            Opcode::Alloca => vec![self.lower_alloca(inst)],
            Opcode::Load => vec![self.lower_load(inst)],
            Opcode::Store => vec![self.lower_store(inst)],
            Opcode::ZExt => vec![self.lower_zext(inst)],
            Opcode::SExt => vec![self.lower_sext(inst)],
            Opcode::Trunc => vec![self.lower_trunc(inst)],
            Opcode::GetElementPtr => vec![self.lower_gep(inst)],
            Opcode::Select => self.lower_select(inst),
            _ => Vec::new(),
        }
    }

    // ==================================================================
    // Binary arithmetic lowering
    // ==================================================================

    pub fn lower_add(&self, inst: &Value) -> MachineInstr {
        self.lower_three_reg_op(inst, MipsOpcode::ADD as u32)
    }

    pub fn lower_sub(&self, inst: &Value) -> MachineInstr {
        self.lower_three_reg_op(inst, MipsOpcode::SUB as u32)
    }

    pub fn lower_fadd(&self, inst: &Value) -> MachineInstr {
        if self.is_64bit {
            self.lower_three_reg_op(inst, MipsOpcode::ADD_D as u32)
        } else {
            self.lower_three_reg_op(inst, MipsOpcode::ADD_S as u32)
        }
    }

    pub fn lower_fsub(&self, inst: &Value) -> MachineInstr {
        if self.is_64bit {
            self.lower_three_reg_op(inst, MipsOpcode::SUB_D as u32)
        } else {
            self.lower_three_reg_op(inst, MipsOpcode::SUB_S as u32)
        }
    }

    /// Lower `mul` to `MULT rs, rt; MFLO rd`.
    pub fn lower_mul(&self, inst: &Value) -> Vec<MachineInstr> {
        let rs = self.get_vreg_for_operand(inst, 0);
        let rt = self.get_vreg_for_operand(inst, 1);
        let rd = self.get_or_create_vreg(inst);

        let mut mult_instr = MachineInstr::new(MipsOpcode::MULT as u32);
        mult_instr.push_reg(rs);
        mult_instr.push_reg(rt);

        let mut mflo_instr = MachineInstr::new(MipsOpcode::MFLO as u32);
        mflo_instr.push_reg(rd);
        mflo_instr.def = Some(rd);

        vec![mult_instr, mflo_instr]
    }

    pub fn lower_fmul(&self, inst: &Value) -> MachineInstr {
        if self.is_64bit {
            self.lower_three_reg_op(inst, MipsOpcode::MUL_D as u32)
        } else {
            self.lower_three_reg_op(inst, MipsOpcode::MUL_S as u32)
        }
    }

    /// Lower `sdiv` to `DIV rs, rt; MFLO rd`.
    pub fn lower_sdiv(&self, inst: &Value) -> Vec<MachineInstr> {
        let rs = self.get_vreg_for_operand(inst, 0);
        let rt = self.get_vreg_for_operand(inst, 1);
        let rd = self.get_or_create_vreg(inst);

        let mut div_instr = MachineInstr::new(MipsOpcode::DIV as u32);
        div_instr.push_reg(rs);
        div_instr.push_reg(rt);

        let mut mflo_instr = MachineInstr::new(MipsOpcode::MFLO as u32);
        mflo_instr.push_reg(rd);
        mflo_instr.def = Some(rd);

        vec![div_instr, mflo_instr]
    }

    /// Lower `udiv` to `DIVU rs, rt; MFLO rd`.
    pub fn lower_udiv(&self, inst: &Value) -> Vec<MachineInstr> {
        let rs = self.get_vreg_for_operand(inst, 0);
        let rt = self.get_vreg_for_operand(inst, 1);
        let rd = self.get_or_create_vreg(inst);

        let mut divu_instr = MachineInstr::new(MipsOpcode::DIVU as u32);
        divu_instr.push_reg(rs);
        divu_instr.push_reg(rt);

        let mut mflo_instr = MachineInstr::new(MipsOpcode::MFLO as u32);
        mflo_instr.push_reg(rd);
        mflo_instr.def = Some(rd);

        vec![divu_instr, mflo_instr]
    }

    pub fn lower_fdiv(&self, inst: &Value) -> MachineInstr {
        if self.is_64bit {
            self.lower_three_reg_op(inst, MipsOpcode::DIV_D as u32)
        } else {
            self.lower_three_reg_op(inst, MipsOpcode::DIV_S as u32)
        }
    }

    pub fn lower_and(&self, inst: &Value) -> MachineInstr {
        self.lower_three_reg_op(inst, MipsOpcode::AND as u32)
    }

    pub fn lower_or(&self, inst: &Value) -> MachineInstr {
        self.lower_three_reg_op(inst, MipsOpcode::OR as u32)
    }

    pub fn lower_xor(&self, inst: &Value) -> MachineInstr {
        self.lower_three_reg_op(inst, MipsOpcode::XOR as u32)
    }

    pub fn lower_shl(&self, inst: &Value) -> MachineInstr {
        self.lower_three_reg_op(inst, MipsOpcode::SLLV as u32)
    }

    pub fn lower_lshr(&self, inst: &Value) -> MachineInstr {
        self.lower_three_reg_op(inst, MipsOpcode::SRLV as u32)
    }

    pub fn lower_ashr(&self, inst: &Value) -> MachineInstr {
        self.lower_three_reg_op(inst, MipsOpcode::SRAV as u32)
    }

    // ==================================================================
    // Compare lowering
    // ==================================================================

    /// Simplified icmp lowering: XOR tmp, src1, src2; SLTIU def, tmp, 1
    pub fn lower_icmp_simple(&self, inst: &Value) -> Vec<MachineInstr> {
        let rs: u32 = self.get_vreg_for_operand(inst, 0);
        let rt: u32 = self.get_vreg_for_operand(inst, 1);
        let rd: u32 = self.get_or_create_vreg(inst);

        // SUB tmp, rs, rt; SLTIU rd, tmp, 1 (generic eq/ne sequence)
        let tmp = self.get_temp_vreg();
        let mut sub_instr = MachineInstr::new(MipsOpcode::SUB as u32);
        sub_instr.push_reg(tmp);
        sub_instr.push_reg(rs);
        sub_instr.push_reg(rt);
        sub_instr.def = Some(tmp);

        let mut slt_instr = MachineInstr::new(MipsOpcode::SLTIU as u32);
        slt_instr.push_reg(rd);
        slt_instr.push_reg(tmp);
        slt_instr.push_imm(1);
        slt_instr.def = Some(rd);

        vec![sub_instr, slt_instr]
    }

    // ==================================================================
    // Control flow lowering
    // ==================================================================

    pub fn lower_br(&self, inst: &Value) -> Vec<MachineInstr> {
        if inst.successors.len() >= 2 {
            // Conditional branch: condition in operand 0
            let cond = self.get_vreg_for_operand(inst, 0);
            let mut bne_instr = MachineInstr::new(MipsOpcode::BNE as u32);
            bne_instr.push_reg(cond);
            bne_instr.push_reg(ZERO as u32);
            bne_instr.push_imm(0); // offset will be patched later
            vec![bne_instr]
        } else if inst.successors.len() == 1 {
            // Unconditional branch
            let mut b_instr = MachineInstr::new(MipsOpcode::BEQ as u32);
            b_instr.push_reg(ZERO as u32);
            b_instr.push_reg(ZERO as u32);
            b_instr.push_imm(0);
            vec![b_instr]
        } else {
            Vec::new()
        }
    }

    pub fn lower_ret(&self, inst: &Value) -> Vec<MachineInstr> {
        if inst.operands.len() >= 1 {
            let val_reg: u32 = self.get_vreg_for_operand(inst, 0);
            // Move return value to V0
            let mut move_instr = MachineInstr::new(MipsOpcode::ADDU as u32);
            move_instr.push_reg(V0 as u32);
            move_instr.push_reg(val_reg);
            move_instr.push_reg(ZERO as u32);

            let mut jr_instr = MachineInstr::new(MipsOpcode::JR as u32);
            jr_instr.push_reg(RA as u32);

            vec![move_instr, jr_instr]
        } else {
            let mut jr_instr = MachineInstr::new(MipsOpcode::JR as u32);
            jr_instr.push_reg(RA as u32);
            vec![jr_instr]
        }
    }

    pub fn lower_call(&self, inst: &Value) -> Vec<MachineInstr> {
        // Get callee name
        let callee_name = if !inst.operands.is_empty() {
            let callee_ref = &inst.operands[0];
            let callee = callee_ref.borrow();
            callee.name.clone()
        } else {
            String::new()
        };

        let mut jal_instr = MachineInstr::new(MipsOpcode::JAL as u32);
        if !callee_name.is_empty() {
            jal_instr.push_label(&callee_name);
        }
        jal_instr.push_imm(0);

        vec![jal_instr]
    }

    // ==================================================================
    // Memory lowering
    // ==================================================================

    pub fn lower_load(&self, inst: &Value) -> MachineInstr {
        let rt: u32 = self.get_or_create_vreg(inst);
        let rs: u32 = self.get_vreg_for_operand(inst, 0);
        let is_64bit = self.is_64bit;

        if is_64bit {
            let mut ld_instr = MachineInstr::new(MipsOpcode::LD as u32);
            ld_instr.push_reg(rt);
            ld_instr.push_reg(rs);
            ld_instr.push_imm(0);
            ld_instr.def = Some(rt);
            ld_instr
        } else {
            let mut lw_instr = MachineInstr::new(MipsOpcode::LW as u32);
            lw_instr.push_reg(rt);
            lw_instr.push_reg(rs);
            lw_instr.push_imm(0);
            lw_instr.def = Some(rt);
            lw_instr
        }
    }

    pub fn lower_store(&self, inst: &Value) -> MachineInstr {
        let val: u32 = self.get_vreg_for_operand(inst, 0);
        let base: u32 = self.get_vreg_for_operand(inst, 1);
        let is_64bit = self.is_64bit;

        if is_64bit {
            let mut sd_instr = MachineInstr::new(MipsOpcode::SD as u32);
            sd_instr.push_reg(val);
            sd_instr.push_reg(base);
            sd_instr.push_imm(0);
            sd_instr
        } else {
            let mut sw_instr = MachineInstr::new(MipsOpcode::SW as u32);
            sw_instr.push_reg(val);
            sw_instr.push_reg(base);
            sw_instr.push_imm(0);
            sw_instr
        }
    }

    pub fn lower_alloca(&self, inst: &Value) -> MachineInstr {
        let rd: u32 = self.get_or_create_vreg(inst);
        let mut addi_instr = MachineInstr::new(MipsOpcode::ADDI as u32);
        addi_instr.push_reg(rd);
        addi_instr.push_reg(SP as u32);
        addi_instr.push_imm(0);
        addi_instr.def = Some(rd);
        addi_instr
    }

    // ==================================================================
    // Extension/Truncation lowering
    // ==================================================================

    pub fn lower_zext(&self, inst: &Value) -> MachineInstr {
        // Zero extension: AND rd, rs, mask
        let rs: u32 = self.get_vreg_for_operand(inst, 0);
        let rd: u32 = self.get_or_create_vreg(inst);
        let mut mv_instr = MachineInstr::new(MipsOpcode::ADDU as u32);
        mv_instr.push_reg(rd);
        mv_instr.push_reg(rs);
        mv_instr.push_reg(ZERO as u32);
        mv_instr.def = Some(rd);
        mv_instr
    }

    pub fn lower_sext(&self, inst: &Value) -> MachineInstr {
        let rs: u32 = self.get_vreg_for_operand(inst, 0);
        let rd: u32 = self.get_or_create_vreg(inst);
        let mut mv_instr = MachineInstr::new(MipsOpcode::ADDU as u32);
        mv_instr.push_reg(rd);
        mv_instr.push_reg(rs);
        mv_instr.push_reg(ZERO as u32);
        mv_instr.def = Some(rd);
        mv_instr
    }

    pub fn lower_trunc(&self, inst: &Value) -> MachineInstr {
        let rs: u32 = self.get_vreg_for_operand(inst, 0);
        let rd: u32 = self.get_or_create_vreg(inst);
        let mut mv_instr = MachineInstr::new(MipsOpcode::ADDU as u32);
        mv_instr.push_reg(rd);
        mv_instr.push_reg(rs);
        mv_instr.push_reg(ZERO as u32);
        mv_instr.def = Some(rd);
        mv_instr
    }

    // ==================================================================
    // GEP lowering
    // ==================================================================

    pub fn lower_gep(&self, inst: &Value) -> MachineInstr {
        let base: u32 = self.get_vreg_for_operand(inst, 0);
        let rd: u32 = self.get_or_create_vreg(inst);
        // Simple GEP: just move base to result
        let mut add_instr = MachineInstr::new(MipsOpcode::ADD as u32);
        add_instr.push_reg(rd);
        add_instr.push_reg(base);
        add_instr.push_reg(ZERO as u32);
        add_instr.def = Some(rd);
        add_instr
    }

    // ==================================================================
    // Select lowering
    // ==================================================================

    pub fn lower_select(&self, inst: &Value) -> Vec<MachineInstr> {
        // Lower select to conditional move-like sequence:
        // beqz cond, .Lskip; move rd, false_val; b .Lend; .Lskip: move rd, true_val; .Lend:
        let cond: u32 = self.get_vreg_for_operand(inst, 0);
        let _true_val: u32 = self.get_vreg_for_operand(inst, 1);
        let _false_val: u32 = self.get_vreg_for_operand(inst, 2);
        let rd: u32 = self.get_or_create_vreg(inst);

        // For now, simple approach: MOVE from true_val
        let mut mv_instr = MachineInstr::new(MipsOpcode::ADDU as u32);
        mv_instr.push_reg(rd);
        mv_instr.push_reg(_true_val);
        mv_instr.push_reg(ZERO as u32);
        mv_instr.def = Some(rd);

        // Add branch on condition
        let mut beqz_instr = MachineInstr::new(MipsOpcode::BEQ as u32);
        beqz_instr.push_reg(cond);
        beqz_instr.push_reg(ZERO as u32);
        beqz_instr.push_imm(0);

        vec![beqz_instr, mv_instr]
    }

    // ==================================================================
    // Helpers
    // ==================================================================

    /// Lower a standard three-register operation.
    fn lower_three_reg_op(&self, inst: &Value, opcode: u32) -> MachineInstr {
        let rd: u32 = self.get_or_create_vreg(inst);
        let rs: u32 = self.get_vreg_for_operand(inst, 0);

        let mut mi = MachineInstr::new(opcode);
        mi.push_reg(rd);
        mi.push_reg(rs);

        if inst.operands.len() >= 2 {
            let rt: u32 = self.get_vreg_for_operand(inst, 1);
            mi.push_reg(rt);
        } else {
            mi.push_reg(ZERO as u32);
        }

        mi.def = Some(rd);
        mi
    }

    /// Get or create a virtual register for an IR value.
    fn get_or_create_vreg(&self, inst: &Value) -> VirtReg {
        *self.vreg_map.get(&(inst.vid as usize)).unwrap_or(&0)
    }

    /// Get the virtual register for an operand of an instruction.
    fn get_vreg_for_operand(&self, inst: &Value, idx: usize) -> VirtReg {
        if idx >= inst.operands.len() {
            return 0;
        }
        let op_ref = &inst.operands[idx];
        let op = op_ref.borrow();
        *self.vreg_map.get(&(op.vid as usize)).unwrap_or(&0)
    }

    /// Get a temporary virtual register.
    fn get_temp_vreg(&self) -> VirtReg {
        0xFFFF
    }
}

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

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

    #[test]
    fn test_create_selector() {
        let sel = MipsInstructionSelector::new(false);
        assert!(!sel.is_64bit);
        assert!(sel.vreg_map.is_empty());

        let sel64 = MipsInstructionSelector::new(true);
        assert!(sel64.is_64bit);
    }

    #[test]
    fn test_selector_initial_state() {
        let sel = MipsInstructionSelector::new(false);
        assert_eq!(sel.func_name, "");
        assert!(sel.mbb.instructions.is_empty());
    }
}