ferric_parser/parser/

parser.rs

//! The ferric_parser module provides tools for parsing a sequence of tokens
//! into an abstract syntax tree (AST) and performing optimizations on it.

use ferric_lexer::{Keyword, Operator, Token};
use crate::parser::code_gen::ast_node::ASTNode;
use std::collections::HashSet;
use crate::parser::symbol_table::SymbolTable;

/// Represents the parser which takes a sequence of tokens and produces an AST.
///
/// # Example
///
/// ```rust
/// use ferric_parser::parser::Parser;
/// use ferric_lexer::{Token, Lexer};
///
/// let code = "let x = 5; let y = x + 3;";
/// let tokens = Lexer::new(code).lex();
/// let mut parser = Parser::new(tokens);
/// let ast = parser.parse().unwrap();
/// ```
pub struct Parser {
    tokens: Vec<Token>,
    position: usize,
    symtable: SymbolTable,
}

/// Propagates constants through the AST, simplifying expressions
/// where possible using the symbol table.
fn propagate_constants(node: ASTNode, symtable: &mut SymbolTable) -> ASTNode {
    match node {
        ASTNode::Program(statements) => {
            let optimized_statements = statements
                .into_iter()
                .map(|stmt| propagate_constants(stmt, symtable))
                .collect();
            ASTNode::Program(optimized_statements)
        }
        ASTNode::Statement(inner) => {
            ASTNode::Statement(Box::new(propagate_constants(*inner, symtable)))
        }
        ASTNode::Assign { name, value } => {
            let new_value = propagate_constants(*value, symtable);

            // If the new_value is a number, update the symbol table
            if let ASTNode::Number(val) = &new_value {
                symtable.set(&name, *val);
            }

            ASTNode::Assign {
                name,
                value: Box::new(new_value),
            }
        }
        ASTNode::BuiltInFunctionCall { name, args } => {
            let optimized_args = args
                .into_iter()
                .map(|arg| propagate_constants(arg, symtable))
                .collect();
            ASTNode::BuiltInFunctionCall {
                name,
                args: optimized_args,
            }
        }
        ASTNode::Variable(name) => {
            if let Some(value) = symtable.get(&name) {
                ASTNode::Number(value)
            } else {
                ASTNode::Variable(name)
            }
        }
        ASTNode::BinaryOp { op, left, right } => ASTNode::BinaryOp {
            op,
            left: Box::new(propagate_constants(*left, symtable)),
            right: Box::new(propagate_constants(*right, symtable)),
        },
        _ => node,
    }
}

/// Evaluates and simplifies binary operations on constants.
fn fold_constants(node: ASTNode) -> ASTNode {
    match node {
        ASTNode::Program(statements) => {
            let folded_statements = statements.into_iter().map(fold_constants).collect();
            ASTNode::Program(folded_statements)
        }
        ASTNode::Statement(inner) => ASTNode::Statement(Box::new(fold_constants(*inner))),
        ASTNode::Assign { name, value } => ASTNode::Assign {
            name,
            value: Box::new(fold_constants(*value)),
        },
        ASTNode::BuiltInFunctionCall { name, args } => {
            let folded_args = args.into_iter().map(fold_constants).collect();
            ASTNode::BuiltInFunctionCall {
                name,
                args: folded_args,
            }
        }
        ASTNode::BinaryOp { op, left, right } => {
            let left = fold_constants(*left);
            let right = fold_constants(*right);
            if let (ASTNode::Number(l_val), ASTNode::Number(r_val)) = (&left, &right) {
                let result = match op {
                    Operator::Plus => l_val + r_val,
                    Operator::Minus => l_val - r_val,
                    Operator::Multiply => l_val * r_val,
                    Operator::Divide => l_val / r_val,
                    Operator::Modulo => l_val % r_val,
                    _ => {
                        return ASTNode::BinaryOp {
                            op,
                            left: Box::new(left),
                            right: Box::new(right),
                        }
                    }
                };
                ASTNode::Number(result)
            } else {
                ASTNode::BinaryOp {
                    op,
                    left: Box::new(left),
                    right: Box::new(right),
                }
            }
        }
        _ => node,
    }
}

/// Collects the variables used within the AST.
fn collect_used_variables(node: &ASTNode, used_vars: &mut HashSet<String>) {
    match node {
        ASTNode::Variable(name) => {
            used_vars.insert(name.clone());
        }
        ASTNode::Program(statements) => {
            for stmt in statements {
                collect_used_variables(stmt, used_vars);
            }
        }
        ASTNode::BinaryOp { left, right, .. } => {
            collect_used_variables(left, used_vars);
            collect_used_variables(right, used_vars);
        }
        ASTNode::Assign { value, .. } => {
            collect_used_variables(value, used_vars);
        }
        ASTNode::BuiltInFunctionCall { args, .. } | ASTNode::UserFunctionCall { args, .. } => {
            for arg in args {
                collect_used_variables(arg, used_vars);
            }
        }
        ASTNode::Statement(inner) => match &**inner {
            ASTNode::Assign { value, .. } => {
                collect_used_variables(value, used_vars);
            }
            _ => collect_used_variables(&*inner, used_vars),
        },
        // Handle other node types as necessary
        _ => {}
    }
}

/// Filters out assignments to variables that are not used in the AST.
fn filter_unused_assignments(node: ASTNode, used_vars: &HashSet<String>) -> Option<ASTNode> {
    let result = match node {
        ASTNode::Program(statements) => {
            let filtered_statements: Vec<ASTNode> = statements
                .into_iter()
                .filter_map(|stmt| {
                    if let ASTNode::Assign { name, .. } = &stmt {
                        if !used_vars.contains(name) {
                            return None; // Filter out unused assignments
                        }
                        return filter_unused_assignments(stmt, used_vars);
                    }
                    Some(stmt)
                })
                .collect();
            Some(ASTNode::Program(filtered_statements))
        }
        ASTNode::Statement(inner) => {
            if let Some(filtered_inner) = filter_unused_assignments(*inner, used_vars) {
                Some(ASTNode::Statement(Box::new(filtered_inner)))
            } else {
                None
            }
        }
        // ... handle other cases
        _ => Some(node),
    };

    result
}

/// Removes any variables from the AST that aren't being used.
fn remove_unused_variables(ast: ASTNode) -> ASTNode {
    // Step 1: Collect used variables
    let mut used_vars = HashSet::new();
    collect_used_variables(&ast, &mut used_vars);

    // Step 2: Filter out unused assignments
    filter_unused_assignments(ast, &used_vars).unwrap()
}

impl Parser {
    /// Creates a new parser for a given sequence of tokens.
    pub fn new(tokens: Vec<Token>) -> Self {
        Parser {
            tokens,
            position: 0,
            symtable: SymbolTable::new(),
        }
    }

    /// Parses the sequence of tokens into an AST.
    pub fn parse(&mut self) -> Result<ASTNode, String> {
        let mut statements = Vec::new();

        while self.position < self.tokens.len() {
            let stmt = self.statement()?;
            statements.push(stmt);

            if let Some(Token::EoL) = self.peek_next_token() {
                self.position += 1; // Consume EoL
            } else {
                return Err(format!("Expected EoL, found {:?}", self.peek_next_token()));
            }
        }
        let mut ast = ASTNode::Program(statements);
        let mut previous_ast = ast.clone();

        loop {
            let optimized_ast = remove_unused_variables(self.optimize(ast.clone()));
            if optimized_ast == previous_ast {
                break;
            }
            previous_ast = ast;
            ast = optimized_ast;
        }

        Ok(ast)
    }

    /// Consumes optional EoL tokens from the token stream.
    fn consume_optional_eol(&mut self) {
        if let Some(Token::EoL) = self.peek_next_token() {
            self.position += 1; // Consume EoL
        }
    }

    /// Parses a statement from the token stream.
    fn statement(&mut self) -> Result<ASTNode, String> {
        let stmt = match self.peek_next_token() {
            Some(Token::Keyword(Keyword::Let)) => parse_variable_declaration!(self),
            Some(Token::BuiltInFunction(_)) => parse_func_call!(self),
            Some(Token::Keyword(Keyword::If)) => parse_if_statement!(self),
            Some(Token::Keyword(Keyword::While)) => parse_while_loop!(self),
            Some(Token::Number(_)) => self.expression(),
            Some(Token::Identifier(_)) => {
                if let Some(Token::Assign) = self.peek_nth_token(1) {
                    self.expect_assignment()
                } else {
                    self.expression()
                }
            }
            _ => Err(format!(
                "Expected a statement, found {:?}",
                self.peek_next_token()
            )),
        }?;

        Ok(ASTNode::Statement(Box::new(stmt)))
    }

    /// Parses function arguments from the token stream.
    fn parse_function_arguments(&mut self) -> Result<Vec<ASTNode>, String> {
        let mut args = Vec::new();

        // Expecting an open parenthesis after function name
        match &self.tokens[self.position] {
            Token::OpenParen => self.position += 1,
            _ => {
                return Err(format!(
                    "Expected '(' after function name, found {:?}",
                    &self.tokens[self.position]
                ))
            }
        }

        // Parse arguments until a close parenthesis is found
        while self.tokens[self.position] != Token::CloseParen {
            let expr = self.expression()?;
            args.push(expr);

            // If the next token is a comma, consume it and continue parsing the next argument
            if self.tokens[self.position] == Token::Comma {
                self.position += 1;
            }
        }

        // Consuming the close parenthesis
        self.position += 1;

        Ok(args)
    }

    /// Parses a block (sequence of statements) from the token stream.
    fn parse_block(&mut self) -> Result<ASTNode, String> {
        let mut statements = Vec::new();

        loop {
            match self.peek_next_token() {
                Some(Token::CloseBrace) => {
                    self.position += 1;
                    break;
                }
                Some(Token::EoL) => {
                    self.position += 1; // consume the EoL
                }
                _ => {
                    let stmt = self.statement()?;
                    statements.push(stmt);
                    self.consume_optional_eol();
                }
            }
        }

        Ok(ASTNode::Program(statements))
    }

    /// Parses an expression from the token stream.
    fn expression(&mut self) -> Result<ASTNode, String> {
        let mut left = self.term()?;

        while let Some(op) = self.peek_next_operator(&[
            Operator::Plus,
            Operator::Minus,
            //Operator::Modulo,
            Operator::Equal,
            Operator::NotEqual,
            Operator::GreaterThan,
            Operator::LessThan,
            Operator::GreaterThanOrEqual,
            Operator::LessThanOrEqual,
        ]) {
            self.position += 1; // Consume Operator
            let right = self.term()?;
            left = ASTNode::BinaryOp {
                op,
                left: Box::new(left),
                right: Box::new(right),
            };
        }

        Ok(left)
    }

    /// Parses a term from the token stream.
    fn term(&mut self) -> Result<ASTNode, String> {
        let mut left = self.factor()?;

        while let Some(op) =
            self.peek_next_operator(&[Operator::Multiply, Operator::Divide, Operator::Modulo])
        {
            self.position += 1; // Consume Operator
            let right = self.factor()?;
            left = ASTNode::BinaryOp {
                op,
                left: Box::new(left),
                right: Box::new(right),
            };
        }

        Ok(left)
    }

    /// Parses a factor from the token stream.
    fn factor(&mut self) -> Result<ASTNode, String> {
        if let Some(Token::Operator(Operator::Plus)) = self.peek_next_token() {
            self.position += 1; // Consume '+'
            return self.factor();
        }

        if let Some(Token::Operator(Operator::Minus)) = self.peek_next_token() {
            self.position += 1; // Consume '-'
            let right = self.primary()?;
            return Ok(ASTNode::BinaryOp {
                op: Operator::Minus,
                left: Box::new(ASTNode::Number(0)), // Represent as 0 - right
                right: Box::new(right),
            });
        }

        self.primary()
    }

    /// Parses a primary expression from the token stream.
    fn primary(&mut self) -> Result<ASTNode, String> {
        let next_token = self.peek_next_token().cloned();
        // First, handle expressions enclosed in parentheses
        if let Some(Token::OpenParen) = next_token {
            self.position += 1; // Consume '('
            let expr = self.expression()?; // Evaluate the enclosed expression
            if let Some(Token::CloseParen) = self.peek_next_token() {
                self.position += 1; // Consume ')'
                return Ok(expr);
            } else {
                return Err(format!("Expected ')', found {:?}", self.peek_next_token()));
            }
        }

        if let Some(Token::BuiltInFunction(_)) = next_token {
            return parse_func_call!(self);
        }

        if let Some(Token::StringLiteral(s)) = next_token {
            self.position += 1;
            return Ok(ASTNode::StringLiteral(s));
        }

        // Next, handle numbers and variables
        match next_token {
            Some(Token::Number(_)) => self.expect_number(),
            Some(Token::Identifier(_)) => {
                if let Some(Token::Assign) = self.peek_nth_token(1) {
                    self.expect_assignment()
                } else {
                    self.expect_variable()
                }
            }
            // Here you can extend to handle other primary expressions
            _ => Err(format!(
                "Expected a primary expression, found {:?}",
                self.peek_next_token()
            )),
        }
    }

    /// Peeks at the next token in the token stream without consuming it.
    fn peek_next_token(&self) -> Option<&Token> {
        if self.position < self.tokens.len() {
            Some(&self.tokens[self.position])
        } else {
            None
        }
    }

    /// Peeks at the next operator token in the token stream without consuming it.
    fn peek_next_operator(&self, ops: &[Operator]) -> Option<Operator> {
        if self.position >= self.tokens.len() {
            return None;
        }

        match self.tokens[self.position] {
            Token::Operator(ref op) if ops.contains(op) => Some(op.clone()),
            _ => None,
        }
    }

    /// Peeks at the nth token in the token stream without consuming it.
    fn peek_nth_token(&self, n: usize) -> Option<&Token> {
        if self.position + n < self.tokens.len() {
            Some(&self.tokens[self.position + n])
        } else {
            None
        }
    }

    /// Expects and parses an assignment from the token stream.
    fn expect_assignment(&mut self) -> Result<ASTNode, String> {
        let name = match &self.tokens[self.position] {
            Token::Identifier(ref ident) => ident.clone(),
            _ => {
                return Err(format!(
                    "Expected an identifier, found {:?}",
                    self.tokens[self.position]
                ))
            }
        };
        self.position += 2; // Consume Identifier and Assign
        let value = self.expression()?;
        Ok(ASTNode::Assign {
            name,
            value: Box::new(value),
        })
    }

    /// Expects and parses a variable from the token stream.
    fn expect_variable(&mut self) -> Result<ASTNode, String> {
        expect_token!(self, Identifier(name) => Variable)
    }

    /// Expects and parses a number from the token stream.
    fn expect_number(&mut self) -> Result<ASTNode, String> {
        expect_token!(self, Number(val) => Number)
    }

    /// Parses and optimizes an AST.
    ///
    /// This method first parses the tokens into an AST and then performs
    /// constant propagation and constant folding to optimize the AST.
    ///
    /// # Example
    ///
    /// Assuming Ferric code looks like:
    /// ```
    /// let x = 5;
    /// let y = x + 3;
    /// ```
    ///
    /// The resulting AST after parsing would represent the above code structure.
    ///
    /// ```rust
    /// use ferric_parser::parser::Parser;
    /// use ferric_lexer::Lexer;
    ///
    /// let code = "let x = 5; let y = x + 3;";
    /// let tokens = Lexer::new(code).lex();
    /// let mut parser = Parser::new(tokens);
    /// let ast = parser.parse().unwrap();
    /// let optimized_ast = parser.optimize(ast);
    /// ```
    pub fn optimize(&mut self, node: ASTNode) -> ASTNode {
        let node = propagate_constants(node, &mut self.symtable);
        fold_constants(node)
    }
}