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 {
            Err(stream.expecting_err("statement"))
        }
    }
}

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

#[derive(Debug, Clone)]
pub enum Expression {
    ValueRef(String),
}

impl Parse for Expression {
    fn parse(mut stream: TokenStream) -> Result<Self, TokenStreamError> {
        Ok(Expression::ValueRef(stream.parse()?))
    }
}