λ-calc interpreter cramb

lc.rs

use std::collections::HashSet;
use std::iter::Peekable;
use std::str::Chars;

#[derive(Debug, PartialEq, Clone)]
enum Token {
    Lambda,
    Dot,
    Variable(String),
    LParen,
    RParen,
}

fn tokenize(input: &str) -> Vec<Token> {
    let mut tokens = Vec::new();
    let mut chars = input.chars().peekable();

    while let Some(&ch) = chars.peek() {
        match ch {
            'λ' => {
                tokens.push(Token::Lambda);
                chars.next();
            },
            '.' => {
                tokens.push(Token::Dot);
                chars.next();
            },
            '(' => {
                tokens.push(Token::LParen);
                chars.next();
            },
            ')' => {
                tokens.push(Token::RParen);
                chars.next();
            },
            'a'..='z' => {
                let mut var = String::new();
                while let Some(&c) = chars.peek() {
                    if c.is_ascii_lowercase() {
                        var.push(c);
                        chars.next();
                    } else {
                        break;
                    }
                }
                tokens.push(Token::Variable(var));
            },
            _ if ch.is_whitespace() => {
                chars.next();
            },
            _ => panic!("Unexpected character: {}", ch),
        }
    }

    tokens
}

#[derive(Debug, PartialEq, Clone)]
enum Expr {
    Var(String),
    Abs(String, Box<Expr>),
    App(Box<Expr>, Box<Expr>),
}

struct Parser<I: Iterator<Item = Token>> {
    tokens: Peekable<I>,
}

impl<I: Iterator<Item = Token>> Parser<I> {
    fn new(tokens: I) -> Self {
        Self {
            tokens: tokens.peekable(),
        }
    }

    fn parse_expr(&mut self) -> Expr {
        let mut expr = self.parse_primary();
        while let Some(&Token::Variable(_)) | Some(&Token::LParen) = self.tokens.peek() {
            let right = self.parse_primary();
            expr = Expr::App(Box::new(expr), Box::new(right));
        }
        expr
    }

    fn parse_primary(&mut self) -> Expr {
        match self.tokens.next() {
            Some(Token::Variable(var)) => Expr::Var(var),
            Some(Token::Lambda) => self.parse_abs(),
            Some(Token::LParen) => {
                let expr = self.parse_expr();
                assert_eq!(self.tokens.next(), Some(Token::RParen));
                expr
            }
            _ => panic!("Unexpected token"),
        }
    }

    fn parse_abs(&mut self) -> Expr {
        if let Some(Token::Variable(var)) = self.tokens.next() {
            assert_eq!(self.tokens.next(), Some(Token::Dot));
            let body = self.parse_expr();
            Expr::Abs(var, Box::new(body))
        } else {
            panic!("Expected variable after λ");
        }
    }
}

fn parse(tokens: Vec<Token>) -> Expr {
    let mut parser = Parser::new(tokens.into_iter());
    let expr = parser.parse_expr();
    if parser.tokens.peek().is_some() {
        panic!("Unexpected tokens after expression");
    }
    expr
}

impl Expr {
    fn alpha_convert(&self, bound_vars: &HashSet<String>, counter: &mut usize) -> Expr {
        match self {
            Expr::Var(x) => Expr::Var(x.clone()),
            Expr::Abs(param, body) => {
                let new_param = format!("{}{}", param, counter);
                *counter += 1;
                let mut new_bound_vars = bound_vars.clone();
                new_bound_vars.insert(new_param.clone());
                Expr::Abs(new_param.clone(), Box::new(body.alpha_convert(&new_bound_vars, counter)))
            }
            Expr::App(e1, e2) => Expr::App(
                Box::new(e1.alpha_convert(bound_vars, counter)),
                Box::new(e2.alpha_convert(bound_vars, counter)),
            ),
        }
    }

    fn substitute(&self, var: &str, val: &Expr) -> Expr {
        match self {
            Expr::Var(x) => {
                if x == var {
                    val.clone()
                } else {
                    Expr::Var(x.clone())
                }
            }
            Expr::Abs(param, body) => {
                if param == var {
                    Expr::Abs(param.clone(), body.clone())
                } else {
                    Expr::Abs(param.clone(), Box::new(body.substitute(var, val)))
                }
            }
            Expr::App(e1, e2) => Expr::App(
                Box::new(e1.substitute(var, val)),
                Box::new(e2.substitute(var, val)),
            ),
        }
    }

    fn beta_reduce(&self) -> Expr {
        match self {
            Expr::Var(_) => self.clone(),
            Expr::Abs(param, body) => Expr::Abs(param.clone(), Box::new(body.beta_reduce())),
            Expr::App(e1, e2) => match &**e1 {
                Expr::Abs(param, body) => {
                    let substituted = body.substitute(param, e2);
                    substituted.beta_reduce()
                }
                _ => Expr::App(Box::new(e1.beta_reduce()), Box::new(e2.beta_reduce())),
            },
        }
    }

    fn evaluate(&self) -> Expr {
        let mut expr = self.clone();
        loop {
            let reduced_expr = expr.beta_reduce();
            if reduced_expr == expr {
                break;
            }
            expr = reduced_expr;
        }
        expr
    }
}

impl std::fmt::Display for Expr {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        match self {
            Expr::Var(var) => write!(f, "{}", var),
            Expr::Abs(var, expr) => write!(f, "\\{}.{}", var, expr),
            Expr::App(expr1, expr2) => write!(f, "({} {})", expr1, expr2),
        }
    }
}

fn main() {
    let input = "(λm.λn.λf.λx.((m f) ((n f) x))) (λf.λx.(f x)) (λf.λx.(f (f x)))";
    let tokens = tokenize(input);
    let ast = parse(tokens);
    let evaluated_expr = ast.evaluate();
    println!("{}", evaluated_expr.to_string());
}