lambda-cat 0.1.0

//! Recursive-descent parser for the lambda calculus surface syntax.
//!
//! The grammar is:
//!
//! ```text
//! expr     ::= lambda | let | fix | app_expr
//! lambda   ::= "\" ident "." expr
//! let      ::= "let" ident "=" expr "in" expr
//! fix      ::= "fix" ident "." expr
//! app_expr ::= atom atom*                  (left-associative application)
//! atom     ::= ident | "(" expr ")"
//! ```
//!
//! Application binds tighter than abstraction and `let`/`fix`, so
//! `\x. x y` parses as `\x. (x y)` and `let id = \x. x in id id` parses as
//! `let id = (\x. x) in (id id)`.

use crate::error::Error;
use crate::lexer::{Token, TokenKind};
use crate::syntax::{Expr, VarName};

/// Look-ahead result for the parser: either a borrowed token or end-of-input.
enum Peek<'a> {
    Eof,
    Tok(&'a Token),
}

fn peek(tokens: &[Token], pos: usize) -> Peek<'_> {
    tokens.get(pos).map_or(Peek::Eof, Peek::Tok)
}

/// Parse a full expression from a token slice.
///
/// # Errors
///
/// Returns [`Error::UnexpectedToken`] if extra tokens trail after the
/// expression, [`Error::UnexpectedEnd`] if input is exhausted mid-production,
/// or other parse errors propagated from sub-productions.
///
/// # Examples
///
/// ```
/// # fn main() -> Result<(), lambda_cat::error::Error> {
/// use lambda_cat::lexer::lex;
/// use lambda_cat::parser::parse;
///
/// let tokens = lex("\\x. x")?;
/// let expr = parse(&tokens)?;
/// assert_eq!(format!("{expr}"), "(\\x. x)");
/// # Ok(())
/// # }
/// ```
///
/// [`Error::UnexpectedToken`]: crate::error::Error::UnexpectedToken
/// [`Error::UnexpectedEnd`]: crate::error::Error::UnexpectedEnd
pub fn parse(tokens: &[Token]) -> Result<Expr, Error> {
    let (expr, end_pos) = parse_expr(tokens, 0)?;
    tokens.get(end_pos).map_or(Ok(expr), |tok| {
        Err(Error::UnexpectedToken {
            at: tok.at(),
            expected: "end of input",
            found: format!("{}", tok.kind()),
        })
    })
}

fn parse_expr(tokens: &[Token], pos: usize) -> Result<(Expr, usize), Error> {
    match peek(tokens, pos) {
        Peek::Eof => Err(Error::UnexpectedEnd {
            expected: "expression",
        }),
        Peek::Tok(tok) => dispatch_expr(tokens, pos, tok),
    }
}

fn dispatch_expr(tokens: &[Token], pos: usize, tok: &Token) -> Result<(Expr, usize), Error> {
    match tok.kind() {
        TokenKind::Lambda => parse_lambda(tokens, pos),
        TokenKind::KwLet => parse_let(tokens, pos),
        TokenKind::KwFix => parse_fix(tokens, pos),
        TokenKind::Ident(_) | TokenKind::LParen => parse_app(tokens, pos),
        TokenKind::KwIn | TokenKind::Dot | TokenKind::Equals | TokenKind::RParen => {
            Err(Error::UnexpectedToken {
                at: tok.at(),
                expected: "expression",
                found: format!("{}", tok.kind()),
            })
        }
    }
}

fn parse_lambda(tokens: &[Token], pos: usize) -> Result<(Expr, usize), Error> {
    let after_lambda = expect_kind(tokens, pos, &TokenKind::Lambda, "`\\`")?;
    let (param, after_ident) = expect_ident(tokens, after_lambda)?;
    let after_dot = expect_kind(tokens, after_ident, &TokenKind::Dot, "`.`")?;
    let (body, after_body) = parse_expr(tokens, after_dot)?;
    Ok((Expr::lam(param, body), after_body))
}

fn parse_let(tokens: &[Token], pos: usize) -> Result<(Expr, usize), Error> {
    let after_let = expect_kind(tokens, pos, &TokenKind::KwLet, "keyword `let`")?;
    let (name, after_name) = expect_ident(tokens, after_let)?;
    let after_eq = expect_kind(tokens, after_name, &TokenKind::Equals, "`=`")?;
    let (value, after_value) = parse_expr(tokens, after_eq)?;
    let after_in = expect_kind(tokens, after_value, &TokenKind::KwIn, "keyword `in`")?;
    let (body, after_body) = parse_expr(tokens, after_in)?;
    Ok((Expr::bind(name, value, body), after_body))
}

fn parse_fix(tokens: &[Token], pos: usize) -> Result<(Expr, usize), Error> {
    let after_fix = expect_kind(tokens, pos, &TokenKind::KwFix, "keyword `fix`")?;
    let (name, after_name) = expect_ident(tokens, after_fix)?;
    let after_dot = expect_kind(tokens, after_name, &TokenKind::Dot, "`.`")?;
    let (body, after_body) = parse_expr(tokens, after_dot)?;
    Ok((Expr::fix(name, body), after_body))
}

fn parse_app(tokens: &[Token], pos: usize) -> Result<(Expr, usize), Error> {
    let (head, after_head) = parse_atom(tokens, pos)?;
    parse_app_tail(tokens, after_head, head)
}

fn parse_app_tail(tokens: &[Token], pos: usize, lhs: Expr) -> Result<(Expr, usize), Error> {
    let can_apply = starts_atom(tokens, pos);
    if can_apply {
        let (arg, after_arg) = parse_atom(tokens, pos)?;
        parse_app_tail(tokens, after_arg, Expr::app(lhs, arg))
    } else {
        Ok((lhs, pos))
    }
}

fn starts_atom(tokens: &[Token], pos: usize) -> bool {
    match peek(tokens, pos) {
        Peek::Eof => false,
        Peek::Tok(tok) => matches!(tok.kind(), TokenKind::Ident(_) | TokenKind::LParen),
    }
}

fn parse_atom(tokens: &[Token], pos: usize) -> Result<(Expr, usize), Error> {
    match peek(tokens, pos) {
        Peek::Eof => Err(Error::UnexpectedEnd {
            expected: "atom (identifier or `(`)",
        }),
        Peek::Tok(tok) => dispatch_atom(tokens, pos, tok),
    }
}

fn dispatch_atom(tokens: &[Token], pos: usize, tok: &Token) -> Result<(Expr, usize), Error> {
    match tok.kind() {
        TokenKind::Ident(name) => Ok((Expr::Var(name.clone()), pos + 1)),
        TokenKind::LParen => parse_paren_inner(tokens, pos + 1),
        TokenKind::Lambda
        | TokenKind::KwLet
        | TokenKind::KwFix
        | TokenKind::KwIn
        | TokenKind::Dot
        | TokenKind::Equals
        | TokenKind::RParen => Err(Error::UnexpectedToken {
            at: tok.at(),
            expected: "atom (identifier or `(`)",
            found: format!("{}", tok.kind()),
        }),
    }
}

fn parse_paren_inner(tokens: &[Token], pos: usize) -> Result<(Expr, usize), Error> {
    let (inner, after_inner) = parse_expr(tokens, pos)?;
    let after_close = expect_kind(tokens, after_inner, &TokenKind::RParen, "`)`")?;
    Ok((inner, after_close))
}

fn expect_ident(tokens: &[Token], pos: usize) -> Result<(VarName, usize), Error> {
    match peek(tokens, pos) {
        Peek::Eof => Err(Error::UnexpectedEnd {
            expected: "identifier",
        }),
        Peek::Tok(tok) => dispatch_expect_ident(tok, pos),
    }
}

fn dispatch_expect_ident(tok: &Token, pos: usize) -> Result<(VarName, usize), Error> {
    match tok.kind() {
        TokenKind::Ident(name) => Ok((name.clone(), pos + 1)),
        TokenKind::Lambda
        | TokenKind::KwLet
        | TokenKind::KwIn
        | TokenKind::KwFix
        | TokenKind::Dot
        | TokenKind::Equals
        | TokenKind::LParen
        | TokenKind::RParen => Err(Error::UnexpectedToken {
            at: tok.at(),
            expected: "identifier",
            found: format!("{}", tok.kind()),
        }),
    }
}

fn expect_kind(
    tokens: &[Token],
    pos: usize,
    expected: &TokenKind,
    name: &'static str,
) -> Result<usize, Error> {
    match peek(tokens, pos) {
        Peek::Eof => Err(Error::UnexpectedEnd { expected: name }),
        Peek::Tok(tok) => {
            let matches = tok.kind() == expected;
            if matches {
                Ok(pos + 1)
            } else {
                Err(Error::UnexpectedToken {
                    at: tok.at(),
                    expected: name,
                    found: format!("{}", tok.kind()),
                })
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::lexer::lex;

    fn parse_str(src: &str) -> Result<Expr, Error> {
        let tokens = lex(src)?;
        parse(&tokens)
    }

    #[test]
    fn identity_function() -> Result<(), Error> {
        let e = parse_str("\\x. x")?;
        let expected = Expr::lam("x", Expr::var("x"));
        (e == expected).then_some(()).ok_or(Error::UnexpectedEnd {
            expected: "identity AST",
        })
    }

    #[test]
    fn left_associative_application() -> Result<(), Error> {
        let e = parse_str("f x y")?;
        let expected = Expr::app(Expr::app(Expr::var("f"), Expr::var("x")), Expr::var("y"));
        (e == expected).then_some(()).ok_or(Error::UnexpectedEnd {
            expected: "left-assoc AST",
        })
    }

    #[test]
    fn let_with_body() -> Result<(), Error> {
        let e = parse_str("let id = \\x. x in id id")?;
        let expected = Expr::bind(
            "id",
            Expr::lam("x", Expr::var("x")),
            Expr::app(Expr::var("id"), Expr::var("id")),
        );
        (e == expected).then_some(()).ok_or(Error::UnexpectedEnd {
            expected: "let AST",
        })
    }
}