Math expression to AST in Rust

main.rs

use std::iter::Peekable;

/// Lexer for parsing basic mathematical expressions.
struct Lexer<'a> {
    /// The input string to tokenize.
    input: &'a str,
    /// The current position of the lexer in the input string.
    current_position: usize,
}

impl<'a> Lexer<'a> {
    /// Creates a new lexer instance with the given input string.
    fn new(input: &'a str) -> Lexer<'a> {
        Lexer {
            input,
            current_position: 0,
        }
    }

    /// Returns the next token in the input string, or None if there are no more tokens.
    fn next_token(&mut self) -> Option<Token> {
        let input_chars: Vec<char> = self.input.chars().collect();

        while self.current_position < input_chars.len() {
            let current_char = input_chars[self.current_position];

            if current_char.is_digit(10) {
                let mut value = current_char.to_string();
                self.current_position += 1;

                while self.current_position < input_chars.len()
                    && (input_chars[self.current_position].is_digit(10)
                        || input_chars[self.current_position] == '.')
                {
                    value.push(input_chars[self.current_position]);
                    self.current_position += 1;
                }

                return Some(Token::Number(value.parse().unwrap()));
            }

            match current_char {
                '+' => {
                    self.current_position += 1;
                    return Some(Token::Plus);
                }
                '-' => {
                    self.current_position += 1;
                    return Some(Token::Minus);
                }
                '*' => {
                    self.current_position += 1;
                    return Some(Token::Multiply);
                }
                '/' => {
                    self.current_position += 1;
                    return Some(Token::Divide);
                }
                '(' => {
                    self.current_position += 1;
                    return Some(Token::LeftParenthesis);
                }
                ')' => {
                    self.current_position += 1;
                    return Some(Token::RightParenthesis);
                }
                _ => {
                    self.current_position += 1;
                }
            }
        }

        None
    }
}

/// Represents a token in a mathematical expression.
#[derive(Debug, PartialEq)]
enum Token {
    Plus,
    Minus,
    Multiply,
    Divide,
    LeftParenthesis,
    RightParenthesis,
    Number(f64),
}

/// Represents an expression in an abstract syntax tree (AST).
#[derive(Debug)]
enum Expr {
    Binary(Box<Expr>, char, Box<Expr>),
    Number(f64),
}

/// Builds an AST from a vector of tokens.
fn build_ast(tokens: Vec<Token>) -> Expr {
    let mut iter = tokens.into_iter().peekable();
    build_expr(&mut iter)
}

/// Builds an expression from a vector of tokens, starting with the first term.
fn build_expr(iter: &mut Peekable<std::vec::IntoIter<Token>>) -> Expr {
    let mut lhs = build_term(iter);

    loop {
        match iter.peek() {
            Some(Token::Plus) => {
                iter.next();
                lhs = Expr::Binary(Box::new(lhs), '+', Box::new(build_term(iter)));
            }
            Some(Token::Minus) => {
                iter.next();
                lhs = Expr::Binary(Box::new(lhs), '-', Box::new(build_term(iter)));
            }
            _ => break,
        }
    }

    lhs
}

/// Builds a term from a vector of tokens, starting with the first factor.
fn build_term(iter: &mut Peekable<std::vec::IntoIter<Token>>) -> Expr {
    let mut lhs = build_factor(iter);

    loop {
        match iter.peek() {
            Some(Token::Multiply) => {
                iter.next();
                lhs = Expr::Binary(Box::new(lhs), '*', Box::new(build_factor(iter)));
            }
            Some(Token::Divide) => {
                iter.next();
                lhs = Expr::Binary(Box::new(lhs), '/', Box::new(build_factor(iter)));
            }
            _ => break,
        }
    }

    lhs
}

/// Builds a factor from a vector of tokens.
fn build_factor(iter: &mut Peekable<std::vec::IntoIter<Token>>) -> Expr {
    match iter.next().unwrap() {
        Token::Number(n) => Expr::Number(n),
        Token::LeftParenthesis => {
            let expr = build_expr(iter);
            assert_eq!(iter.next(), Some(Token::RightParenthesis));
            expr
        }
        _ => panic!("Unexpected token!"),
    }
}

/// Parses a mathematical expression and returns its AST.
macro_rules! expression_to_ast {
    ($input:expr) => {{
        let mut lexer = Lexer::new($input);
        let mut tokens = vec![];

        loop {
            match lexer.next_token() {
                Some(token) => tokens.push(token),
                None => break,
            }
        }

        build_ast(tokens)
    }};
}

fn main() {
    let input = "1.5 + 2.5 * (3 - 1)";

    println!("{:?}", expression_to_ast!(input));
}