ferrite-lua/src/ast.rs

use std::{fmt::Debug, hash::Hash, ops::Add, path::PathBuf};

use crate::token_stream::{
    Parse, TokenRange, TokenStream, TokenStreamError,
    lexer::{Keyword, Position, Token},
};

pub type LuaNumber = f64;

#[derive(Debug, Clone)]
pub struct Node<T: Clone + Debug> {
    pub kind: T,
    pub meta: Metadata,
}

impl<T: Clone + Debug + PartialEq> PartialEq for Node<T> {
    fn eq(&self, other: &Self) -> bool {
        self.kind == other.kind
    }
}

impl<T: Clone + Debug + PartialEq> Eq for Node<T> {}

impl<T: Clone + Debug + PartialEq + Hash> Hash for Node<T> {
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        self.kind.hash(state);
    }
}

impl<T: Clone + Debug> Node<T> {
    pub fn empty(kind: T) -> Node<T> {
        Node {
            kind,
            meta: Metadata::empty(),
        }
    }

    pub fn with<OtherT: Clone + Debug>(&self, other: OtherT) -> Node<OtherT> {
        Node {
            kind: other,
            meta: self.meta.clone(),
        }
    }
}

#[derive(Clone, PartialEq, Hash, Eq)]
pub struct Metadata {
    pub documentation: Option<String>,
    pub token_range: TokenRange,
    pub position: Position,
    pub file_path: PathBuf,
}

impl Metadata {
    pub fn empty() -> Metadata {
        Metadata {
            documentation: None,
            token_range: Default::default(),
            position: Position(0, 0),
            file_path: PathBuf::new(),
        }
    }
}

impl Add<Metadata> for Metadata {
    type Output = Metadata;

    fn add(self, rhs: Metadata) -> Self::Output {
        Metadata {
            documentation: self.documentation,
            token_range: TokenRange {
                start: self.token_range.start,
                end: rhs.token_range.end,
            },
            position: self.position,
            file_path: self.file_path,
        }
    }
}

impl Debug for Metadata {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{:?}", self.token_range)
    }
}

impl<T: Parse + Clone + Debug> Parse for Node<T> {
    fn parse(mut stream: TokenStream) -> Result<Self, TokenStreamError> {
        let position = stream
            .get_position()
            .ok_or(stream.expecting_err(std::any::type_name::<T>()))?;
        let documentation = stream.find_documentation("/").into_iter().last();
        Ok(Node {
            kind: stream.parse()?,
            meta: Metadata {
                documentation,
                token_range: stream.get_range(),
                position,
                file_path: stream.file_path.clone(),
            },
        })
    }
}

impl Metadata {
    fn pre(
        stream: &mut TokenStream,
        expecting: &str,
    ) -> Result<(Option<String>, Position), TokenStreamError> {
        Ok((
            stream.find_documentation("/").into_iter().last(),
            stream
                .get_position()
                .ok_or(stream.expecting_err(expecting))?,
        ))
    }

    fn produce(
        stream: &mut TokenStream,
        (documentation, position): (Option<String>, Position),
    ) -> Metadata {
        Metadata {
            documentation: documentation,
            token_range: stream.get_range(),
            position,
            file_path: stream.file_path.clone(),
        }
    }
}

#[derive(Debug, Clone)]
pub struct Function {
    pub name: Option<Node<String>>,
    pub params: Vec<Node<String>>,
    pub block: Block,
    pub meta: Metadata,
}

impl Parse for Function {
    fn parse(mut stream: TokenStream) -> Result<Self, TokenStreamError> {
        let pre = Metadata::pre(&mut stream, "function")?;

        stream.expect(Token::Keyword(Keyword::Function))?;
        let name = stream.parse::<Node<String>>().ok();
        stream.expect(Token::Symbol('('))?;

        let mut params = Vec::new();

        if let Ok(param) = stream.parse() {
            params.push(param);
            while stream.peek() == Some(Token::Symbol(',')) {
                stream.next();
                params.push(stream.parse()?);
            }
        }

        stream.expect(Token::Symbol(')'))?;

        let block = stream.parse()?;

        stream.expect(Token::Keyword(Keyword::End))?;

        Ok(Function {
            name,
            params,
            block,
            meta: Metadata::produce(&mut stream, pre),
        })
    }
}

impl Parse for String {
    fn parse(mut stream: TokenStream) -> Result<Self, TokenStreamError> {
        if let Some(Token::Word(text)) = stream.next() {
            Ok(text)
        } else {
            Err(stream.expected_err("identifier"))
        }
    }
}

#[derive(Debug, Clone)]
pub struct Block {
    pub statements: Vec<Node<Statement>>,
    pub meta: Metadata,
}

impl Parse for Block {
    fn parse(mut stream: TokenStream) -> Result<Self, TokenStreamError> {
        let pre = Metadata::pre(&mut stream, "block")?;
        let mut statements = Vec::new();

        while !matches!(
            stream.peek(),
            Some(Token::Keyword(Keyword::End) | Token::Eof)
        ) {
            statements.push(stream.parse()?);
        }

        Ok(Block {
            statements,
            meta: Metadata::produce(&mut stream, pre),
        })
    }
}

#[derive(Debug, Clone)]
pub enum Statement {
    Assignment(
        Option<AccessModifier>,
        Vec<(Node<String>, Vec<Node<Expression>>)>,
        ExpressionList,
    ),
    Return(ExpressionList),
    If(Node<Expression>, Block),
    Expression(Node<Expression>),
}

impl Parse for Statement {
    fn parse(mut stream: TokenStream) -> Result<Self, TokenStreamError> {
        let peeked = stream.peek();
        if peeked == Some(Token::Keyword(Keyword::Function)) {
            let function = stream.parse::<Node<Function>>()?;
            if let Some(name) = function.kind.name {
                Ok(Self::Assignment(
                    None,
                    vec![(name, Vec::new())],
                    ExpressionList(vec![Node {
                        kind: Expression::FunctionDefinition(
                            function.kind.params,
                            function.kind.block,
                        ),
                        meta: function.meta,
                    }]),
                ))
            } else {
                Ok(Self::Expression(Node {
                    kind: Expression::FunctionDefinition(function.kind.params, function.kind.block),
                    meta: function.meta,
                }))
            }
        } else if peeked == Some(Token::Keyword(Keyword::Return)) {
            stream.next();
            Ok(Statement::Return(stream.parse()?))
        } else if peeked == Some(Token::Keyword(Keyword::If)) {
            stream.next(); // Consume if
            let cond = stream.parse()?;
            stream.expect(Token::Keyword(Keyword::Then))?;
            let then = stream.parse()?;
            stream.expect(Token::Keyword(Keyword::End))?;
            Ok(Self::If(cond, then))
        } else if let Some(Token::Keyword(Keyword::Local | Keyword::Global)) = peeked {
            let access_modifier = match stream.next() {
                Some(Token::Keyword(Keyword::Local)) => AccessModifier::Local,
                Some(Token::Keyword(Keyword::Global)) => AccessModifier::Global,
                _ => panic!(),
            };
            let mut names = Vec::new();
            names.push((stream.parse()?, Vec::new()));
            while stream.peek() == Some(Token::Symbol(',')) {
                stream.next();
                names.push((stream.parse()?, Vec::new()));
            }
            stream.expect(Token::Symbol('='))?;
            let expr = stream.parse()?;
            Ok(Statement::Assignment(Some(access_modifier), names, expr))
        } else if stream.parse_peek::<IndexedAssignment>().is_ok() {
            let access = stream.parse::<IndexedAssignment>().unwrap();

            let expression = stream.parse()?;

            Ok(Self::Assignment(
                None,
                vec![(access.0.0, access.0.1)],
                ExpressionList(vec![expression]),
            ))
        } else if let Ok(expr) = stream.parse() {
            Ok(Self::Expression(expr))
        } else {
            Err(stream.expecting_err("statement"))
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum AccessModifier {
    Local,
    Global,
}

#[derive(Debug, Clone)]
pub struct IndexedAccess(Node<String>, Vec<Node<Expression>>);

impl Parse for IndexedAccess {
    fn parse(mut stream: TokenStream) -> Result<Self, TokenStreamError> {
        let name = stream.parse()?;

        let mut expressions = Vec::new();
        while let Some(Token::Symbol('[') | Token::Symbol('.')) = stream.peek() {
            match stream.next().unwrap() {
                Token::Symbol('[') => {
                    let expression = stream.parse()?;
                    stream.expect_symbol(']')?;
                    expressions.push(expression);
                }
                Token::Symbol('.') => {
                    let word = stream.parse::<Node<String>>()?;
                    expressions.push(Node {
                        kind: Expression::Literal(Literal::String(word.kind)),
                        meta: word.meta,
                    });
                }
                _ => panic!(),
            }
        }

        Ok(IndexedAccess(name, expressions))
    }
}

#[derive(Debug, Clone)]
pub struct IndexedAssignment(IndexedAccess);

impl Parse for IndexedAssignment {
    fn parse(mut stream: TokenStream) -> Result<Self, TokenStreamError> {
        let access = stream.parse()?;
        stream.expect_symbol('=')?;
        Ok(IndexedAssignment(access))
    }
}

#[derive(Debug, Clone)]
pub enum Expression {
    ValueRef(String),
    UnOp(UnaryOperator, Box<Node<Expression>>),
    BinOp(BinaryOperator, Box<Node<Expression>>, Box<Node<Expression>>),
    FunctionDefinition(Vec<Node<String>>, Block),
    FunctionCall(Box<Node<Expression>>, Node<ExpressionList>),
    Literal(Literal),
    TableConstructor,
}

impl Parse for Expression {
    fn parse(mut stream: TokenStream) -> Result<Self, TokenStreamError> {
        let primary_expr = stream.parse::<PrimaryExpression>()?;
        parse_binop_rhs(&mut stream, primary_expr, None).map(|v| v.kind)
    }
}

#[derive(Debug, Clone)]
pub struct ExpressionList(pub Vec<Node<Expression>>);

impl Parse for ExpressionList {
    fn parse(mut stream: TokenStream) -> Result<Self, TokenStreamError> {
        let mut list = Vec::new();
        if let Ok(value) = stream.parse() {
            list.push(value);
            while stream.peek() == Some(Token::Symbol(',')) {
                stream.next();
                list.push(stream.parse()?);
            }
        }

        Ok(ExpressionList(list))
    }
}

#[derive(Debug, Clone)]
pub struct PrimaryExpression(Node<Expression>);

#[derive(Debug, Clone, Copy)]
pub enum UnaryOperator {
    Negation,
}

impl Parse for UnaryOperator {
    fn parse(mut stream: TokenStream) -> Result<Self, TokenStreamError> {
        if let Some(token) = stream.next() {
            match token {
                Token::Symbol('-') => Ok(UnaryOperator::Negation),
                _ => Err(stream.expected_err("unop")),
            }
        } else {
            Err(stream.expected_err("unop"))
        }
    }
}

#[derive(Debug, Clone, Copy)]
pub enum BinaryOperator {
    And,
    Or,
    LessThan,
    LessThanOrEqual,
    GreaterThan,
    GreaterThanOrEqual,
    Equal,
    Add,
    Sub,
}

impl BinaryOperator {
    pub fn precedence(&self) -> u32 {
        match self {
            BinaryOperator::Or => 0,
            BinaryOperator::And => 1,
            BinaryOperator::LessThan => 10,
            BinaryOperator::LessThanOrEqual => 10,
            BinaryOperator::GreaterThan => 10,
            BinaryOperator::GreaterThanOrEqual => 10,
            BinaryOperator::Equal => 10,
            BinaryOperator::Add => 20,
            BinaryOperator::Sub => 20,
        }
    }
}

impl Parse for BinaryOperator {
    fn parse(mut stream: TokenStream) -> Result<Self, TokenStreamError> {
        if let Some(token) = stream.next() {
            match (token, stream.peek()) {
                (Token::Word(word), _) => match word.as_str() {
                    "and" => Ok(BinaryOperator::And),
                    "or" => Ok(BinaryOperator::Or),
                    _ => Err(stream.expected_err("binop")),
                },
                (Token::Symbol('<'), Some(Token::Symbol('='))) => {
                    stream.next();
                    Ok(BinaryOperator::LessThanOrEqual)
                }
                (Token::Symbol('>'), Some(Token::Symbol('='))) => {
                    stream.next();
                    Ok(BinaryOperator::GreaterThanOrEqual)
                }
                (Token::Symbol('='), Some(Token::Symbol('='))) => {
                    stream.next();
                    Ok(BinaryOperator::Equal)
                }
                (Token::Symbol('<'), _) => Ok(BinaryOperator::LessThan),
                (Token::Symbol('>'), _) => Ok(BinaryOperator::GreaterThan),
                (Token::Symbol('+'), _) => Ok(BinaryOperator::Add),
                (Token::Symbol('-'), _) => Ok(BinaryOperator::Sub),
                _ => Err(stream.expected_err("binop")),
            }
        } else {
            Err(stream.expected_err("binop"))
        }
    }
}

impl Parse for PrimaryExpression {
    fn parse(mut stream: TokenStream) -> Result<Self, TokenStreamError> {
        let pre = Metadata::pre(&mut stream, "expression")?;

        let mut unary_operators = Vec::new();

        while let Ok(unop) = stream.parse::<UnaryOperator>() {
            unary_operators.push(unop);
        }

        let peeked = stream.peek();
        let mut expression = if peeked == Some(Token::Keyword(Keyword::Function)) {
            let function = stream.parse::<Node<Function>>()?;
            Expression::FunctionDefinition(function.kind.params, function.kind.block)
        } else if let Some(Token::DecimalValue(value)) = peeked {
            stream.next(); // Consume decimal value
            Expression::Literal(Literal::Number(
                u64::from_str_radix(&value, 10).unwrap() as f64
            ))
        } else if let Some(Token::StringLit(value)) = peeked {
            stream.next(); // Consume string-literal
            Expression::Literal(Literal::String(value))
        } else if let Some(Token::Symbol('{')) = peeked {
            stream.next();
            stream.expect_symbol('}')?;
            Expression::TableConstructor
        } else {
            Expression::ValueRef(stream.parse()?)
        };

        while let Some(Token::Symbol('(') | Token::Symbol('[')) = stream.peek() {
            match stream.next().unwrap() {
                Token::Symbol('(') => {
                    let expression_list = stream.parse::<Node<ExpressionList>>()?;
                    stream.expect(Token::Symbol(')'))?;
                    expression = Expression::FunctionCall(
                        Box::new(Node {
                            kind: expression,
                            meta: Metadata::produce(&mut stream, pre.clone()),
                        }),
                        expression_list,
                    );
                }
                Token::Symbol('[') => todo!(),
                _ => panic!(),
            }
        }

        for unop in unary_operators.into_iter().rev() {
            expression = Expression::UnOp(
                unop,
                Box::new(Node {
                    kind: expression,
                    meta: Metadata::produce(&mut stream, pre.clone()),
                }),
            );
        }

        Ok(PrimaryExpression(Node {
            kind: expression,
            meta: Metadata::produce(&mut stream, pre),
        }))
    }
}

fn parse_binop_rhs(
    stream: &mut TokenStream,
    mut lhs: PrimaryExpression,
    prev_op: Option<BinaryOperator>,
) -> Result<Node<Expression>, TokenStreamError> {
    let meta_pre = Metadata::pre(stream, "binary expression")?;

    let precedence = if let Some(op) = prev_op {
        op.precedence()
    } else {
        0
    };

    while let Ok(curr_operator) =
        stream.parse_if::<BinaryOperator, _>(|op| op.precedence() >= precedence)
    {
        let mut rhs = stream.parse::<PrimaryExpression>()?;

        if let Ok(next_op) = stream.parse_peek::<BinaryOperator>() {
            if curr_operator.precedence() < next_op.precedence() {
                // Operator on the right of rhs has more precedence, turn
                // rhs into lhs for new binop
                let rhs_expr = stream
                    .parse_with(|mut st| parse_binop_rhs(&mut st, rhs, Some(curr_operator)))?;
                rhs = PrimaryExpression(rhs_expr);
            }
        }

        lhs = PrimaryExpression(Node {
            kind: Expression::BinOp(curr_operator, Box::new(lhs.0), Box::new(rhs.0)),
            meta: Metadata::produce(stream, meta_pre.clone()),
        });
    }

    Ok(lhs.0)
}

#[derive(Debug, Clone)]
pub enum Literal {
    Number(LuaNumber),
    String(String),
}