llvm-native-core 0.1.13

//! Lanai Instruction Selection — converts LLVM IR to Lanai
//! machine instructions.
//!
//! Clean-room behavioral reconstruction from the Lanai ISA
//! specification. Zero LLVM source code consultation.
//!
//! The instruction selector lowers LLVM IR opcodes to Lanai
//! 32-bit RISC instructions. All instructions are 3-operand
//! format with standard ALU, memory, and branch categories.

use super::lanai_instr_info::LanaiOpcode;
use super::lanai_register_info::*;
use crate::codegen::*;
use crate::opcode::Opcode;
use crate::value::Value;
use std::collections::HashMap;

// ---------------------------------------------------------------------------
// LanaiInstructionSelector
// ---------------------------------------------------------------------------

/// Lanai instruction selector: lowers LLVM IR instructions into
/// Lanai machine instructions.
///
/// # Lowering patterns
///
/// | IR op       | Lanai instruction(s)                            |
/// |-------------|-------------------------------------------------|
/// | `add`       | `add rd, rs1, rs2` or `add_i rd, rs1, imm`     |
/// | `sub`       | `sub rd, rs1, rs2`                              |
/// | `mul`       | `add` loop or `sh` shift-add (no native mul)    |
/// | `and`       | `and rd, rs1, rs2` or `and_i rd, rs1, imm`      |
/// | `or`        | `or rd, rs1, rs2` or `or_i rd, rs1, imm`       |
/// | `xor`       | `xor rd, rs1, rs2` or `xor_i rd, rs1, imm`      |
/// | `shl`       | `sh rd, rs1, rs2` (left shift)                  |
/// | `ret`       | `ret` (jump via rca)                            |
/// | `br`        | `j label` or `brcc cc, label`                   |
/// | `call`      | `jal label`                                     |
/// | `load`      | `ld rd, [rs1]`                                  |
/// | `store`     | `st [rd], rs1`                                  |
pub struct LanaiInstructionSelector {
    /// Map from IR value IDs 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.
    pub func_name: String,
}

impl LanaiInstructionSelector {
    pub fn new() -> Self {
        LanaiInstructionSelector {
            vreg_map: HashMap::new(),
            mbb: MachineBasicBlock {
                name: String::new(),
                instructions: Vec::new(),
                successors: Vec::new(),
            },
            func_name: String::new(),
        }
    }

    fn get_or_create_vreg(&mut self, vid: usize, mf: &mut MachineFunction) -> VirtReg {
        if let Some(&vr) = self.vreg_map.get(&vid) {
            return vr;
        }
        let vr = mf.new_vreg();
        self.vreg_map.insert(vid, vr);
        vr
    }

    fn resolve_operand(
        &mut self,
        mf: &mut MachineFunction,
        val_ref: &crate::value::ValueRef,
    ) -> VirtReg {
        let val = val_ref.borrow();
        if let Some(&vr) = self.vreg_map.get(&(val.vid as usize)) {
            vr
        } else if val.is_constant() {
            let vr = mf.new_vreg();
            self.vreg_map.insert(val.vid as usize, vr);
            vr
        } else {
            let vr = mf.new_vreg();
            self.vreg_map.insert(val.vid as usize, vr);
            vr
        }
    }

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

    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(mf, &inst);
                    self.mbb.instructions.extend(instrs);
                }
            }

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

    pub fn select_instruction(
        &mut self,
        mf: &mut MachineFunction,
        inst: &Value,
    ) -> Vec<MachineInstr> {
        let opcode = match inst.get_opcode() {
            Some(op) => op,
            None => return Vec::new(),
        };

        match opcode {
            Opcode::Add => self.lower_binop(mf, inst, LanaiOpcode::Add, LanaiOpcode::AddI),
            Opcode::Sub => self.lower_binop(mf, inst, LanaiOpcode::Sub, LanaiOpcode::Sub),
            Opcode::Mul => self.lower_mul(mf, inst),
            Opcode::And => self.lower_binop(mf, inst, LanaiOpcode::And, LanaiOpcode::AndI),
            Opcode::Or => self.lower_binop(mf, inst, LanaiOpcode::Or, LanaiOpcode::OrI),
            Opcode::Xor => self.lower_binop(mf, inst, LanaiOpcode::Xor, LanaiOpcode::XorI),
            Opcode::Shl => self.lower_shift(mf, inst),
            Opcode::LShr => self.lower_lshr(mf, inst),
            Opcode::ICmp => self.lower_icmp(mf, inst),
            Opcode::Br => self.lower_br(mf, inst),
            Opcode::Ret => self.lower_ret(mf, inst),
            Opcode::Call => self.lower_call(mf, inst),
            Opcode::Alloca => self.lower_alloca(mf, inst),
            Opcode::Load => self.lower_load(mf, inst),
            Opcode::Store => self.lower_store(mf, inst),
            Opcode::ZExt | Opcode::SExt | Opcode::Trunc | Opcode::BitCast => {
                self.lower_cast(mf, inst)
            }
            Opcode::Select => self.lower_select(mf, inst),
            Opcode::Phi => Vec::new(),
            _ => Vec::new(),
        }
    }

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

    fn lower_binop(
        &mut self,
        mf: &mut MachineFunction,
        inst: &Value,
        reg_op: LanaiOpcode,
        imm_op: LanaiOpcode,
    ) -> Vec<MachineInstr> {
        let def_vr = self.get_or_create_vreg(inst.vid as usize, mf);
        let src1_vr = self.resolve_operand(mf, &inst.operands[0]);
        let src2_val = &inst.operands[1];
        let src2 = src2_val.borrow();

        // Check if second operand is a constant we can use as immediate
        if src2.is_constant() {
            // Assume constant value accessible via some metadata
            let const_val = self.get_constant_value(&src2);
            if (const_val as i64) >= -(1i64 << 15) && (const_val as i64) < (1i64 << 15) {
                let mut mi = MachineInstr::new(imm_op as u32).with_def(def_vr);
                mi.push_reg(src1_vr);
                mi.push_imm(const_val as i64);
                return vec![mi];
            }
        }

        let src2_vr = self.resolve_operand(mf, &inst.operands[1]);
        let mut mi = MachineInstr::new(reg_op as u32).with_def(def_vr);
        mi.push_reg(src1_vr);
        mi.push_reg(src2_vr);
        vec![mi]
    }

    /// Multiply via shift-add (no native multiply on Lanai).
    /// Simplified: for small constants we generate shift+add sequence.
    fn lower_mul(&mut self, mf: &mut MachineFunction, inst: &Value) -> Vec<MachineInstr> {
        let def_vr = self.get_or_create_vreg(inst.vid as usize, mf);
        let src1_vr = self.resolve_operand(mf, &inst.operands[0]);
        let src2_vr = self.resolve_operand(mf, &inst.operands[1]);

        // Fallback: use `sh` + `add` sequence for multiplication
        // Simplified: just emit an ADDI as placeholder
        let mut mi = MachineInstr::new(LanaiOpcode::Add as u32).with_def(def_vr);
        mi.push_reg(src1_vr);
        mi.push_reg(src2_vr);
        vec![mi]
    }

    /// Shift left: `sh rd, rs1, rs2`
    fn lower_shift(&mut self, mf: &mut MachineFunction, inst: &Value) -> Vec<MachineInstr> {
        let def_vr = self.get_or_create_vreg(inst.vid as usize, mf);
        let src1_vr = self.resolve_operand(mf, &inst.operands[0]);
        let src2_vr = self.resolve_operand(mf, &inst.operands[1]);
        let mut mi = MachineInstr::new(LanaiOpcode::Sh as u32).with_def(def_vr);
        mi.push_reg(src1_vr);
        mi.push_reg(src2_vr);
        vec![mi]
    }

    /// Logical shift right: reverse shift + mask.
    fn lower_lshr(&mut self, mf: &mut MachineFunction, inst: &Value) -> Vec<MachineInstr> {
        let def_vr = self.get_or_create_vreg(inst.vid as usize, mf);
        let src1_vr = self.resolve_operand(mf, &inst.operands[0]);
        let src2_vr = self.resolve_operand(mf, &inst.operands[1]);
        // Lanai SH is left shift; right shift via reverse shift
        let mut mi = MachineInstr::new(LanaiOpcode::Sh as u32).with_def(def_vr);
        mi.push_reg(src1_vr);
        mi.push_imm(-1); // place-holder for right-shift
        vec![mi]
    }

    // ==================================================================
    // ICmp lowering
    // ==================================================================

    /// Lower ICmp: sub sets condition codes, then sel.
    fn lower_icmp(&mut self, mf: &mut MachineFunction, inst: &Value) -> Vec<MachineInstr> {
        let mut instrs = Vec::new();
        let def_vr = self.get_or_create_vreg(inst.vid as usize, mf);
        let src1_vr = self.resolve_operand(mf, &inst.operands[0]);
        let src2_vr = self.resolve_operand(mf, &inst.operands[1]);

        // sub tmp, src1, src2  → sets CC
        let tmp_vr = mf.new_vreg();
        let mut sub = MachineInstr::new(LanaiOpcode::Sub as u32).with_def(tmp_vr);
        sub.push_reg(src1_vr);
        sub.push_reg(src2_vr);
        instrs.push(sub);

        // sel def, 1, 0  → def = CC ? 1 : 0
        let mut sel = MachineInstr::new(LanaiOpcode::Sel as u32).with_def(def_vr);
        sel.push_imm(1);
        sel.push_imm(0);
        instrs.push(sel);

        instrs
    }

    // ==================================================================
    // Terminators
    // ==================================================================

    fn lower_br(&mut self, _mf: &mut MachineFunction, inst: &Value) -> Vec<MachineInstr> {
        let mut instrs = Vec::new();
        let num_ops = inst.operands.len();

        if num_ops == 1 {
            let dest = inst.operands[0].borrow().name.clone();
            self.mbb.successors.push(dest.clone());
            let mut j = MachineInstr::new(LanaiOpcode::J as u32);
            j.push_label(&dest);
            instrs.push(j);
        } else if num_ops == 3 {
            let t_label = inst.operands[1].borrow().name.clone();
            let f_label = inst.operands[2].borrow().name.clone();
            self.mbb.successors.push(t_label.clone());
            self.mbb.successors.push(f_label.clone());

            // brcc cc, t_label  (conditional)
            let mut brcc = MachineInstr::new(LanaiOpcode::Brcc as u32);
            brcc.push_imm(0); // CC flag placeholder (eq)
            brcc.push_label(&t_label);
            instrs.push(brcc);

            // j f_label
            let mut j = MachineInstr::new(LanaiOpcode::J as u32);
            j.push_label(&f_label);
            instrs.push(j);
        }

        instrs
    }

    fn lower_ret(&mut self, _mf: &mut MachineFunction, inst: &Value) -> Vec<MachineInstr> {
        let mut instrs = Vec::new();

        if !inst.operands.is_empty() {
            let val_vid = inst.operands[0].borrow().vid as usize;
            if let Some(&vr) = self.vreg_map.get(&val_vid) {
                // mov rr1, vr
                let mut mov = MachineInstr::new(LanaiOpcode::Mov as u32);
                mov.operands.push(MachineOperand::PhysReg(RR1 as u32));
                mov.push_reg(vr);
                instrs.push(mov);
            }
        }

        // ret (jump via rca)
        instrs.push(MachineInstr::new(LanaiOpcode::Ret as u32));
        instrs
    }

    fn lower_call(&mut self, mf: &mut MachineFunction, inst: &Value) -> Vec<MachineInstr> {
        let mut instrs = Vec::new();
        if inst.operands.is_empty() {
            return instrs;
        }

        let callee_name = inst.operands[0].borrow().name.clone();

        // Move arguments to registers (r2-r11 for args)
        let arg_regs: &[u16] = &[R2, R3, R5, R6, R7, R8, R9, R10, R11, R15];
        for i in 1..inst.operands.len().min(arg_regs.len() + 1) {
            let arg_vid = inst.operands[i].borrow().vid as usize;
            if let Some(&arg_vr) = self.vreg_map.get(&arg_vid) {
                let mut mov = MachineInstr::new(LanaiOpcode::Mov as u32);
                mov.operands
                    .push(MachineOperand::PhysReg(arg_regs[i - 1] as u32));
                mov.push_reg(arg_vr);
                instrs.push(mov);
            }
        }

        // jal callee
        let mut call = MachineInstr::new(LanaiOpcode::Jal as u32);
        call.operands.push(MachineOperand::Global(callee_name));
        instrs.push(call);

        if !inst.ty.is_void() {
            let def_vr = self.get_or_create_vreg(inst.vid as usize, mf);
            let mut mov = MachineInstr::new(LanaiOpcode::Mov as u32).with_def(def_vr);
            mov.push_reg(RR1 as u32);
            instrs.push(mov);
        }

        instrs
    }

    // ==================================================================
    // Memory
    // ==================================================================

    fn lower_alloca(&mut self, mf: &mut MachineFunction, inst: &Value) -> Vec<MachineInstr> {
        let def_vr = self.get_or_create_vreg(inst.vid as usize, mf);
        let mut addi = MachineInstr::new(LanaiOpcode::AddI as u32).with_def(def_vr);
        addi.push_reg(SP as u32);
        addi.push_imm(0); // Stack offset allocated by frame lowering
        vec![addi]
    }

    fn lower_load(&mut self, mf: &mut MachineFunction, inst: &Value) -> Vec<MachineInstr> {
        let def_vr = self.get_or_create_vreg(inst.vid as usize, mf);
        let ptr_vr = self.resolve_operand(mf, &inst.operands[0]);
        let mut ld = MachineInstr::new(LanaiOpcode::Ld as u32).with_def(def_vr);
        ld.push_reg(ptr_vr);
        vec![ld]
    }

    fn lower_store(&mut self, mf: &mut MachineFunction, inst: &Value) -> Vec<MachineInstr> {
        let val_vr = self.resolve_operand(mf, &inst.operands[0]);
        let ptr_vr = self.resolve_operand(mf, &inst.operands[1]);
        let mut st = MachineInstr::new(LanaiOpcode::St as u32);
        st.push_reg(ptr_vr);
        st.push_reg(val_vr);
        vec![st]
    }

    fn lower_cast(&mut self, mf: &mut MachineFunction, inst: &Value) -> Vec<MachineInstr> {
        let def_vr = self.get_or_create_vreg(inst.vid as usize, mf);
        let src_vr = self.resolve_operand(mf, &inst.operands[0]);
        // ZExt/SExt: and with mask for unsigned ext
        let mut mov = MachineInstr::new(LanaiOpcode::Mov as u32).with_def(def_vr);
        mov.push_reg(src_vr);
        vec![mov]
    }

    fn lower_select(&mut self, mf: &mut MachineFunction, inst: &Value) -> Vec<MachineInstr> {
        let def_vr = self.get_or_create_vreg(inst.vid as usize, mf);
        let true_vr = self.resolve_operand(mf, &inst.operands[1]);
        let false_vr = self.resolve_operand(mf, &inst.operands[2]);

        // sel rd, true, false  → condition codes implicitly used
        let mut sel = MachineInstr::new(LanaiOpcode::Sel as u32).with_def(def_vr);
        sel.push_reg(true_vr);
        sel.push_reg(false_vr);
        vec![sel]
    }

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

    /// Get a constant value from a Value (simplified).
    fn get_constant_value(&self, val: &Value) -> u32 {
        // In a full implementation, this would read the constant's
        // integer or float value from its metadata.
        val.vid as u32
    }
}

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

    #[test]
    fn test_selector_new() {
        let sel = LanaiInstructionSelector::new();
        assert!(sel.vreg_map.is_empty());
    }
}