rusty-minic/src/ast/types.rs

use crate::{diagnostic::span::Span, frontend::types::{Token, 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 data_type: Type,
    pub type_span: Span,
}
pub struct VarDeclStmtValue {
    pub name: String,
    pub name_span: Span,

}

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

pub enum Statement {
    Return(ReturnStmt),
    Block(BlockStmt),
    Expr(Expr),
    VarDecl(VarDeclStmt),
    If(IfStmt),
    While(WhileStmt),
    Break(BreakStmt),
    Continue(ContinueStmt),
}
// 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,
//             Statement::If(s) => s.span,
//             Statement::While(s) => s.span,
//             Statement::Break(s) => s.span,
//             Statement::Continue(s) => s.span,
//         }
//     }
// }
pub struct IfStmt {
    pub condition: Expr,
    pub then_branch: BlockStmt,
    pub ifelse_branch: Vec<IfElseBranch>,
    pub else_branch: Option<BlockStmt>,
    // pub span: Span,
}
pub struct IfElseBranch {
    pub condition: Expr,
    pub then_branch: BlockStmt,
}
pub struct WhileStmt {
    pub condition: Expr,
    pub body: BlockStmt,
    // pub span: Span,
}
pub struct BreakStmt {
    pub span: Span,
}
pub struct ContinueStmt {
    pub span: 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>
    },
    UnaryOp {
        op: UnaryOp,
        operand: Box<Expr>,
    },
    FuncCall(String, Vec<Expr>),
    Assign {
        lvalue: Box<Expr>, 
        rvalue: Box<Expr>
    },
}
#[derive(Clone, Copy)]
pub enum BinaryOp {
    Add, Sub, Mul, Div, Mod,
    Equal, NotEqual, Less, LessEqual, Greater, GreaterEqual,
    And, Or,
}

#[derive(Clone, Copy)]
pub enum UnaryOp {
    Add, Sub, Not,
}
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),
            TokenValue::And => Some(BinaryOp::And),
            TokenValue::Or => Some(BinaryOp::Or),
            _ => None,
        }
    }
    pub fn is_logical(&self) -> bool {
        matches!(self, BinaryOp::And | BinaryOp::Or)
    }
    pub fn is_cmp(&self) -> bool {
        matches!(self, BinaryOp::Equal | BinaryOp::NotEqual | BinaryOp::Less | BinaryOp::LessEqual | BinaryOp::Greater | BinaryOp::GreaterEqual)
    }
}


#[derive(Clone, Copy)]
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 name_span: Span,
    pub type_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.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"),
            Statement::If(_) => write!(f, "IfStmt"),
            Statement::While(_) => write!(f, "WhileStmt"),
            Statement::Break(_) => write!(f, "BreakStmt"),
            Statement::Continue(_) => write!(f, "ContinueStmt"),
        }
    }
}
impl fmt::Display for IfStmt {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "IfStmt")
    }
}
impl fmt::Display for ReturnStmt {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "ReturnStmt")
    }
}

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

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

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

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"),
            ExprValue::UnaryOp { op, .. } => write!(f, "UnaryOp({})", op),
        }
    }
}

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 => ">=",
            BinaryOp::And => "&&",
            BinaryOp::Or => "||",
        };
        write!(f, "{}", op)
    }
}
impl fmt::Display for UnaryOp {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let op = match self {
            UnaryOp::Add => "+",
            UnaryOp::Sub => "-",
            UnaryOp::Not => "!",
        };
        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"),
        }
    }
}