math-parser-rs 0.1.0

use super::{
    ast_types::{AstNode, OperatorKind, ParserError},
    optimizations::{general::constant_folding, trig::rewrite_cos_as_sin},
    types::{LexResult, Token, TokenType},
};

pub struct Parser {
    tokens: Vec<Token>,
    pos: usize,
}

impl Parser {
    pub fn new(tokens: LexResult) -> Parser {
        Parser {
            pos: 0,
            tokens: tokens
                .rhs
                .into_iter()
                .filter(|tok| tok.token_type != TokenType::Unimplemented)
                .collect(),
        }
    }

    fn optimize_ast(&self, mut ast: AstNode) -> AstNode {
        ast = rewrite_cos_as_sin(&ast);
        ast = constant_folding(&ast);
        ast
    }

    pub fn parse(&mut self) -> Result<AstNode, ParserError> {
        let ast = self.parse_expression()?;
        Ok(self.optimize_ast(ast))
    }

    /// Handlest lowest-precedance operations like `+` and `-`
    fn parse_expression(&mut self) -> Result<AstNode, ParserError> {
        let mut node = self.parse_term()?;

        while self.current_token_is(TokenType::Operator, Some("+"))
            || self.current_token_is(TokenType::Operator, Some("-"))
        {
            let op = if self.current_token_is(TokenType::Operator, Some("+")) {
                self.advance();
                OperatorKind::Plus
            } else {
                self.advance();
                OperatorKind::Minus
            };

            let right = self.parse_term()?;

            node = AstNode::BinaryOp {
                op,
                left: Box::new(node),
                right: Box::new(right),
            };
        }
        Ok(node)
    }
    /// Handlest higer-precedance operations like `*` and `/`
    fn parse_term(&mut self) -> Result<AstNode, ParserError> {
        let mut node = self.parse_factor()?;

        while self.current_token_is(TokenType::Operator, Some("*"))
            || self.current_token_is(TokenType::Operator, Some("/"))
        {
            let op = if self.current_token_is(TokenType::Operator, Some("*")) {
                self.advance();
                OperatorKind::Multiply
            } else {
                self.advance();
                OperatorKind::Divide
            };

            let right = self.parse_factor()?;

            node = AstNode::BinaryOp {
                op,
                left: Box::new(node),
                right: Box::new(right),
            };
        }
        Ok(node)
    }
    /// Handles number, parantheses
    fn parse_factor(&mut self) -> Result<AstNode, ParserError> {
        let current = self.current();
        match current {
            Some(tok) if self.current_token_is(TokenType::Number, None) => {
                let number = tok.text.parse::<f64>();
                match number {
                    Ok(number) => {
                        self.advance();
                        Ok(AstNode::Number(number))
                    }
                    Err(_) => Err(ParserError {
                        message: "Could not convert number to float".to_string(),
                    }),
                }
            }
            Some(tok) if self.current_token_is(TokenType::Identifier, None) => {
                let var_name = tok.text.clone();
                if let Some(next) = self.peek() {
                    if next.token_type == TokenType::ParaOpen {
                        self.advance(); // consume identifier
                        self.advance(); // consume '('
                        let mut args = Vec::new();
                        if !self.current_token_is(TokenType::ParaClose, None) {
                            loop {
                                let expr = self.parse_expression()?;
                                args.push(expr);

                                if self.current_token_is(TokenType::Comma, None) {
                                    self.advance(); // consume ','
                                    if self.current_token_is(TokenType::ParaClose, None) {
                                        return Err(ParserError {
                                                       message: "Unexpected ',' before ')' in function argument list".to_string(),
                                                   });
                                    }
                                } else if self.current_token_is(TokenType::ParaClose, None) {
                                    break;
                                } else {
                                    return Err(ParserError {
                                        message: "Expected ',' or ')' in function argument list"
                                            .to_string(),
                                    });
                                }
                            }
                        }
                        self.eat(TokenType::ParaClose, None);
                        Ok(AstNode::FunctionCall {
                            name: var_name,
                            args,
                        })
                    } else {
                        self.advance();
                        Ok(AstNode::Variable(var_name))
                    }
                } else {
                    self.advance();
                    Ok(AstNode::Variable(var_name))
                }
            }
            Some(_) if self.current_token_is(TokenType::ParaOpen, None) => {
                self.advance();
                let expr = self.parse_expression()?;
                if self.eat(TokenType::ParaClose, None) {
                    Ok(expr)
                } else {
                    Err(ParserError {
                        message: "Expected ')'".to_string(),
                    })
                }
            }
            Some(tok) => Err(ParserError {
                message: format!(
                    "Unexpected token: {:?} '{}', expected number or '('",
                    tok.token_type, tok.text
                ),
            }),
            None => Err(ParserError {
                message: "Unexpected end of input".to_string(),
            }),
        }
    }

    /// Pops and returns the current token, advancing the position.
    fn advance(&mut self) -> Option<&Token> {
        if self.pos < self.tokens.len() {
            let token = &self.tokens[self.pos];
            self.pos += 1;
            Some(token)
        } else {
            None
        }
    }

    /// Returns a reference to the next token without consuming it.
    fn peek(&self) -> Option<&Token> {
        if self.pos + 1 < self.tokens.len() {
            let token = &self.tokens[self.pos + 1];
            Some(token)
        } else {
            None
        }
    }

    /// Returns a reference to the current token.
    fn current(&self) -> Option<&Token> {
        if self.pos < self.tokens.len() {
            let token = &self.tokens[self.pos];
            Some(token)
        } else {
            None
        }
    }

    fn current_token_is(&self, ttype: TokenType, text: Option<&str>) -> bool {
        let current_token = self.current();
        match current_token {
            Some(token) => {
                if token.token_type != ttype {
                    return false;
                }
                match text {
                    Some(text) => text == token.text,
                    None => true,
                }
            }
            None => false,
        }
    }

    fn eat(&mut self, ttype: TokenType, text: Option<&str>) -> bool {
        let is_match = self.current_token_is(ttype, text);
        if is_match {
            self.advance();
        }
        is_match
    }
}