λ-calc interpreter in crablang

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),
        }
    }
}


impl Expr {
    fn evaluate_with_depth(&self, max_depth: usize) -> Result<Expr, String> {
        self.evaluate_internal(max_depth, 0)
    }

    fn evaluate_internal(&self, max_depth: usize, current_depth: usize) -> Result<Expr, String> {
        if current_depth > max_depth {
            return Err("Exceeded maximum recursion depth".to_string());
        }

        match self {
            Expr::Var(_) => Ok(self.clone()),
            Expr::Abs(param, body) => {
                let reduced_body = body.evaluate_internal(max_depth, current_depth + 1)?;
                Ok(Expr::Abs(param.clone(), Box::new(reduced_body)))
            }
            Expr::App(e1, e2) => {
                let reduced_e1 = e1.evaluate_internal(max_depth, current_depth + 1)?;
                let reduced_e2 = e2.evaluate_internal(max_depth, current_depth + 1)?;

                match reduced_e1 {
                    Expr::Abs(param, body) => {
                        let substituted = body.substitute(&param, &reduced_e2);
                        substituted.evaluate_internal(max_depth, current_depth + 1)
                    }
                    _ => Ok(Expr::App(Box::new(reduced_e1), Box::new(reduced_e2))),
                }
            }
        }
    }
}

impl Expr {
    fn evaluate_with_history(&self) -> Result<Expr, String> {
        let mut seen = HashSet::new();
        self.evaluate_internal_history(&mut seen)
    }

    fn evaluate_internal_history(&self, seen: &mut HashSet<String>) -> Result<Expr, String> {
        let expr_string = format!("{:?}", self);
        if seen.contains(&expr_string) {
            return Err("Detected infinite recursion".to_string());
        }

        seen.insert(expr_string);

        match self {
            Expr::Var(_) => Ok(self.clone()),
            Expr::Abs(param, body) => {
                let reduced_body = body.evaluate_internal_history(seen)?;
                Ok(Expr::Abs(param.clone(), Box::new(reduced_body)))
            }
            Expr::App(e1, e2) => {
                let reduced_e1 = e1.evaluate_internal_history(seen)?;
                let reduced_e2 = e2.evaluate_internal_history(seen)?;

                match reduced_e1 {
                    Expr::Abs(param, body) => {
                        let substituted = body.substitute(&param, &reduced_e2);
                        substituted.evaluate_internal_history(seen)
                    }
                    _ => Ok(Expr::App(Box::new(reduced_e1), Box::new(reduced_e2))),
                }
            }
        }
    }
}


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);
    match ast.evaluate_with_depth(1000) {
        Ok(result) => println!("Result: {}", result),
        Err(e) => println!("Error: {}", e),
    }
}