aster/aster-core/src/parser/parser.rs

use crate::ast::*;
use crate::lexer::{Token, TokenKind};
use crate::error::RuntimeError;

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) -> (Vec<Stmt>, Vec<RuntimeError>) {
        let mut statements = Vec::new();
        let mut errors = Vec::new();

        while !self.is_at_end() {
            match self.declaration() {
                Ok(stmt) => statements.push(stmt),
                Err(err) => {
                    errors.push(err);
                    self.synchronize();
                }
            }
        }

        (statements, errors)
    }

    fn declaration(&mut self) -> Result<Stmt, RuntimeError> {
        if self.match_kind(&[TokenKind::Let]) {
            self.let_declaration(true)
        } else if self.match_kind(&[TokenKind::Const]) {
            self.let_declaration(false)
        } else if self.match_kind(&[TokenKind::Fn]) {
            self.fn_declaration()
        } else {
            self.statement()
        }
    }

    fn statement(&mut self) -> Result<Stmt, RuntimeError> {
        if self.match_kind(&[TokenKind::If]) {
            self.if_statement()
        } else if self.match_kind(&[TokenKind::While]) {
            self.while_statement()
        } else if self.match_kind(&[TokenKind::For]) {
            self.for_statement()
        } else if self.match_kind(&[TokenKind::Break]) {
            self.break_statement()
        } else if self.match_kind(&[TokenKind::Continue]) {
            self.continue_statement()
        } else if self.match_kind(&[TokenKind::Return]) {
            self.return_statement()
        } else if self.match_kind(&[TokenKind::Try]) {
            self.try_statement()
        } else if self.match_kind(&[TokenKind::Throw]) {
            self.throw_statement()
        } else if self.match_kind(&[TokenKind::LeftBrace]) {
            Ok(Stmt::Block(self.block()?))
        } else {
            self.expression_statement()
        }
    }

    fn let_declaration(&mut self, mutable: bool) -> Result<Stmt, RuntimeError> {
        let name = self.consume_ident("Expected variable name.")?;
        self.consume(TokenKind::Equal, "Expected '=' after variable name.")?;
        let initializer = self.expression()?;
        self.match_kind(&[TokenKind::Semicolon]); // 分号可选

        Ok(Stmt::Let { name, initializer, mutable })
    }

    fn fn_declaration(&mut self) -> Result<Stmt, RuntimeError> {
        let name = self.consume_ident("Expected function name.")?;
        self.consume(TokenKind::LeftParen, "Expected '(' after function name.")?;

        let mut params = Vec::new();
        if !self.check(&TokenKind::RightParen) {
            loop {
                params.push(self.consume_ident("Expected parameter name.")?);
                if !self.match_kind(&[TokenKind::Comma]) {
                    break;
                }
            }
        }

        self.consume(TokenKind::RightParen, "Expected ')' after parameters.")?;
        self.consume(TokenKind::LeftBrace, "Expected '{' before function body.")?;
        let body = self.block()?;

        Ok(Stmt::Function { name, params, body })
    }

    fn lambda_expr(&mut self) -> Result<Expr, RuntimeError> {
        self.consume(TokenKind::LeftParen, "Expected '(' after 'fn'.")?;
        let mut params = Vec::new();
        if !self.check(&TokenKind::RightParen) {
            loop {
                params.push(self.consume_ident("Expected parameter name.")?);
                if !self.match_kind(&[TokenKind::Comma]) {
                    break;
                }
            }
        }
        self.consume(TokenKind::RightParen, "Expected ')' after parameters.")?;
        self.consume(TokenKind::LeftBrace, "Expected '{' before function body.")?;
        let body = self.block()?;
        Ok(Expr::Lambda { params, body })
    }

    fn if_statement(&mut self) -> Result<Stmt, RuntimeError> {
        self.consume(TokenKind::LeftParen, "Expected '(' after 'if'.")?;
        let condition = self.expression()?;
        self.consume(TokenKind::RightParen, "Expected ')' after condition.")?;

        let then_branch = Box::new(self.statement()?);
        let else_branch = if self.match_kind(&[TokenKind::Else]) {
            Some(Box::new(self.statement()?))
        } else {
            None
        };

        Ok(Stmt::If { condition, then_branch, else_branch })
    }

    fn while_statement(&mut self) -> Result<Stmt, RuntimeError> {
        self.consume(TokenKind::LeftParen, "Expected '(' after 'while'.")?;
        let condition = self.expression()?;
        self.consume(TokenKind::RightParen, "Expected ')' after condition.")?;

        let body = Box::new(self.statement()?);
        Ok(Stmt::While { condition, body })
    }

    fn for_statement(&mut self) -> Result<Stmt, RuntimeError> {
        self.consume(TokenKind::LeftParen, "Expected '(' after 'for'.")?;

        // Detect for-in: for (Identifier in expr) body
        if let TokenKind::Identifier(var_name) = &self.peek().kind {
            let var_name = var_name.clone();
            if matches!(self.peek_next().kind, TokenKind::In) {
                self.advance(); // consume identifier
                self.advance(); // consume 'in'
                let iterable = self.expression()?;
                self.consume(TokenKind::RightParen, "Expected ')' after for-in expression.")?;
                let body = Box::new(self.statement()?);
                return Ok(Stmt::ForIn { var_name, iterable, body });
            }
        }

        let initializer = if self.match_kind(&[TokenKind::Semicolon]) {
            None
        } else if self.match_kind(&[TokenKind::Let]) {
            Some(self.let_declaration(true)?)
        } else {
            Some(self.expression_statement()?)
        };

        let condition = if !self.check(&TokenKind::Semicolon) {
            Some(self.expression()?)
        } else {
            None
        };

        self.consume(TokenKind::Semicolon, "Expected ';' after loop condition.")?;
        let step = if !self.check(&TokenKind::RightParen) {
            Some(self.expression()?)
        } else {
            None
        };

        self.consume(TokenKind::RightParen, "Expected ')' after for clauses.")?;
        let body = Box::new(self.statement()?);

        Ok(Stmt::For {
            initializer: initializer.map(Box::new),
            condition,
            step,
            body
        })
    }

    fn break_statement(&mut self) -> Result<Stmt, RuntimeError> {
        self.match_kind(&[TokenKind::Semicolon]); // 分号可选
        Ok(Stmt::Break)
    }

    fn continue_statement(&mut self) -> Result<Stmt, RuntimeError> {
        self.match_kind(&[TokenKind::Semicolon]); // 分号可选
        Ok(Stmt::Continue)
    }

    fn return_statement(&mut self) -> Result<Stmt, RuntimeError> {
        let value = if !self.check(&TokenKind::Semicolon) && !self.check(&TokenKind::RightBrace) {
            Some(self.expression()?)
        } else {
            None
        };

        self.match_kind(&[TokenKind::Semicolon]); // 分号可选
        Ok(Stmt::Return(value))
    }

    fn try_statement(&mut self) -> Result<Stmt, RuntimeError> {
        self.consume(TokenKind::LeftBrace, "Expected '{' after 'try'.")?;
        let body = self.block()?;

        let mut catch_var = None;
        let mut catch_body = None;
        if self.match_kind(&[TokenKind::Catch]) {
            self.consume(TokenKind::LeftParen, "Expected '(' after 'catch'.")?;
            catch_var = Some(self.consume_ident("Expected exception variable.")?);
            self.consume(TokenKind::RightParen, "Expected ')' after catch variable.")?;
            self.consume(TokenKind::LeftBrace, "Expected '{' before catch body.")?;
            catch_body = Some(self.block()?);
        }

        let mut finally_body = None;
        if self.match_kind(&[TokenKind::Finally]) {
            self.consume(TokenKind::LeftBrace, "Expected '{' after 'finally'.")?;
            finally_body = Some(self.block()?);
        }

        if catch_var.is_none() && finally_body.is_none() {
            return Err(RuntimeError::parse("Expected 'catch' or 'finally' after 'try'.", self.peek().clone()));
        }

        Ok(Stmt::Try { body, catch_var, catch_body, finally_body })
    }

    fn throw_statement(&mut self) -> Result<Stmt, RuntimeError> {
        let expr = self.expression()?;
        self.match_kind(&[TokenKind::Semicolon]); // 分号可选
        Ok(Stmt::Throw(expr))
    }

    fn block(&mut self) -> Result<Vec<Stmt>, RuntimeError> {
        let mut statements = Vec::new();

        while !self.check(&TokenKind::RightBrace) && !self.is_at_end() {
            statements.push(self.declaration()?);
        }

        self.consume(TokenKind::RightBrace, "Expected '}' after block.")?;
        Ok(statements)
    }

    fn expression_statement(&mut self) -> Result<Stmt, RuntimeError> {
        let expr = self.expression()?;
        self.match_kind(&[TokenKind::Semicolon]); // 分号可选
        Ok(Stmt::ExprStmt(expr))
    }

    fn expression(&mut self) -> Result<Expr, RuntimeError> {
        self.assignment()
    }

    fn ternary(&mut self) -> Result<Expr, RuntimeError> {
        let expr = self.logical_or()?;
        if self.match_kind(&[TokenKind::Question]) {
            let then_branch = self.assignment()?;
            self.consume(TokenKind::Colon, "Expected ':' after then branch of ternary.")?;
            let else_branch = self.assignment()?;
            Ok(Expr::Ternary {
                condition: Box::new(expr),
                then_branch: Box::new(then_branch),
                else_branch: Box::new(else_branch),
            })
        } else {
            Ok(expr)
        }
    }

    fn assignment(&mut self) -> Result<Expr, RuntimeError> {
        let expr = self.ternary()?;

        let op = if self.match_kind(&[TokenKind::Equal]) {
            Some(AssignOp::Equal)
        } else if self.match_kind(&[TokenKind::PlusEqual]) {
            Some(AssignOp::PlusEqual)
        } else if self.match_kind(&[TokenKind::MinusEqual]) {
            Some(AssignOp::MinusEqual)
        } else if self.match_kind(&[TokenKind::StarEqual]) {
            Some(AssignOp::StarEqual)
        } else if self.match_kind(&[TokenKind::SlashEqual]) {
            Some(AssignOp::SlashEqual)
        } else if self.match_kind(&[TokenKind::PercentEqual]) {
            Some(AssignOp::PercentEqual)
        } else {
            None
        };

        if let Some(op) = op {
            let value = self.assignment()?;
            return match expr {
                Expr::Variable(name) => Ok(Expr::Assign { name, op, value: Box::new(value) }),
                Expr::Get { object, name } => Ok(Expr::Set { object, name, op, value: Box::new(value) }),
                Expr::IndexGet { array, index } => Ok(Expr::IndexSet { array, index, op, value: Box::new(value) }),
                _ => Err(RuntimeError::parse("Invalid assignment target.", self.previous().clone())),
            };
        }

        Ok(expr)
    }

    fn logical_or(&mut self) -> Result<Expr, RuntimeError> {
        let mut expr = self.logical_and()?;

        while self.match_kind(&[TokenKind::OrOr]) {
            let right = self.logical_and()?;
            expr = Expr::Logical {
                left: Box::new(expr),
                op: LogicalOp::Or,
                right: Box::new(right),
            };
        }

        Ok(expr)
    }

    fn logical_and(&mut self) -> Result<Expr, RuntimeError> {
        let mut expr = self.equality()?;

        while self.match_kind(&[TokenKind::AndAnd]) {
            let right = self.equality()?;
            expr = Expr::Logical {
                left: Box::new(expr),
                op: LogicalOp::And,
                right: Box::new(right),
            };
        }

        Ok(expr)
    }

    fn equality(&mut self) -> Result<Expr, RuntimeError> {
        let mut expr = self.comparison()?;

        while self.match_kind(&[TokenKind::EqualEqual, TokenKind::BangEqual]) {
            let op = match self.previous().kind {
                TokenKind::EqualEqual => BinaryOp::Equal,
                TokenKind::BangEqual => BinaryOp::NotEqual,
                _ => unreachable!(),
            };
            let right = self.comparison()?;
            expr = Expr::Binary {
                left: Box::new(expr),
                op,
                right: Box::new(right),
            };
        }

        Ok(expr)
    }

    fn comparison(&mut self) -> Result<Expr, RuntimeError> {
        let mut expr = self.term()?;

        while self.match_kind(&[TokenKind::Greater, TokenKind::GreaterEqual, TokenKind::Less, TokenKind::LessEqual]) {
            let op = match self.previous().kind {
                TokenKind::Greater => BinaryOp::Greater,
                TokenKind::GreaterEqual => BinaryOp::GreaterEqual,
                TokenKind::Less => BinaryOp::Less,
                TokenKind::LessEqual => BinaryOp::LessEqual,
                _ => unreachable!(),
            };
            let right = self.term()?;
            expr = Expr::Binary {
                left: Box::new(expr),
                op,
                right: Box::new(right),
            };
        }

        Ok(expr)
    }

    fn term(&mut self) -> Result<Expr, RuntimeError> {
        let mut expr = self.factor()?;

        while self.match_kind(&[TokenKind::Plus, TokenKind::Minus]) {
            let op = match self.previous().kind {
                TokenKind::Plus => BinaryOp::Add,
                TokenKind::Minus => BinaryOp::Sub,
                _ => unreachable!(),
            };
            let right = self.factor()?;
            expr = Expr::Binary {
                left: Box::new(expr),
                op,
                right: Box::new(right),
            };
        }

        Ok(expr)
    }

    fn factor(&mut self) -> Result<Expr, RuntimeError> {
        let mut expr = self.unary()?;

        while self.match_kind(&[TokenKind::Star, TokenKind::Slash, TokenKind::Percent]) {
            let op = match self.previous().kind {
                TokenKind::Star => BinaryOp::Mul,
                TokenKind::Slash => BinaryOp::Div,
                TokenKind::Percent => BinaryOp::Mod,
                _ => unreachable!(),
            };
            let right = self.unary()?;
            expr = Expr::Binary {
                left: Box::new(expr),
                op,
                right: Box::new(right),
            };
        }

        Ok(expr)
    }

    fn unary(&mut self) -> Result<Expr, RuntimeError> {
        if self.match_kind(&[TokenKind::Bang, TokenKind::Minus]) {
            let op = match self.previous().kind {
                TokenKind::Bang => UnaryOp::Not,
                TokenKind::Minus => UnaryOp::Negate,
                _ => unreachable!(),
            };
            let right = self.unary()?;
            return Ok(Expr::Unary {
                op,
                right: Box::new(right),
            });
        }

        self.call()
    }

    fn call(&mut self) -> Result<Expr, RuntimeError> {
        let mut expr = self.primary()?;

        loop {
            if self.match_kind(&[TokenKind::LeftParen]) {
                let mut arguments = Vec::new();
                if !self.check(&TokenKind::RightParen) {
                    loop {
                        arguments.push(self.expression()?);
                        if !self.match_kind(&[TokenKind::Comma]) {
                            break;
                        }
                    }
                }
                self.consume(TokenKind::RightParen, "Expected ')' after arguments.")?;
                expr = Expr::Call {
                    callee: Box::new(expr),
                    arguments,
                };
            } else if self.match_kind(&[TokenKind::Dot]) {
                let name = self.consume_ident("Expected property name after '.'.")?;
                expr = Expr::Get {
                    object: Box::new(expr),
                    name,
                };
            } else if self.match_kind(&[TokenKind::LeftBracket]) {
                let index = self.expression()?;
                self.consume(TokenKind::RightBracket, "Expected ']' after index.")?;
                expr = Expr::IndexGet {
                    array: Box::new(expr),
                    index: Box::new(index),
                };
            } else {
                break;
            }
        }

        Ok(expr)
    }

    fn primary(&mut self) -> Result<Expr, RuntimeError> {
        // 布尔值
        if self.match_kind(&[TokenKind::True]) {
            return Ok(Expr::Literal(Literal::Bool(true)));
        }
        if self.match_kind(&[TokenKind::False]) {
            return Ok(Expr::Literal(Literal::Bool(false)));
        }
        if self.match_kind(&[TokenKind::Nil]) {
            return Ok(Expr::Literal(Literal::Nil));
        }

        // 数字
        if let TokenKind::Number(n) = self.peek().kind {
            self.advance();
            return Ok(Expr::Literal(Literal::Number(n)));
        }

        // 字符串
        if let TokenKind::String(s) = &self.peek().kind {
            let s = s.clone();
            self.advance();
            return Ok(Expr::Literal(Literal::String(s)));
        }

        // 标识符（变量）
        if let TokenKind::Identifier(name) = &self.peek().kind {
            let name = name.clone();
            self.advance();
            return Ok(Expr::Variable(name));
        }

        // 匿名函数（闭包）: fn(params) { body }
        if self.match_kind(&[TokenKind::Fn]) {
            return self.lambda_expr();
        }

        // 括号表达式
        if self.match_kind(&[TokenKind::LeftParen]) {
            let expr = self.expression()?;
            self.consume(TokenKind::RightParen, "Expected ')' after expression.")?;
            return Ok(expr);
        }

        if self.match_kind(&[TokenKind::LeftBrace]) {
            return self.object_literal();
        }

        if self.match_kind(&[TokenKind::LeftBracket]) {
            return self.array_literal();
        }

        Err(RuntimeError::parse("Expected expression.", self.peek().clone()))
    }

    fn match_kind(&mut self, kinds: &[TokenKind]) -> bool {
        for kind in kinds {
            if self.check(kind) {
                self.advance();
                return true;
            }
        }
        false
    }

    fn consume(&mut self, kind: TokenKind, msg: &str) -> Result<(), RuntimeError> {
        if self.check(&kind) {
            self.advance();
            Ok(())
        } else {
            Err(RuntimeError::parse(msg, self.peek().clone()))
        }
    }

    fn consume_ident(&mut self, msg: &str) -> Result<String, RuntimeError> {
        match &self.peek().kind {
            TokenKind::Identifier(name) => {
                let name = name.clone();
                self.advance();
                Ok(name)
            }
            _ => Err(RuntimeError::parse(msg, self.peek().clone())),
        }
    }

    fn check(&self, kind: &TokenKind) -> bool {
        if self.is_at_end() {
            return false;
        }
        std::mem::discriminant(&self.peek().kind)
            == std::mem::discriminant(kind)
    }

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

    fn is_at_end(&self) -> bool {
        matches!(self.peek().kind, TokenKind::EOF)
    }

    fn peek(&self) -> &Token {
        &self.tokens[self.current]
    }

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

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

    fn object_literal(&mut self) -> Result<Expr, RuntimeError> {
        let mut properties = Vec::new();
        if !self.check(&TokenKind::RightBrace) {
            loop {
                let name = self.consume_ident("Expected property name in object literal")?;
                self.consume(TokenKind::Colon, "Expected ':' after property name in object literal")?;
                let value = self.expression()?;
                properties.push((name, value));
                if !self.match_kind(&[TokenKind::Comma]) {
                    break;
                }
                // Allow trailing comma
                if self.check(&TokenKind::RightBrace) {
                    break;
                }
            }
        }
        self.consume(TokenKind::RightBrace, "Expected '}' after object literal")?;
        Ok(Expr::ObjectLiteral { properties })
    }

    fn array_literal(&mut self) -> Result<Expr, RuntimeError> {
        let mut elements = Vec::new();
        if !self.check(&TokenKind::RightBracket) {
            loop {
                elements.push(self.expression()?);
                if !self.match_kind(&[TokenKind::Comma]) {
                    break;
                }
                // Allow trailing comma
                if self.check(&TokenKind::RightBracket) {
                    break;
                }
            }
        }
        self.consume(TokenKind::RightBracket, "Expected ']' after array literal")?;
        Ok(Expr::ArrayLiteral { elements })
    }

    fn synchronize(&mut self) {
        // Always consume at least one token to guarantee forward progress,
        // preventing infinite loops when previous() is already a semicolon.
        if !self.is_at_end() {
            self.advance();
        }

        while !self.is_at_end() {
            if self.current > 0 && matches!(self.previous().kind, TokenKind::Semicolon) {
                return;
            }

            match self.peek().kind {
                TokenKind::Let
                | TokenKind::Const
                | TokenKind::Fn
                | TokenKind::If
                | TokenKind::While
                | TokenKind::For
                | TokenKind::Return
                | TokenKind::Break
                | TokenKind::Continue
                | TokenKind::Try
                | TokenKind::Throw => return,
                _ => {
                    self.advance();
                }
            }
        }
    }
}

#[cfg(test)]
#[path = "tests.rs"]
mod tests;