rlox/src/frontend/parser.rs

use crate::{
    frontend::{
        ast::{AstNode, Expr, Stmt},
        source_registry::SourceSlice,
        tokens::{LiteralValue, Token, TokenType},
    },
    logging::display_ast::{pretty_print_with_config, PrettyConfig},
    result::{parse_error, LoxError, LoxResult},
};

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

impl Parser {
    pub fn new(tokens: Vec<Token>) -> Self {
        Self { tokens, current: 0 }
    }

    pub fn parse(&mut self) -> LoxResult<Vec<AstNode<Stmt>>> {
        let mut statements = Vec::new();
        let mut errors = Vec::new();
        println!("\n \n ++++++++++++++++++++++++++++++++++++\nStarting parsing");
        // Con configurazione personalizzata
        let config = PrettyConfig {
            indent: "    ".to_string(),
            max_depth: None,
            show_positions: false,
            ..Default::default()
        };
        while !self.is_at_end() {
            match self.statement() {
                Ok(stmt) => {
                    println!(
                        "Parsed Successfully: {}",
                        pretty_print_with_config(&stmt.node, &config)
                    );

                    statements.push(stmt.clone());
                }
                Err(err) => {
                    println!("Error: {:?}", err);
                    errors.push(err);
                    let old_current = self.current;
                    self.synchronize();
                    println!("Synchronized from {} to {}", old_current, self.current);
                }
            }
        }

        if !errors.is_empty() {
            return Err(errors.into_iter().next().unwrap());
        }

        Ok(statements)
    }

    fn statement(&mut self) -> LoxResult<AstNode<Stmt>> {
        match (&self.peek().token_type, &self.peek_next().token_type) {
            (TokenType::Var, _) => self.var_statement(),
            (TokenType::Print, _) => self.print_statement(),
            (TokenType::Return, _) => self.return_statement(),
            (TokenType::StartBlock, _) => self.block_statement(),
            (TokenType::Identifier, TokenType::Equal) => self.assignment_statement(),
            (TokenType::If, _) => self.if_statement(),
            (TokenType::While, _) => self.while_statement(),
            (TokenType::For, _) => self.for_statement(),
            _ => self.expression_statement(),
        }
    }

    fn for_statement(&mut self) -> LoxResult<AstNode<Stmt>> {
        let start_slice = self.peek().source_slice.clone();
        self.advance();
        let variable = if let Token {
            token_type: TokenType::Identifier,
            literal: Some(variable_name),
            source_slice,
            ..
        } = self.peek()
        {
            AstNode {
                node: Expr::Literal {
                    value: variable_name.clone(),
                },
                source_slice: source_slice.clone(),
            }
        } else {
            return parse_error(
                self.peek().source_slice.clone(),
                "Expect variable name after 'for' keyword.",
            );
        };
        self.advance();
        self.consume(TokenType::In, "Expect 'in' after for variable.")?;
        let iterable = self.expression()?;
        let body = self.statement()?;
        let end_slice = body.source_slice.clone();
        let combined_slice = SourceSlice::from_positions(
            start_slice.source_id,
            start_slice.start_position,
            end_slice.end_position,
        );
        Ok(AstNode::new(
            Stmt::For {
                variable: Box::new(variable),
                iterable: Box::new(iterable),
                body: Box::new(body),
            },
            combined_slice,
        ))
    }

    fn while_statement(&mut self) -> LoxResult<AstNode<Stmt>> {
        let start_slice = self.peek().source_slice.clone();
        self.advance();
        let condition = self.expression()?;
        let body = self.statement()?;
        let end_slice = body.source_slice.clone();
        let combined_slice = SourceSlice::from_positions(
            start_slice.source_id,
            start_slice.start_position,
            end_slice.end_position,
        );
        Ok(AstNode::new(
            Stmt::While {
                condition: Box::new(condition),
                body: Box::new(body),
            },
            combined_slice,
        ))
    }

    fn if_statement(&mut self) -> LoxResult<AstNode<Stmt>> {
        let start_slice = self.peek().source_slice.clone();
        self.advance(); // consume 'if'
        let condition = self.expression()?;
        self.consume(TokenType::Then, "Expect 'then' after if condition.")?;
        let then_branch = self.statement()?;
        let mut elif_branches = Vec::new();
        let mut last_slice = then_branch.source_slice.clone();

        while self.peek().token_type == TokenType::Elif {
            self.advance(); // consume 'elif'
            let elif_condition = self.expression()?;
            self.consume(TokenType::Then, "Expect 'then' after elif condition.")?;
            let elif_branch = self.statement()?;
            last_slice = elif_branch.source_slice.clone();
            elif_branches.push((Box::new(elif_condition), Box::new(elif_branch)));
        }

        let else_branch = if self.peek().token_type == TokenType::Else {
            self.advance(); // consume 'else'
            let else_stmt = self.statement()?;
            last_slice = else_stmt.source_slice.clone();
            Some(Box::new(else_stmt))
        } else {
            None
        };

        let combined_slice = SourceSlice::from_positions(
            start_slice.source_id,
            start_slice.start_position,
            last_slice.end_position,
        );

        Ok(AstNode::new(
            Stmt::If {
                condition: Box::new(condition),
                then_branch: Box::new(then_branch),
                elif_branch: elif_branches,
                else_branch: else_branch,
            },
            combined_slice,
        ))
    }

    fn var_statement(&mut self) -> LoxResult<AstNode<Stmt>> {
        let start_slice = self.peek().source_slice.clone();
        self.advance();
        let name = self.consume(TokenType::Identifier, "Expect variable name.")?;
        let name_lexeme = name.lexeme.clone();
        self.consume(TokenType::Equal, "Expect '=' after variable name.")?;
        let value = self.expression()?;
        let semicolon = self.consume(
            TokenType::Semicolon,
            "Expect ';' after variable declaration.",
        )?;
        let end_slice = semicolon.source_slice.clone();
        let combined_slice = SourceSlice::from_positions(
            start_slice.source_id,
            start_slice.start_position,
            end_slice.end_position,
        );
        Ok(AstNode::new(
            Stmt::Var {
                name: name_lexeme,
                initializer: Some(Box::new(value)),
            },
            combined_slice,
        ))
    }

    fn assignment_statement(&mut self) -> LoxResult<AstNode<Stmt>> {
        let start_slice = self.peek().source_slice.clone();
        let name = self.consume(TokenType::Identifier, "Expect variable name.")?;
        let name_lexeme = name.lexeme.clone();
        self.consume(TokenType::Equal, "Expect '=' after variable name.")?;
        let value = self.expression()?;
        let semicolon = self.consume(
            TokenType::Semicolon,
            "Expect ';' after variable declaration.",
        )?;
        let end_slice = semicolon.source_slice.clone();
        let combined_slice = SourceSlice::from_positions(
            start_slice.source_id,
            start_slice.start_position,
            end_slice.end_position,
        );
        Ok(AstNode::new(
            Stmt::Assign {
                name: name_lexeme,
                value: Box::new(value),
            },
            combined_slice,
        ))
    }

    fn print_statement(&mut self) -> LoxResult<AstNode<Stmt>> {
        // consume the print keyword
        let start_slice = self.peek().source_slice.clone();
        self.advance();
        let expr = self.expression()?;
        let semicolon = self.consume(TokenType::Semicolon, "Expect ';' after value.")?;
        let end_slice = semicolon.source_slice.clone();
        let combined_slice = SourceSlice::from_positions(
            start_slice.source_id,
            start_slice.start_position,
            end_slice.end_position,
        );
        Ok(AstNode::new(
            Stmt::Print {
                expression: Box::new(expr),
            },
            combined_slice,
        ))
    }

    fn block_statement(&mut self) -> LoxResult<AstNode<Stmt>> {
        let start_slice = self.peek().source_slice.clone();
        self.advance();
        let mut statements = Vec::new();
        while self.peek().token_type != TokenType::EndBlock && !self.is_at_end() {
            let stmt = self.statement()?;
            let should_break = matches!(stmt.node, Stmt::Expression { .. });
            statements.push(stmt);
            if should_break {
                break;
            }
        }
        let end_token = self.consume(TokenType::EndBlock, "Expect 'end' after block.")?;
        let end_slice = end_token.source_slice.clone();
        let combined_slice = SourceSlice::from_positions(
            start_slice.source_id,
            start_slice.start_position,
            end_slice.end_position,
        );
        Ok(AstNode::new(
            Stmt::Block {
                statements: Box::new(statements),
            },
            combined_slice,
        ))
    }

    fn expression_statement(&mut self) -> LoxResult<AstNode<Stmt>> {
        let start_slice = self.peek().source_slice.clone();
        let expr = self.expression()?;
        if self.peek().token_type == TokenType::Semicolon {
            let semicolon = self.advance();
            let end_slice = semicolon.source_slice.clone();
            let combined_slice = SourceSlice::from_positions(
                start_slice.source_id,
                start_slice.start_position,
                end_slice.end_position,
            );
            return Ok(AstNode::new(
                Stmt::Stmt {
                    expression: Box::new(expr),
                },
                combined_slice,
            ));
        }
        // Use the expression's source slice for expression statements without semicolon
        let expr_slice = expr.source_slice.clone();
        Ok(AstNode::new(
            Stmt::Expression {
                expression: Box::new(expr),
            },
            expr_slice,
        ))
    }

    fn return_statement(&mut self) -> LoxResult<AstNode<Stmt>> {
        let start_slice = self.peek().source_slice.clone();
        self.advance();
        let expr = self.expression()?;
        let semicolon = self.consume(TokenType::Semicolon, "Expect ';' after return value.")?;
        let end_slice = semicolon.source_slice.clone();
        let combined_slice = SourceSlice::from_positions(
            start_slice.source_id,
            start_slice.start_position,
            end_slice.end_position,
        );
        Ok(AstNode::new(
            Stmt::Return {
                expression: Box::new(expr),
            },
            combined_slice,
        ))
    }

    fn expression(&mut self) -> LoxResult<AstNode<Expr>> {
        self.or_and()
    }

    fn or_and(&mut self) -> LoxResult<AstNode<Expr>> {
        let mut expr = self.equality()?;

        while [TokenType::Or, TokenType::And].contains(&self.peek().token_type) {
            let start_slice = expr.source_slice.clone();
            let operator = self.peek().token_type.clone();
            self.advance();
            let right = self.equality()?;
            let end_slice = right.source_slice.clone();
            let combined_slice = SourceSlice::from_positions(
                start_slice.source_id,
                start_slice.start_position,
                end_slice.end_position,
            );
            expr = AstNode::new(
                Expr::Binary {
                    left: Box::new(expr),
                    operator,
                    right: Box::new(right),
                },
                combined_slice,
            );
        }

        Ok(expr)
    }
    fn equality(&mut self) -> LoxResult<AstNode<Expr>> {
        let mut expr = self.comparison()?;

        while [TokenType::BangEqual, TokenType::EqualEqual].contains(&self.peek().token_type) {
            let start_slice = expr.source_slice.clone();
            let operator = self.peek().token_type.clone();
            self.advance();
            let right = self.comparison()?;
            let end_slice = right.source_slice.clone();
            let combined_slice = SourceSlice::from_positions(
                start_slice.source_id,
                start_slice.start_position,
                end_slice.end_position,
            );
            expr = AstNode::new(
                Expr::Binary {
                    left: Box::new(expr),
                    operator,
                    right: Box::new(right),
                },
                combined_slice,
            );
        }

        Ok(expr)
    }

    fn comparison(&mut self) -> LoxResult<AstNode<Expr>> {
        let mut expr = self.term()?;

        while [
            TokenType::Greater,
            TokenType::GreaterEqual,
            TokenType::Less,
            TokenType::LessEqual,
        ]
        .contains(&self.peek().token_type)
        {
            let start_slice = expr.source_slice.clone();
            let operator = self.peek().token_type.clone();
            self.advance();
            let right = self.term()?;
            let end_slice = right.source_slice.clone();
            let combined_slice = SourceSlice::from_positions(
                start_slice.source_id,
                start_slice.start_position,
                end_slice.end_position,
            );
            expr = AstNode::new(
                Expr::Binary {
                    left: Box::new(expr),
                    operator,
                    right: Box::new(right),
                },
                combined_slice,
            );
        }

        Ok(expr)
    }

    fn term(&mut self) -> LoxResult<AstNode<Expr>> {
        let mut expr = self.factor()?;

        while [TokenType::Plus, TokenType::Minus].contains(&self.peek().token_type) {
            let start_slice = expr.source_slice.clone();
            let operator = self.peek().token_type.clone();
            self.advance();
            let right = self.factor()?;
            let end_slice = right.source_slice.clone();
            let combined_slice = SourceSlice::from_positions(
                start_slice.source_id,
                start_slice.start_position,
                end_slice.end_position,
            );
            expr = AstNode::new(
                Expr::Binary {
                    left: Box::new(expr),
                    operator,
                    right: Box::new(right),
                },
                combined_slice,
            );
        }

        Ok(expr)
    }

    fn factor(&mut self) -> LoxResult<AstNode<Expr>> {
        let mut expr = self.unary()?;

        while [TokenType::Slash, TokenType::Star].contains(&self.peek().token_type) {
            let start_slice = expr.source_slice.clone();
            let operator = self.peek().token_type.clone();
            self.advance();
            let right = self.unary()?;
            let end_slice = right.source_slice.clone();
            let combined_slice = SourceSlice::from_positions(
                start_slice.source_id,
                start_slice.start_position,
                end_slice.end_position,
            );
            expr = AstNode::new(
                Expr::Binary {
                    left: Box::new(expr),
                    operator,
                    right: Box::new(right),
                },
                combined_slice,
            );
        }

        Ok(expr)
    }

    fn unary(&mut self) -> LoxResult<AstNode<Expr>> {
        if [TokenType::Bang, TokenType::Minus].contains(&self.peek().token_type) {
            let operator = self.peek().token_type.clone();
            let source_slice = self.peek().source_slice.clone();
            self.advance();
            let right = self.unary()?;
            return Ok(AstNode::new(
                Expr::Unary {
                    operator,
                    operand: Box::new(right),
                },
                source_slice,
            ));
        }

        self.primary()
    }

    fn primary(&mut self) -> LoxResult<AstNode<Expr>> {
        match self.peek().token_type {
            TokenType::False => {
                let source_slice = self.peek().source_slice.clone();
                self.advance();
                Ok(AstNode::new(
                    Expr::Literal {
                        value: LiteralValue::Boolean(false),
                    },
                    source_slice,
                ))
            }
            TokenType::True => {
                let source_slice = self.peek().source_slice.clone();
                self.advance();
                Ok(AstNode::new(
                    Expr::Literal {
                        value: LiteralValue::Boolean(true),
                    },
                    source_slice,
                ))
            }
            TokenType::Nil => {
                let source_slice = self.peek().source_slice.clone();
                self.advance();
                Ok(AstNode::new(
                    Expr::Literal {
                        value: LiteralValue::Nil,
                    },
                    source_slice,
                ))
            }
            TokenType::Number => {
                let source_slice = self.peek().source_slice.clone();
                let literal = self.peek().literal.clone();
                self.advance();
                if let Some(literal) = literal {
                    Ok(AstNode::new(Expr::Literal { value: literal }, source_slice))
                } else {
                    Err(self
                        .error(self.previous(), "Expected number literal")
                        .into())
                }
            }
            TokenType::String => {
                let source_slice = self.peek().source_slice.clone();
                let literal = self.peek().literal.clone();
                self.advance();
                if let Some(literal) = literal {
                    Ok(AstNode::new(Expr::Literal { value: literal }, source_slice))
                } else {
                    Err(self
                        .error(self.previous(), "Expected string literal")
                        .into())
                }
            }
            TokenType::LeftParen => {
                let source_slice = self.peek().source_slice.clone();
                self.advance();
                let expr = self.expression()?;
                self.consume(TokenType::RightParen, "Expect ')' after expression.")?;
                Ok(AstNode::new(
                    Expr::Grouping {
                        expression: Box::new(expr),
                    },
                    source_slice,
                ))
            }
            TokenType::Identifier => {
                let name = self.peek().lexeme.clone();
                let source_slice = self.peek().source_slice.clone();
                self.advance();
                Ok(AstNode::new(Expr::Variable { name }, source_slice))
            }
            _ => Err(self.error(self.peek(), "Expect expression.").into()),
        }
    }

    fn is_at_end(&self) -> bool {
        if self.current >= self.tokens.len() {
            true
        } else {
            self.tokens[self.current].token_type == TokenType::Eof
        }
    }

    fn advance(&mut self) -> &Token {
        if !self.is_at_end() {
            self.current += 1;
        }
        self.previous()
    }

    fn peek(&self) -> &Token {
        if self.is_at_end() {
            &self.tokens.last().unwrap()
        } else {
            &self.tokens[self.current]
        }
    }

    fn peek_next(&self) -> &Token {
        if self.is_at_end() {
            &self.peek()
        } else {
            &self.tokens[self.current + 1]
        }
    }

    fn previous(&self) -> &Token {
        if self.current == 0 {
            &self.tokens[0]
        } else {
            &self.tokens[self.current - 1]
        }
    }

    fn consume(&mut self, token_type: TokenType, message: &str) -> LoxResult<&Token> {
        if self.peek().token_type == token_type {
            self.advance();
            Ok(self.previous())
        } else {
            Err(self.error(self.peek(), message).into())
        }
    }

    fn error(&self, token: &Token, message: &str) -> LoxError {
        LoxError::ParseError {
            source_slice: token.source_slice.clone(),
            message: message.to_string(),
        }
    }

    fn synchronize(&mut self) {
        self.advance();

        while !self.is_at_end() {
            if self.previous().token_type == TokenType::Semicolon {
                return;
            }

            match self.peek().token_type {
                TokenType::Class
                | TokenType::Fun
                | TokenType::Var
                | TokenType::For
                | TokenType::If
                | TokenType::While
                | TokenType::Print
                | TokenType::Return => return,
                _ => {}
            }

            self.advance();
        }
    }
}