rusty-minic/src/ir/types.rs

use std::fmt::Display;

use crate::ast::types::Type as AstType;
use crate::ast::types::BinaryOp as AstBinaryOp;
use crate::ir::err::IRError;

pub enum VariableOrIntLit {
    Var(Variable),
    IntLit(i32),
}

impl Display for VariableOrIntLit {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            VariableOrIntLit::Var(v) => write!(f, "{}", v),
            VariableOrIntLit::IntLit(i) => write!(f, "{}", i),
        }
    }
}
pub enum IRInstr {
    Declare(Variable),
    DefineFunc(Function, Vec<Variable>, Vec<IRInstr>),
    Entry,
    Binary(Variable, Variable, BinaryOp, Variable),
    Cmp(Variable, VariableOrIntLit, CmpOp, VariableOrIntLit),
    Unary(Variable, UnaryOp, Variable),
    Exit(Option<Variable>),
    FuncCall(Function, Vec<Variable>, Option<Variable>),
    Goto(usize),
    CondGoto(Variable, usize, usize), // condition, true label, false label
    Label(usize),
    Move(Variable, MoveRValue),
}

impl Display for IRInstr {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            IRInstr::Entry => write!(f, "entry"),
            IRInstr::Binary(dest, left, op, right) => write!(f, "{} = {} {},{}", dest, op, left, right),
            IRInstr::Cmp(dest, left, op, right) => write!(f, "{} = icmp {} {},{}", dest, op, left, right),
            IRInstr::Exit(v) => if let Some(v) = v { write!(f, "exit {}", v) } else { write!(f, "exit") },
            IRInstr::FuncCall(func, args, dest) => {
                if let Some(dest) = dest {
                    write!(f, "{} = call {}", dest, func.to_call_string(args))
                } else {
                    write!(f, "call {}", func.to_call_string(args))
                }
            }
            IRInstr::Move(dest, src) => write!(f, "{} = {}", dest, src),
            IRInstr::Declare(var) => write!(f, "declare {} {}", var.data_type, var),
            IRInstr::DefineFunc(func, args, body) => {
                let body_str = body.iter().map(|instr| format!("  {}", instr)).collect::<Vec<_>>().join("\n");
                write!(f, "define {} {{\n{}\n}}", func.to_decl_string(args), body_str)
            },
            IRInstr::Goto(label) => write!(f, "br label .L{}", label),
            IRInstr::CondGoto(cond, true_label, false_label) => write!(f, "bc {}, .L{}, .L{}", cond, true_label, false_label),
            IRInstr::Label(label) => write!(f, ".L{}:", label),
            IRInstr::Unary(dest, op, src) => write!(f, "{} = {} {}", dest, op, src),
        }
    }
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum IRType {
    I32,
    I1,
    Void,
}
impl IRType {
    pub fn size_in_bytes(&self) -> usize {
        match self {
            IRType::I32 => 4,
            IRType::I1 => 1,
            IRType::Void => 0,
        }
    }

    pub fn get_elevate_result(lhs: IRType, rhs: IRType) -> Option<IRType> {
        if lhs == rhs {
            Some(lhs)
        } else if (lhs == IRType::I32 && rhs == IRType::I1) || (lhs == IRType::I1 && rhs == IRType::I32) {
            Some(IRType::I32)
        } else {
            None
        }
    }
}
impl Display for IRType {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            IRType::I32 => write!(f, "i32"),
            IRType::Void => write!(f, "void"),
            IRType::I1 => write!(f, "i1"),
        }
    }
}
impl From<AstType> for IRType {
    fn from(ast_type: AstType) -> Self {
        match ast_type {
            AstType::Int => IRType::I32,
            AstType::Void => IRType::Void,
        }
    }
}
pub enum MoveRValue {
    Var(Variable),
    ConstInt(i32),
}
impl Display for MoveRValue {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            MoveRValue::Var(v) => write!(f, "{}", v),
            MoveRValue::ConstInt(i) => write!(f, "{}", i),
        }
    }
}
#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Debug, Hash)]
pub enum VariableType {
    Global,
    ParamTemp(usize),
    Local,
    Temp,
}
#[derive(Clone, Copy)]
pub struct Variable {
    // pub name: String,
    pub index: usize,
    pub var_type: VariableType,
    pub data_type: IRType,
}
impl PartialEq for Variable {
    fn eq(&self, other: &Self) -> bool {
        self.index == other.index && self.var_type == other.var_type
    }
}
impl Eq for Variable {}
impl PartialOrd for Variable {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        Some(match self.index.cmp(&other.index) {
            std::cmp::Ordering::Equal => self.var_type.cmp(&other.var_type),
            ord => ord,
        })
    }
}
impl Ord for Variable {
    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
        match self.index.cmp(&other.index) {
            std::cmp::Ordering::Equal => self.var_type.cmp(&other.var_type),
            ord => ord,
        }
    }
}
impl Display for Variable {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let prefix = match self.var_type {
            VariableType::Global => "@g",
            VariableType::Local => "%l",
            VariableType::Temp => "%t",
            VariableType::ParamTemp(_) => "%t",
        };
        write!(f, "{}{}", prefix, self.index)
    }
}
#[derive(Debug, Clone)]
pub struct Function {
    pub name: String,
    pub return_type: IRType,
    pub parameter_types: Vec<IRType>,
    // pub parameters: Vec<Variable>,
}
impl Function {
    pub fn to_call_string(&self, args: &Vec<Variable>) -> String {
        assert!(args.len() == self.parameter_types.len());
        let args_str = args.iter().zip(self.parameter_types.iter()).map(|(arg, param)| format!("{} {}", param, arg)).collect::<Vec<_>>().join(", ");
        format!("{} @{}({})", self.return_type, self.name, args_str)
    }

    pub fn to_decl_string(&self, args: &Vec<Variable>) -> String {
        let params_str = args.iter().zip(self.parameter_types.iter()).map(|(arg, param_type)| format!("{} {}", param_type, arg)).collect::<Vec<_>>().join(", ");
        format!("{} @{}({})", self.return_type, self.name, params_str)
    }
}
pub enum BinaryOp {
    Add,
    Sub,
    Mul,
    Div,
    Mod,
}
pub enum CmpOp {
    Le,
    Lt,
    Gt,
    Ge,
    Ne,
    Eq,
}
pub enum UnaryOp {
    Neg
}

impl Display for BinaryOp {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let op_str = match self {
            BinaryOp::Add => "add",
            BinaryOp::Sub => "sub",
            BinaryOp::Mul => "mul",
            BinaryOp::Div => "div",
            BinaryOp::Mod => "mod",
            // BinaryOp::Le => "cmp le",
            // BinaryOp::Lt => "cmp lt",
            // BinaryOp::Gt => "cmp gt",
            // BinaryOp::Ge => "cmp ge",
            // BinaryOp::Ne => "cmp ne",
            // BinaryOp::Eq => "cmp eq",
        };
        write!(f, "{}", op_str)
    }
}

impl Display for CmpOp {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let op_str = match self {
            CmpOp::Le => "le",
            CmpOp::Lt => "lt",
            CmpOp::Gt => "gt",
            CmpOp::Ge => "ge",
            CmpOp::Ne => "ne",
            CmpOp::Eq => "eq",
        };
        write!(f, "{}", op_str)
    }
}

impl Display for UnaryOp {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let op_str = match self {
            UnaryOp::Neg => "neg",
        };
        write!(f, "{}", op_str)
    }
}
impl From<AstBinaryOp> for BinaryOp {
    fn from(ast_op: AstBinaryOp) -> Self {
        match ast_op {
            AstBinaryOp::Add => BinaryOp::Add,
            AstBinaryOp::Sub => BinaryOp::Sub,
            AstBinaryOp::Mul => BinaryOp::Mul,
            AstBinaryOp::Div => BinaryOp::Div,
            AstBinaryOp::Mod => BinaryOp::Mod,
            // AstBinaryOp::Equal => BinaryOp::Eq,
            // AstBinaryOp::NotEqual => BinaryOp::Ne,
            // AstBinaryOp::Less => BinaryOp::Lt,
            // AstBinaryOp::LessEqual => BinaryOp::Le,
            // AstBinaryOp::Greater => BinaryOp::Gt,
            // AstBinaryOp::GreaterEqual => BinaryOp::Ge,
            _ => unreachable!(),
        }
    }
}
impl From<AstBinaryOp> for CmpOp {
    fn from(ast_op: AstBinaryOp) -> Self {
        match ast_op {
            AstBinaryOp::Equal => CmpOp::Eq,
            AstBinaryOp::NotEqual => CmpOp::Ne,
            AstBinaryOp::Less => CmpOp::Lt,
            AstBinaryOp::LessEqual => CmpOp::Le,
            AstBinaryOp::Greater => CmpOp::Gt,
            AstBinaryOp::GreaterEqual => CmpOp::Ge,
            _ => unreachable!(),
        }
    }
}