RustAst/src/ast/parser.rs

use crate::ast::diagnostics::Diagnostics;
use crate::ast::lexer::{Lexer, Token, TokenKind};
use crate::ast::nodes::{Symbol, SyntaxKind, SyntaxNode};
use crate::ast::types::{Identity, Keyword, NodeIdentity, SourceLocation, Value};
use std::rc::Rc;

pub struct Parser<'a> {
    lexer: Lexer<'a>,
    current_token: Token,
    pub diagnostics: Diagnostics,
}

impl<'a> Parser<'a> {
    pub fn new(input: &'a str) -> Self {
        let mut lexer = Lexer::new(input);
        let mut diagnostics = Diagnostics::new();
        let current_token = match lexer.next_token() {
            Ok(t) => t,
            Err(e) => {
                diagnostics.push_error(e, None);
                Token {
                    kind: TokenKind::EOF,
                    location: SourceLocation { line: 1, col: 1 },
                }
            }
        };
        Self {
            lexer,
            current_token,
            diagnostics,
        }
    }

    fn advance(&mut self) -> Token {
        let next = match self.lexer.next_token() {
            Ok(t) => t,
            Err(e) => {
                self.diagnostics
                    .push_error(e, Some(NodeIdentity::new(self.current_token.location)));
                Token {
                    kind: TokenKind::EOF,
                    location: self.current_token.location,
                }
            }
        };
        std::mem::replace(&mut self.current_token, next)
    }

    fn peek(&self) -> &TokenKind {
        &self.current_token.kind
    }

    pub fn parse_expression(&mut self) -> SyntaxNode {
        let token_loc = self.current_token.location;
        let identity = NodeIdentity::new(token_loc);

        match self.peek() {
            TokenKind::LeftParen => self.parse_list(),
            TokenKind::LeftBracket => self.parse_vector_literal(),
            TokenKind::LeftBrace => self.parse_record_literal(),
            TokenKind::Quote => {
                self.advance(); // consume '
                let expr = self.parse_expression();
                SyntaxNode {
                    identity: identity.clone(),
                    kind: SyntaxKind::Call {
                        callee: Rc::new(self.make_id_node("quote", identity.clone())),
                        args: Rc::new(SyntaxNode {
                            identity,
                            kind: SyntaxKind::Tuple {
                                elements: vec![Rc::new(expr)],
                            },
                            ty: (),
                        }),
                    },
                    ty: (),
                }
            }
            TokenKind::Backtick => {
                self.advance(); // consume `
                let expr = self.parse_expression();
                SyntaxNode {
                    identity,
                    kind: SyntaxKind::Template(Rc::new(expr)),
                    ty: (),
                }
            }
            TokenKind::Tilde => {
                self.advance(); // consume ~
                if *self.peek() == TokenKind::At {
                    self.advance(); // consume @
                    let expr = self.parse_expression();
                    SyntaxNode {
                        identity,
                        kind: SyntaxKind::Splice(Rc::new(expr)),
                        ty: (),
                    }
                } else {
                    let expr = self.parse_expression();
                    SyntaxNode {
                        identity,
                        kind: SyntaxKind::Placeholder(Rc::new(expr)),
                        ty: (),
                    }
                }
            }
            _ => self.parse_atom(),
        }
    }

    pub fn at_eof(&self) -> bool {
        matches!(self.current_token.kind, TokenKind::EOF)
    }

    fn synchronize(&mut self) {
        while !self.at_eof() {
            match self.peek() {
                TokenKind::RightParen | TokenKind::RightBracket | TokenKind::RightBrace => {
                    self.advance();
                    return;
                }
                _ => {
                    self.advance();
                }
            }
        }
    }

    fn parse_atom(&mut self) -> SyntaxNode {
        let token = self.advance();
        let identity = NodeIdentity::new(token.location);

        let kind = match token.kind {
            TokenKind::Integer(n) => SyntaxKind::Constant(Value::Int(n)),
            TokenKind::Float(n) => SyntaxKind::Constant(Value::Float(n)),
            TokenKind::String(s) => SyntaxKind::Constant(Value::Text(s)),
            TokenKind::Keyword(k) => SyntaxKind::Constant(Value::Keyword(Keyword::intern(&k))),
            TokenKind::Identifier(id) => match id.as_ref() {
                "..." => SyntaxKind::Nop,
                s if s.starts_with('.') && s.len() > 1 => {
                    SyntaxKind::FieldAccessor(Keyword::intern(&s[1..]))
                }
                _ => SyntaxKind::Identifier {
                    symbol: id.into(),
                    binding: (),
                },
            },
            TokenKind::EOF => SyntaxKind::Error, // Error already logged by advance
            _ => {
                self.diagnostics.push_error(
                    format!("Unexpected token in atom: {:?}", token.kind),
                    Some(identity.clone()),
                );
                SyntaxKind::Error
            }
        };

        SyntaxNode {
            identity,
            kind,
            ty: (),
        }
    }

    fn parse_list(&mut self) -> SyntaxNode {
        let start_loc = self.advance().location; // consume '('
        let identity = NodeIdentity::new(start_loc);

        if *self.peek() == TokenKind::RightParen {
            self.diagnostics.push_error(
                "Empty list () is not a valid expression",
                Some(identity.clone()),
            );
            self.advance(); // consume )
            return SyntaxNode {
                identity,
                kind: SyntaxKind::Error,
                ty: (),
            };
        }

        let head = self.parse_expression();

        let node = if let SyntaxKind::Identifier { ref symbol, .. } = head.kind {
            match symbol.name.as_ref() {
                "if" => self.parse_if(identity),
                "fn" => self.parse_fn(identity),
                "pipe" => self.parse_pipe(identity),
                "again" => self.parse_again(identity),
                "def" => self.parse_def(identity),
                "assign" => self.parse_assign(identity),
                "do" => self.parse_do(identity),
                "macro" => self.parse_macro_decl(identity),
                _ => self.parse_call(head, identity),
            }
        } else {
            self.parse_call(head, identity)
        };

        self.expect(TokenKind::RightParen);
        node
    }

    fn parse_again(&mut self, identity: Identity) -> SyntaxNode {
        let mut elements = Vec::new();

        while *self.peek() != TokenKind::RightParen && *self.peek() != TokenKind::EOF {
            elements.push(Rc::new(self.parse_expression()));
        }

        let args_node = SyntaxNode {
            identity: identity.clone(),
            kind: SyntaxKind::Tuple { elements },
            ty: (),
        };

        SyntaxNode {
            identity,
            kind: SyntaxKind::Again {
                args: Rc::new(args_node),
            },
            ty: (),
        }
    }

    fn parse_if(&mut self, identity: Identity) -> SyntaxNode {
        let cond = Rc::new(self.parse_expression());
        let then_br = Rc::new(self.parse_expression());
        let mut else_br = None;

        if *self.peek() != TokenKind::RightParen {
            else_br = Some(Rc::new(self.parse_expression()));
        }

        SyntaxNode {
            identity,
            kind: SyntaxKind::If {
                cond,
                then_br,
                else_br,
            },
            ty: (),
        }
    }

    fn parse_def(&mut self, identity: Identity) -> SyntaxNode {
        let pattern = Rc::new(self.parse_pattern());
        let value = Rc::new(self.parse_expression());

        SyntaxNode {
            identity,
            kind: SyntaxKind::Def {
                pattern,
                value,
                info: (),
            },
            ty: (),
        }
    }

    fn parse_assign(&mut self, identity: Identity) -> SyntaxNode {
        // (assign target value)
        let target = Rc::new(self.parse_expression());
        let value = Rc::new(self.parse_expression());

        SyntaxNode {
            identity,
            kind: SyntaxKind::Assign {
                target,
                value,
                info: (),
            },
            ty: (),
        }
    }

    fn parse_do(&mut self, identity: Identity) -> SyntaxNode {
        let mut exprs = Vec::new();
        while *self.peek() != TokenKind::RightParen && *self.peek() != TokenKind::EOF {
            exprs.push(Rc::new(self.parse_expression()));
        }
        SyntaxNode {
            identity,
            kind: SyntaxKind::Block { exprs },
            ty: (),
        }
    }

    fn parse_pipe(&mut self, identity: Identity) -> SyntaxNode {
        let inputs_node = self.parse_expression();
        let inputs = match inputs_node.kind {
            SyntaxKind::Tuple { elements } => elements,
            _ => vec![Rc::new(inputs_node)],
        };
        let lambda = Rc::new(self.parse_expression());

        SyntaxNode {
            identity,
            kind: SyntaxKind::Pipe { inputs, lambda },
            ty: (),
        }
    }

    fn parse_fn(&mut self, identity: Identity) -> SyntaxNode {
        let params = Rc::new(self.parse_param_vector());
        let body = self.parse_expression();

        SyntaxNode {
            identity,
            kind: SyntaxKind::Lambda {
                params,
                body: Rc::new(body),
                info: (),
            },
            ty: (),
        }
    }

    fn parse_macro_decl(&mut self, identity: Identity) -> SyntaxNode {
        let name_node = self.parse_expression();
        let name = match name_node.kind {
            SyntaxKind::Identifier { symbol, .. } => symbol,
            _ => {
                self.diagnostics.push_error(
                    "Expected identifier for macro name",
                    Some(name_node.identity.clone()),
                );
                Symbol::from("error")
            }
        };
        let params = Rc::new(self.parse_param_vector());
        let body = self.parse_expression();
        SyntaxNode {
            identity,
            kind: SyntaxKind::MacroDecl {
                name,
                params,
                body: Rc::new(body),
            },
            ty: (),
        }
    }

    fn parse_param_vector(&mut self) -> SyntaxNode {
        if *self.peek() != TokenKind::LeftBracket {
            self.diagnostics.push_error(
                format!(
                    "Expected parameter vector [...] for fn, found {:?}",
                    self.peek()
                ),
                Some(NodeIdentity::new(self.current_token.location)),
            );
            return SyntaxNode {
                identity: NodeIdentity::new(self.current_token.location),
                kind: SyntaxKind::Error,
                ty: (),
            };
        }
        self.parse_pattern()
    }

    fn parse_pattern(&mut self) -> SyntaxNode {
        let next = self.peek();
        match next {
            TokenKind::Identifier(_) => {
                let token = self.advance();
                let sym: Symbol = match token.kind {
                    TokenKind::Identifier(s) => s.into(),
                    _ => unreachable!(),
                };
                SyntaxNode {
                    identity: NodeIdentity::new(token.location),
                    kind: SyntaxKind::Identifier {
                        symbol: sym,
                        binding: (),
                    },
                    ty: (),
                }
            }
            TokenKind::LeftBracket => {
                let token = self.advance();
                let identity = NodeIdentity::new(token.location);

                let mut elements = Vec::new();
                while *self.peek() != TokenKind::RightBracket && *self.peek() != TokenKind::EOF {
                    elements.push(Rc::new(self.parse_pattern()));
                }
                self.expect(TokenKind::RightBracket);

                SyntaxNode {
                    identity,
                    kind: SyntaxKind::Tuple { elements },
                    ty: (),
                }
            }
            TokenKind::Tilde => {
                let token = self.advance(); // consume ~
                if *self.peek() == TokenKind::At {
                    self.advance(); // consume @
                    let expr = self.parse_expression();
                    SyntaxNode {
                        identity: NodeIdentity::new(token.location),
                        kind: SyntaxKind::Splice(Rc::new(expr)),
                        ty: (),
                    }
                } else {
                    let expr = self.parse_expression();
                    SyntaxNode {
                        identity: NodeIdentity::new(token.location),
                        kind: SyntaxKind::Placeholder(Rc::new(expr)),
                        ty: (),
                    }
                }
            }
            _ => {
                self.diagnostics.push_error(
                    format!(
                        "Expected identifier or pattern vector [...] for definition, found {:?}",
                        next
                    ),
                    Some(NodeIdentity::new(self.current_token.location)),
                );
                self.advance();
                SyntaxNode {
                    identity: NodeIdentity::new(self.current_token.location),
                    kind: SyntaxKind::Error,
                    ty: (),
                }
            }
        }
    }

    fn parse_call(&mut self, callee: SyntaxNode, identity: Identity) -> SyntaxNode {
        let mut elements = Vec::new();

        while *self.peek() != TokenKind::RightParen && *self.peek() != TokenKind::EOF {
            elements.push(Rc::new(self.parse_expression()));
        }

        // The arguments are wrapped in a Tuple node, reusing the call's identity/location.
        let args_node = SyntaxNode {
            identity: identity.clone(),
            kind: SyntaxKind::Tuple { elements },
            ty: (),
        };

        SyntaxNode {
            identity,
            kind: SyntaxKind::Call {
                callee: Rc::new(callee),
                args: Rc::new(args_node),
            },
            ty: (),
        }
    }

    fn parse_vector_literal(&mut self) -> SyntaxNode {
        let token = self.advance();
        let mut elements = Vec::new();

        while *self.peek() != TokenKind::RightBracket && *self.peek() != TokenKind::EOF {
            let expr = self.parse_expression();
            elements.push(Rc::new(expr));
        }

        self.expect(TokenKind::RightBracket);

        SyntaxNode {
            identity: NodeIdentity::new(token.location),
            kind: SyntaxKind::Tuple { elements },
            ty: (),
        }
    }

    fn parse_record_literal(&mut self) -> SyntaxNode {
        let token = self.advance();
        let mut fields = Vec::new();

        while *self.peek() != TokenKind::RightBrace && *self.peek() != TokenKind::EOF {
            let key_node = self.parse_expression();
            // We check for keyword kind here (syntactically) to avoid ambiguity, but
            // strictly we could allow any expression and check at runtime.
            // Delphi enforces keywords. We can do minimal check here.
            match &key_node.kind {
                SyntaxKind::Constant(Value::Keyword(_)) => {}
                _ => {
                    self.diagnostics.push_error(
                        "Record keys must be keywords (syntactically)",
                        Some(key_node.identity.clone()),
                    );
                }
            }

            if *self.peek() == TokenKind::RightBrace || *self.peek() == TokenKind::EOF {
                self.diagnostics.push_error(
                    "Record literal must have even number of forms",
                    Some(NodeIdentity::new(self.current_token.location)),
                );
                break;
            }
            let val_node = self.parse_expression();

            fields.push((Rc::new(key_node), Rc::new(val_node)));
        }
        self.expect(TokenKind::RightBrace);

        SyntaxNode {
            identity: NodeIdentity::new(token.location),
            kind: SyntaxKind::Record {
                fields,
                layout: (),
            },
            ty: (),
        }
    }

    fn expect(&mut self, kind: TokenKind) -> Token {
        if self.peek() == &kind {
            self.advance()
        } else {
            self.diagnostics.push_error(
                format!("Expected {:?}, but found {:?}", kind, self.peek()),
                Some(NodeIdentity::new(self.current_token.location)),
            );
            // Recovery: skip until we find what we expected or a synchronization point
            self.synchronize();
            Token {
                kind,
                location: self.current_token.location,
            }
        }
    }

    fn make_id_node(&self, name: &str, identity: Identity) -> SyntaxNode {
        SyntaxNode {
            identity,
            kind: SyntaxKind::Identifier {
                symbol: Symbol::from(name),
                binding: (),
            },
            ty: (),
        }
    }
}