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},
};

#[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 stream.peek() != Some(Token::Keyword(Keyword::End)) {
            statements.push(stream.parse()?);
        }

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

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

impl Parse for Statement {
    fn parse(mut stream: TokenStream) -> Result<Self, TokenStreamError> {
        let peeked = stream.peek();
        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 peeked == Some(Token::Keyword(Keyword::Local)) {
            stream.next();
            let name = stream.parse()?;
            stream.expect(Token::Symbol('='))?;
            let expr = stream.parse()?;
            Ok(Statement::Assignment(
                Some(DefinitionKind::Local),
                name,
                expr,
            ))
        } else if let Some(Token::Word(_)) = peeked
            && stream.peek2() == Some(Token::Symbol('='))
        {
            let name = stream.parse()?;
            stream.expect(Token::Symbol('='))?;
            Ok(Self::Assignment(None, name, stream.parse()?))
        } else {
            Err(stream.expecting_err("statement"))
        }
    }
}

#[derive(Debug, Clone)]
pub enum DefinitionKind {
    Local,
    Global,
}

#[derive(Debug, Clone)]
pub enum Expression {
    ValueRef(String),
    BinOp(BinaryOperator, Box<Node<Expression>>, Box<Node<Expression>>),
}

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 PrimaryExpression(Node<Expression>);

#[derive(Debug, Clone, Copy)]
pub enum BinaryOperator {
    Lt,
    Gt,
}

impl BinaryOperator {
    pub fn precedence(&self) -> u32 {
        match self {
            BinaryOperator::Lt => 100,
            BinaryOperator::Gt => 105,
        }
    }
}

impl Parse for PrimaryExpression {
    fn parse(mut stream: TokenStream) -> Result<Self, TokenStreamError> {
        let pre = Metadata::pre(&mut stream, "expression")?;
        Ok(PrimaryExpression(Node {
            kind: Expression::ValueRef(stream.parse()?),
            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)
}

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