gaybuild/src/ast/tree.rs

use std::fmt;

use crate::error::Error;

use crate::lex::token;
use crate::lex::token::{Position, Token};

pub enum Node {
    Break(Position),
    Ident(Position, String),
    DepList(Position, Box<Node>),
    SourceList(Position, Box<Node>),
    Bool(Position, bool),
    Int(Position, i128),
    String(Position, String),
    Array(Position, Vec<Node>),
    Block(Position, Vec<Node>),
    Fn(Position, FnStmt),
    ArrayExpr(Position, ArrayExpr),
    CallExpr(Position, CallExpr),
    IfStmt(Position, IfStmt),
    ReturnStmt(Position, Box<Node>),
    ForStmt(Position, ForStmt),
    WhileStmt(Position, WhileStmt),
    UnaryExpr(Position, UnaryExpr),
    BinaryExpr(Position, BinaryExpr),
    TypeStmt(Position, Box<Node>),
    SetStmt(Position, SetStmt),
    TargetStmt(Position, TargetStmt),
    File(Position, File),
}

pub struct File {
    pub name: String,
    pub content: Vec<Node>,
}

pub struct FnStmt {
    pub name: Option<Box<Node>>,
    pub params: Vec<Node>,
    pub body: Box<Node>,
}

pub struct IfStmt {
    pub condition: Box<Node>,
    pub then_body: Box<Node>,
    pub else_body: Option<Box<Node>>,
}

pub struct ForStmt {
    pub setup: Option<Box<Node>>,
    pub condition: Option<Box<Node>>,
    pub exec: Option<Box<Node>>,
    pub body: Box<Node>,
}

pub struct WhileStmt {
    pub condition: Box<Node>,
    pub body: Box<Node>,
}

pub struct SetStmt {
    pub name: Box<Node>,
    pub val: Box<Node>,
}

pub struct TargetStmt {
    pub name: Box<Node>,
    pub val: Box<Node>,
}

pub struct BinaryExpr {
    pub lhs: Box<Node>,
    pub op: Operator,
    pub rhs: Box<Node>,
}

pub struct UnaryExpr {
    pub op: Operator,
    pub rhs: Box<Node>,
}

pub struct CallExpr {
    pub func: Box<Node>,
    pub params: Vec<Node>,
}

pub struct ArrayExpr {
    pub array: Box<Node>,
    pub index: Box<Node>,
}

#[derive(Debug, PartialEq)]
pub enum Operator {
    Dot,
    Eq,
    EqEq,
    Bang,
    BangEq,
    Gt,
    GtGt,
    GtEq,
    GtGtEq,
    Lt,
    LtLt,
    LtEq,
    LtLtEq,
    Plus,
    PlusEq,
    Minus,
    MinusEq,
    Asterisk,
    AsteriskEq,
    Slash,
    SlashEq,
    Percent,
    PercentEq,
    Pipe,
    PipePipe,
    PipeEq,
    Amp,
    AmpAmp,
    AmpEq,
    Caret,
    CaretEq,
}

impl Node {
    pub fn assert_ident(&self) -> Result<&String, Error> {
        match self {
            Node::Ident(_, name) => Ok(name),
            _ => Err(Error::syntax_error(
                self.pos().clone(),
                format!("Expected an identifier, got {}", self),
            )),
        }
    }

    pub fn walk(&self, cb: fn(node: &Node, depth: u32)) {
        self.visit(cb, 0);
    }

    pub fn pos(&self) -> &Position {
        match self {
            Node::Break(pos) => pos,
            Node::Ident(pos, _) => pos,
            Node::Bool(pos, _) => pos,
            Node::Int(pos, _) => pos,
            Node::String(pos, _) => pos,
            Node::Block(pos, _) => pos,
            Node::Fn(pos, _) => pos,
            Node::ReturnStmt(pos, _) => pos,
            Node::DepList(pos, _) => pos,
            Node::SourceList(pos, _) => pos,
            Node::Array(pos, _) => pos,
            Node::ArrayExpr(pos, _) => pos,
            Node::CallExpr(pos, _) => pos,
            Node::IfStmt(pos, _) => pos,
            Node::ForStmt(pos, _) => pos,
            Node::WhileStmt(pos, _) => pos,
            Node::UnaryExpr(pos, _) => pos,
            Node::BinaryExpr(pos, _) => pos,
            Node::TypeStmt(pos, _) => pos,
            Node::SetStmt(pos, _) => pos,
            Node::TargetStmt(pos, _) => pos,
            Node::File(pos, _) => pos,
        }
    }

    fn visit(&self, cb: fn(node: &Node, depth: u32), current_depth: u32) {
        cb(self, current_depth);
        let depth = current_depth + 1;
        match self {
            Node::DepList(_, list) => list.visit(cb, depth),
            Node::SourceList(_, list) => list.visit(cb, depth),
            Node::Array(_, elements) => {
                for node in elements {
                    node.visit(cb, depth);
                }
            }
            Node::Block(_, statements) => {
                for node in statements {
                    node.visit(cb, depth);
                }
            }
            Node::Fn(_, func) => {
                if let Some(n) = &func.name {
                    n.visit(cb, depth);
                }
                for p in &func.params {
                    p.visit(cb, depth);
                }
                func.body.visit(cb, depth);
            }
            Node::ArrayExpr(_, expr) => {
                expr.array.visit(cb, depth);
                expr.index.visit(cb, depth);
            }
            Node::CallExpr(_, expr) => {
                expr.func.visit(cb, depth);
                for p in &expr.params {
                    p.visit(cb, depth);
                }
            }
            Node::IfStmt(_, stmt) => {
                stmt.condition.visit(cb, depth);
                stmt.then_body.visit(cb, depth);
                if let Some(eb) = &stmt.else_body {
                    eb.visit(cb, depth);
                }
            }
            Node::ForStmt(_, stmt) => {
                if let Some(s) = &stmt.setup {
                    s.visit(cb, depth);
                }
                if let Some(c) = &stmt.condition {
                    c.visit(cb, depth);
                }
                if let Some(e) = &stmt.exec {
                    e.visit(cb, depth);
                }
                stmt.body.visit(cb, depth);
            }
            Node::WhileStmt(_, stmt) => {
                stmt.condition.visit(cb, depth);
                stmt.body.visit(cb, depth);
            }
            Node::ReturnStmt(_, node) => node.visit(cb, depth),
            Node::UnaryExpr(_, expr) => expr.rhs.visit(cb, depth),
            Node::BinaryExpr(_, expr) => {
                expr.lhs.visit(cb, depth);
                expr.rhs.visit(cb, depth);
            }
            Node::TypeStmt(_, expr) => expr.visit(cb, depth),
            Node::SetStmt(_, expr) => {
                expr.name.visit(cb, depth);
                expr.val.visit(cb, depth);
            }
            Node::TargetStmt(_, stmt) => {
                stmt.name.visit(cb, depth);
                stmt.val.visit(cb, depth);
            }
            Node::File(_, file) => {
                for n in &file.content {
                    n.visit(cb, depth);
                }
            }
            _ => return,
        }
    }

    pub fn make_fn(pos: Position, name: Option<Node>, params: Vec<Node>, body: Node) -> Node {
        let name = if let Some(name) = name {
            Some(Box::new(name))
        } else {
            None
        };
        Node::Fn(pos, FnStmt { name, params, body: Box::new(body) })
    }

    pub fn make_binary_expr(pos: Position, lhs: Node, op: Operator, rhs: Node) -> Node {
        Node::BinaryExpr(pos, BinaryExpr { lhs: Box::new(lhs), op, rhs: Box::new(rhs) })
    }

    pub fn make_unary_expr(pos: Position, op: Operator, rhs: Node) -> Node {
        Node::UnaryExpr(pos, UnaryExpr { op, rhs: Box::new(rhs) })
    }

    pub fn make_call_expr(pos: Position, func: Node, params: Vec<Node>) -> Node {
        Node::CallExpr(pos, CallExpr { func: Box::new(func), params })
    }

    pub fn make_array_expr(pos: Position, array: Node, index: Node) -> Node {
        Node::ArrayExpr(pos, ArrayExpr { array: Box::new(array), index: Box::new(index) })
    }

    pub fn make_set_stmt(pos: Position, name: Node, val: Node) -> Node {
        Node::SetStmt(pos, SetStmt { name: Box::new(name), val: Box::new(val) })
    }

    pub fn make_target_stmt(pos: Position, name: Node, val: Node) -> Node {
        Node::TargetStmt(pos, TargetStmt { name: Box::new(name), val: Box::new(val) })
    }

    pub fn make_if_stmt(pos: Position, condition: Node, then_body: Node, else_body: Option<Node>) -> Node {
        let else_body = if let Some(else_body) = else_body {
            Some(Box::new(else_body))
        } else {
            None
        };
        Node::IfStmt(pos, IfStmt {
            condition: Box::new(condition),
            then_body: Box::new(then_body),
            else_body,
        })
    }

    pub fn make_while_stmt(pos: Position, condition: Node, body: Node) -> Node {
        Node::WhileStmt(pos, WhileStmt {
            condition: Box::new(condition),
            body: Box::new(body)
        })
    }

    pub fn make_for_stmt(pos: Position, setup: Option<Node>, condition: Option<Node>, exec: Option<Node>, body: Node) -> Node {
        let setup = if let Some(setup) = setup {
            Some(Box::new(setup))
        } else {
            None
        };
        let condition = if let Some(condition) = condition {
            Some(Box::new(condition))
        } else {
            None
        };
        let exec = if let Some(exec) = exec {
            Some(Box::new(exec))
        } else {
            None
        };
        Node::ForStmt(pos, ForStmt { setup, condition, exec, body: Box::new(body) })
    }
}

impl fmt::Display for Node {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let tmp: String;
        write!(
            f,
            "{}",
            match self {
                Node::Break(_) => "<break>",
                Node::Ident(_, name) => name.as_str(),
                Node::Bool(_, b) => {
                    tmp = format!("{}", b);
                    tmp.as_str()
                }
                Node::Int(_, i) => {
                    tmp = format!("{}", i);
                    tmp.as_str()
                }
                Node::String(_, s) => s.as_str(),
                Node::Block(_, _) => "<block>",
                Node::Fn(_, _) => "fn",
                Node::ReturnStmt(_, _) => "return",
                Node::DepList(_, _) => "depend",
                Node::SourceList(_, _) => "source",
                Node::Array(_, _) => "<array>",
                Node::ArrayExpr(_, _) => "<array-access>",
                Node::CallExpr(_, _) => "<call>",
                Node::IfStmt(_, _) => "<if>",
                Node::ForStmt(_, _) => "<for>",
                Node::WhileStmt(_, _) => "<while>",
                Node::UnaryExpr(_, expr) => expr.op.raw(),
                Node::BinaryExpr(_, expr) => expr.op.raw(),
                Node::TypeStmt(_, _) => "type",
                Node::SetStmt(_, _) => "set",
                Node::TargetStmt(_, _) => "target",
                Node::File(pos, _) => pos.file.as_str(),
            }
        )
    }
}

impl Operator {
    pub fn from_token(token: &Token) -> Result<Operator, Error> {
        match token.kind {
            token::Kind::Eq => Ok(Operator::Eq),
            token::Kind::EqEq => Ok(Operator::EqEq),
            token::Kind::BangEq => Ok(Operator::BangEq),
            token::Kind::Bang => Ok(Operator::Bang),
            token::Kind::Gt => Ok(Operator::Gt),
            token::Kind::GtGt => Ok(Operator::GtGt),
            token::Kind::GtEq => Ok(Operator::GtEq),
            token::Kind::GtGtEq => Ok(Operator::GtGtEq),
            token::Kind::Lt => Ok(Operator::Lt),
            token::Kind::LtLt => Ok(Operator::LtLt),
            token::Kind::LtEq => Ok(Operator::LtEq),
            token::Kind::LtLtEq => Ok(Operator::LtLtEq),
            token::Kind::Plus => Ok(Operator::Plus),
            token::Kind::PlusEq => Ok(Operator::PlusEq),
            token::Kind::Minus => Ok(Operator::Minus),
            token::Kind::MinusEq => Ok(Operator::MinusEq),
            token::Kind::Asterisk => Ok(Operator::Asterisk),
            token::Kind::AsteriskEq => Ok(Operator::AsteriskEq),
            token::Kind::Slash => Ok(Operator::Slash),
            token::Kind::SlashEq => Ok(Operator::SlashEq),
            token::Kind::Percent => Ok(Operator::Percent),
            token::Kind::PercentEq => Ok(Operator::PercentEq),
            token::Kind::Pipe => Ok(Operator::Pipe),
            token::Kind::PipePipe => Ok(Operator::PipePipe),
            token::Kind::PipeEq => Ok(Operator::PipeEq),
            token::Kind::Amp => Ok(Operator::Amp),
            token::Kind::AmpAmp => Ok(Operator::AmpAmp),
            token::Kind::AmpEq => Ok(Operator::AmpEq),
            token::Kind::Caret => Ok(Operator::Caret),
            token::Kind::CaretEq => Ok(Operator::CaretEq),
            _ => Err(Error::syntax_error(
                token.pos.clone(),
                format!("\"{}\" is not an operator", token.raw),
            )),
        }
    }

    pub fn raw(&self) -> &'static str {
        match self {
            Operator::Dot => ".",
            Operator::Eq => "=",
            Operator::EqEq => "==",
            Operator::BangEq => "!=",
            Operator::Gt => ">",
            Operator::GtGt => ">>",
            Operator::GtEq => ">=",
            Operator::GtGtEq => ">>=",
            Operator::Lt => "<",
            Operator::LtLt => "<<",
            Operator::LtEq => "<=",
            Operator::LtLtEq => "<<=",
            Operator::Plus => "+",
            Operator::PlusEq => "+=",
            Operator::Minus => "-",
            Operator::MinusEq => "-=",
            Operator::Asterisk => "*",
            Operator::AsteriskEq => "*=",
            Operator::Slash => "/",
            Operator::SlashEq => "/=",
            Operator::Percent => "%",
            Operator::PercentEq => "%=",
            Operator::Bang => "!",
            Operator::Pipe => "|",
            Operator::PipePipe => "||",
            Operator::PipeEq => "|=",
            Operator::Amp => "&",
            Operator::AmpAmp => "&&",
            Operator::AmpEq => "&=",
            Operator::Caret => "^",
            Operator::CaretEq => "^=",
        }
    }
}

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

#[derive(Debug, Clone, Copy, PartialEq)]
pub enum Type {
    /// For identifier tokens (evaluates at runtime)
    Unknown,
    /// For expressions that don't emit a value
    Void,
    Bool,
    Int,
    String,
    Array,
    Fn,
}

impl Type {
    pub fn from_token(token: &Token) -> Option<Type> {
        match token.kind {
            token::Kind::IntLiteral => Some(Type::Int),
            token::Kind::StringLiteral => Some(Type::String),
            _ => None,
        }
    }
}

impl fmt::Display for Type {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(
            f,
            "{}",
            match self {
                Type::Unknown => "<unknown>",
                Type::Void => "()",
                Type::Bool => "bool",
                Type::Int => "int",
                Type::String => "string",
                Type::Array => "array",
                Type::Fn => "fn",
            }
        )
    }
}