rusty-minic/src/ast/types.rs

use crate::{diagnostic::span::Span, frontend::types::{TokenValue, TypeIdent}};
use std::fmt;

pub struct CompileUnit {
    pub global_decls: Vec<GlobalDeclStmt>,
}
pub enum GlobalDeclStmt {
    VarDecl(VarDeclStmt),
    FuncDecl(FuncDeclStmt),
}

pub struct VarDeclStmt {
    pub values: Vec<VarDeclStmtValue>,
    pub span: Span,
}
pub struct VarDeclStmtValue {
    pub name: String,
    pub var_type: Type,
    pub span: Span,

}

pub struct FuncDeclStmt {
    pub name: String,
    pub return_type: Type,
    pub params: Vec<Param>,
    pub body: BlockStmt,
    pub span: Span,
}
pub struct BlockStmt {
    pub statements: Vec<Statement>,
    pub span: Span,
}

pub enum Statement {
    Return(ReturnStmt),
    Block(BlockStmt),
    Expr(Expr),
    VarDecl(VarDeclStmt),
}
impl Statement {
    pub fn span(&self) -> Span {
        match self {
            Statement::Return(s) => s.span,
            Statement::Block(s) => s.span,
            Statement::Expr(s) => s.span,
            Statement::VarDecl(s) => s.span,
        }
    }
}
pub struct ReturnStmt {
    pub value: Option<Expr>,
    pub span: Span,
}
pub struct Expr {
    pub value: ExprValue,
    pub span: Span,
}
pub enum ExprValue {
    IntLit(i64),
    Var(String),
    BinaryOp {
        lhs: Box<Expr>, 
        op: BinaryOp, 
        rhs: Box<Expr>
    },
    FuncCall(String, Vec<Expr>),
    Assign {
        lvalue: Box<Expr>, 
        rvalue: Box<Expr>
    },
}

pub enum BinaryOp {
    Add, Sub, Mul, Div, Mod,
    Equal, NotEqual, Less, LessEqual, Greater, GreaterEqual,
}

impl BinaryOp {
    pub fn from_token_value(token_value: &TokenValue) -> Option<Self> {
        match token_value {
            TokenValue::Plus => Some(BinaryOp::Add),
            TokenValue::Minus => Some(BinaryOp::Sub),
            TokenValue::Star => Some(BinaryOp::Mul),
            TokenValue::Slash => Some(BinaryOp::Div),
            TokenValue::Percent => Some(BinaryOp::Mod),
            TokenValue::DoubleEqual => Some(BinaryOp::Equal),
            TokenValue::NotEqual => Some(BinaryOp::NotEqual),
            TokenValue::Less => Some(BinaryOp::Less),
            TokenValue::LessEqual => Some(BinaryOp::LessEqual),
            TokenValue::Greater => Some(BinaryOp::Greater),
            TokenValue::GreaterEqual => Some(BinaryOp::GreaterEqual),
            _ => None,
        }
    }
}
pub enum Type {
    Int,
    Void,
}
impl From<TypeIdent> for Type {
    fn from(value: TypeIdent) -> Self {
        match value {
            TypeIdent::Int => Type::Int,
            TypeIdent::Void => Type::Void,
        }
    }
}
pub struct Param {
    pub name: String,
    pub param_type: Type,
    pub span: Span,
}

impl fmt::Display for CompileUnit {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "CompileUnit")
    }
}

impl fmt::Display for GlobalDeclStmt {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            GlobalDeclStmt::VarDecl(_) => write!(f, "GlobalVarDecl"),
            GlobalDeclStmt::FuncDecl(_) => write!(f, "FuncDecl"),
        }
    }
}

impl fmt::Display for VarDeclStmt {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "VarDecl")
    }
}

impl fmt::Display for VarDeclStmtValue {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{} {}", self.var_type, self.name)
    }
}

impl fmt::Display for FuncDeclStmt {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{} {}", self.return_type, self.name)
    }
}

impl fmt::Display for BlockStmt {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "Block")
    }
}

impl fmt::Display for Statement {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Statement::Return(_) => write!(f, "ReturnStmt"),
            Statement::Block(_) => write!(f, "BlockStmt"),
            Statement::Expr(_) => write!(f, "ExprStmt"),
            Statement::VarDecl(_) => write!(f, "VarDeclStmt"),
        }
    }
}

impl fmt::Display for ReturnStmt {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "ReturnStmt")
    }
}

impl fmt::Display for Expr {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}", self.value)
    }
}

impl fmt::Display for ExprValue {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            ExprValue::IntLit(value) => write!(f, "IntLit({})", value),
            ExprValue::Var(name) => write!(f, "Var({})", name),
            ExprValue::BinaryOp { op, .. } => write!(f, "BinaryOp({})", op),
            ExprValue::FuncCall(name, _) => write!(f, "FuncCall({})", name),
            ExprValue::Assign { .. } => write!(f, "Assign"),
        }
    }
}

impl fmt::Display for Param {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{} {}", self.param_type, self.name)
    }
}

impl fmt::Display for BinaryOp {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let op = match self {
            BinaryOp::Add => "+",
            BinaryOp::Sub => "-",
            BinaryOp::Mul => "*",
            BinaryOp::Div => "/",
            BinaryOp::Mod => "%",
            BinaryOp::Equal => "==",
            BinaryOp::NotEqual => "!=",
            BinaryOp::Less => "<",
            BinaryOp::LessEqual => "<=",
            BinaryOp::Greater => ">",
            BinaryOp::GreaterEqual => ">=",
        };
        write!(f, "{}", op)
    }
}

impl fmt::Display for Type {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Type::Int => write!(f, "int"),
            Type::Void => write!(f, "void"),
        }
    }
}