RustAst/src/ast/parser.rs

use crate::ast::lexer::{Lexer, Token, TokenKind};
use crate::ast::nodes::{Node, Symbol, UntypedKind};
use crate::ast::types::{Identity, Keyword, NodeIdentity, Value};
use std::rc::Rc;

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

impl<'a> Parser<'a> {
    pub fn new(input: &'a str) -> Result<Self, String> {
        let mut lexer = Lexer::new(input);
        let current_token = lexer.next_token()?;
        Ok(Self {
            lexer,
            current_token,
        })
    }

    fn advance(&mut self) -> Result<Token, String> {
        let prev = std::mem::replace(&mut self.current_token, self.lexer.next_token()?);
        Ok(prev)
    }

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

    pub fn parse_expression(&mut self) -> Result<Node<UntypedKind>, String> {
        let token_loc = self.current_token.location;
        let identity = Rc::new(NodeIdentity {
            location: 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()?;
                Ok(Node {
                    identity: identity.clone(),
                    kind: UntypedKind::Call {
                        callee: Box::new(self.make_id_node("quote", identity.clone())),
                        args: Box::new(Node {
                            identity,
                            kind: UntypedKind::Tuple {
                                elements: vec![expr],
                            },
                            ty: (),
                        }),
                    },
                    ty: (),
                })
            }
            TokenKind::Backtick => {
                self.advance()?; // consume `
                let expr = self.parse_expression()?;
                Ok(Node {
                    identity,
                    kind: UntypedKind::Template(Box::new(expr)),
                    ty: (),
                })
            }
            TokenKind::Tilde => {
                self.advance()?; // consume ~
                if *self.peek() == TokenKind::At {
                    self.advance()?; // consume @
                    let expr = self.parse_expression()?;
                    Ok(Node {
                        identity,
                        kind: UntypedKind::Splice(Box::new(expr)),
                        ty: (),
                    })
                } else {
                    let expr = self.parse_expression()?;
                    Ok(Node {
                        identity,
                        kind: UntypedKind::Placeholder(Box::new(expr)),
                        ty: (),
                    })
                }
            }
            _ => self.parse_atom(),
        }
    }

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

    fn parse_atom(&mut self) -> Result<Node<UntypedKind>, String> {
        let token = self.advance()?;
        let identity = Rc::new(NodeIdentity {
            location: token.location,
        });

        let kind = match token.kind {
            TokenKind::Integer(n) => UntypedKind::Constant(Value::Int(n)),
            TokenKind::Float(n) => UntypedKind::Constant(Value::Float(n)),
            TokenKind::String(s) => UntypedKind::Constant(Value::Text(s)),
            TokenKind::Keyword(k) => UntypedKind::Constant(Value::Keyword(Keyword::intern(&k))),
            TokenKind::Identifier(id) => match id.as_ref() {
                "..." => UntypedKind::Nop,
                _ => UntypedKind::Identifier(id.into()),
            },
            _ => {
                return Err(format!(
                    "Unexpected token in atom: {:?} at {:?}",
                    token.kind, token.location
                ));
            }
        };

        Ok(Node {
            identity,
            kind,
            ty: (),
        })
    }

    fn parse_list(&mut self) -> Result<Node<UntypedKind>, String> {
        let start_loc = self.advance()?.location; // consume '('
        let identity = Rc::new(NodeIdentity {
            location: start_loc,
        });

        if *self.peek() == TokenKind::RightParen {
            return Err(format!(
                "Empty list () is not a valid expression at {:?}",
                start_loc
            ));
        }

        let head = self.parse_expression()?;

        let result = if let UntypedKind::Identifier(ref sym) = head.kind {
            match sym.name.as_ref() {
                "if" => self.parse_if(identity),
                "fn" => self.parse_fn(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)?;
        result
    }

    fn parse_again(&mut self, identity: Identity) -> Result<Node<UntypedKind>, String> {
        let args = Box::new(self.parse_expression()?);
        Ok(Node {
            identity,
            kind: UntypedKind::Again { args },
            ty: (),
        })
    }

    fn parse_if(&mut self, identity: Identity) -> Result<Node<UntypedKind>, String> {
        let cond = Box::new(self.parse_expression()?);
        let then_br = Box::new(self.parse_expression()?);
        let mut else_br = None;

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

        Ok(Node {
            identity,
            kind: UntypedKind::If {
                cond,
                then_br,
                else_br,
            },
            ty: (),
        })
    }

    fn parse_def(&mut self, identity: Identity) -> Result<Node<UntypedKind>, String> {
        let target = Box::new(self.parse_pattern()?);
        let value = Box::new(self.parse_expression()?);

        Ok(Node {
            identity,
            kind: UntypedKind::Def { target, value },
            ty: (),
        })
    }

    fn parse_assign(&mut self, identity: Identity) -> Result<Node<UntypedKind>, String> {
        // (assign target value)
        let target = Box::new(self.parse_expression()?);
        let value = Box::new(self.parse_expression()?);

        Ok(Node {
            identity,
            kind: UntypedKind::Assign { target, value },
            ty: (),
        })
    }

    fn parse_do(&mut self, identity: Identity) -> Result<Node<UntypedKind>, String> {
        let mut exprs = Vec::new();
        while *self.peek() != TokenKind::RightParen && *self.peek() != TokenKind::EOF {
            exprs.push(self.parse_expression()?);
        }
        Ok(Node {
            identity,
            kind: UntypedKind::Block { exprs },
            ty: (),
        })
    }

    fn parse_fn(&mut self, identity: Identity) -> Result<Node<UntypedKind>, String> {
        let params = Box::new(self.parse_param_vector()?);
        let body = self.parse_expression()?;

        Ok(Node {
            identity,
            kind: UntypedKind::Lambda {
                params,
                body: Rc::new(body),
            },
            ty: (),
        })
    }

    fn parse_macro_decl(&mut self, identity: Identity) -> Result<Node<UntypedKind>, String> {
        let name_node = self.parse_expression()?;
        let name = match name_node.kind {
            UntypedKind::Identifier(sym) => sym,
            _ => return Err("Expected identifier for macro name".to_string()),
        };
        let params = Box::new(self.parse_param_vector()?);
        let body = self.parse_expression()?;
        Ok(Node {
            identity,
            kind: UntypedKind::MacroDecl {
                name,
                params,
                body: Box::new(body),
            },
            ty: (),
        })
    }

    fn parse_param_vector(&mut self) -> Result<Node<UntypedKind>, String> {
        if *self.peek() != TokenKind::LeftBracket {
            return Err(format!(
                "Expected parameter vector [...] for fn, found {:?}",
                self.peek()
            ));
        }
        self.parse_pattern()
    }

    fn parse_pattern(&mut self) -> Result<Node<UntypedKind>, String> {
        let next = self.peek();
        match next {
            TokenKind::Identifier(_) => {
                let token = self.advance()?;
                let sym: Symbol = match token.kind {
                    TokenKind::Identifier(s) => s.into(),
                    _ => unreachable!(),
                };
                Ok(Node {
                    identity: Rc::new(NodeIdentity {
                        location: token.location,
                    }),
                    kind: UntypedKind::Parameter(sym),
                    ty: (),
                })
            }
            TokenKind::LeftBracket => {
                let token = self.advance()?;
                let identity = Rc::new(NodeIdentity {
                    location: token.location,
                });

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

                Ok(Node {
                    identity,
                    kind: UntypedKind::Tuple { elements },
                    ty: (),
                })
            }
            _ => Err(format!(
                "Expected identifier or pattern vector [...] for definition, found {:?}",
                next
            )),
        }
    }

    fn parse_call(
        &mut self,
        callee: Node<UntypedKind>,
        identity: Identity,
    ) -> Result<Node<UntypedKind>, String> {
        let mut elements = Vec::new();

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

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

        Ok(Node {
            identity,
            kind: UntypedKind::Call {
                callee: Box::new(callee),
                args: Box::new(args_node),
            },
            ty: (),
        })
    }

    fn parse_vector_literal(&mut self) -> Result<Node<UntypedKind>, String> {
        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(expr);
        }

        self.expect(TokenKind::RightBracket)?;

        Ok(Node {
            identity: Rc::new(NodeIdentity {
                location: token.location,
            }),
            kind: UntypedKind::Tuple { elements },
            ty: (),
        })
    }

    fn parse_record_literal(&mut self) -> Result<Node<UntypedKind>, String> {
        let token = self.advance()?;
        let mut fields = Vec::new();

        while *self.peek() != TokenKind::RightBrace {
            if *self.peek() == TokenKind::EOF {
                return Err("Unexpected EOF in record literal".to_string());
            }

            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 {
                UntypedKind::Constant(Value::Keyword(_)) => {}
                _ => return Err("Record keys must be keywords (syntactically)".to_string()),
            }

            if *self.peek() == TokenKind::RightBrace {
                return Err("Record literal must have even number of forms".to_string());
            }
            let val_node = self.parse_expression()?;

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

        Ok(Node {
            identity: Rc::new(NodeIdentity {
                location: token.location,
            }),
            kind: UntypedKind::Record { fields },
            ty: (),
        })
    }

    fn expect(&mut self, kind: TokenKind) -> Result<(), String> {
        let token = self.advance()?;
        if token.kind == kind {
            Ok(())
        } else {
            Err(format!(
                "Expected {:?}, but found {:?} at {:?}",
                kind, token.kind, token.location
            ))
        }
    }

    fn make_id_node(&self, name: &str, identity: Identity) -> Node<UntypedKind> {
        Node {
            identity,
            kind: UntypedKind::Identifier(Symbol::from(name)),
            ty: (),
        }
    }
}