Simple Lisp in Rust

lisp.rs

use std::{cell::RefCell, collections::HashMap, fmt::Display, rc::Rc};

// The `Expr` type is a type that represents a lisp S-expression.
// It can be either a `Nil` value, a `Cons` cell, a `Number, or an `Atom`.
#[derive(Clone)]
pub enum Expr {
    Nil,
    Cons(Rc<Expr>, Rc<Expr>),
    Num(u64),
    Str(String),
    Atom(String),

    Prim(Prim),
    Closure(Closure),
}

pub type Prim = fn(Env, Vec<Rc<Expr>>) -> Rc<Expr>;

#[derive(Clone)]
pub struct Closure {
    pub env: Env,
    pub params: Vec<String>,
    pub body: Rc<Expr>,
}

impl Expr {
    pub fn as_atom(&self) -> String {
        match self {
            Expr::Atom(s) => s.clone(),
            _ => panic!("not an atom"),
        }
    }

    pub fn as_cons(&self) {
        match self {
            Expr::Cons(car, cdr) => (),
            _ => panic!("not a cons"),
        }
    }

    pub fn accumulate_cons(&self) -> (Vec<Rc<Expr>>, Option<&Expr>) {
        let mut items = Vec::new();
        let mut current = self;
        loop {
            match current {
                Expr::Cons(car, cdr) => {
                    items.push(car.clone());
                    current = cdr;
                }
                Expr::Nil => {
                    return (items, None);
                }
                _ => {
                    return (items, Some(current));
                }
            }
        }
    }
}

impl Display for Expr {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Expr::Nil => write!(f, "nil"),
            Expr::Cons(_, _) => {
                let (items, end) = self.accumulate_cons();

                write!(f, "(")?;

                if !items.is_empty() {
                    write!(f, "{}", items[0])?;
                    for item in &items[1..] {
                        write!(f, " {}", item)?;
                    }
                }

                if let Some(end) = end {
                    write!(f, ". {}", end)?;
                }

                write!(f, ")")
            }
            Expr::Num(n) => write!(f, "{}", n),
            Expr::Atom(s) => write!(f, "{}", s),
            Expr::Str(s) => write!(f, "\"{}\"", s),
            Expr::Prim(_) => write!(f, "<fun>"),
            Expr::Closure(_) => write!(f, "<closure>"),
        }
    }
}

// Parser

pub fn parse(code: &str) -> Vec<Expr> {
    let mut stack = Vec::new();
    let mut mark = Vec::new();

    let mut iter = code.chars().peekable();

    while let Some(char) = iter.next() {
        match char {
            ' ' | '\n' | '\t' | '\r' => (),
            '0'..='9' => {
                let mut num = char.to_digit(10).unwrap() as u64;
                while let Some('0'..='9') = iter.peek() {
                    num = num * 10 + iter.next().unwrap().to_digit(10).unwrap() as u64;
                }
                stack.push(Expr::Num(num));
            }
            '(' => {
                mark.push(stack.len());
            }
            ')' => {
                let start = mark.pop().expect("unmatched ')'");
                let items = stack.split_off(start);

                let result = items.into_iter().rfold(Expr::Nil, |item, acc| {
                    Expr::Cons(Rc::new(acc), Rc::new(item))
                });

                stack.push(result)
            }
            '"' => {
                let mut string = String::new();
                while let Some(x) = iter.peek() {
                    if x == &'"' {
                        iter.next();
                        break;
                    }
                    string.push(iter.next().unwrap());
                }
                stack.push(Expr::Str(string));
            }
            char => {
                let mut atom = String::new();
                atom.push(char);
                while !matches!(
                    iter.peek(),
                    Some(' ' | '\n' | '\t' | '\r' | '(' | ')' | '"')
                ) {
                    atom.push(iter.next().unwrap());
                }
                stack.push(Expr::Atom(atom));
            }
        }
    }

    if !mark.is_empty() {
        panic!("unmatched '('");
    }

    stack
}

// Interpreter

#[derive(Clone)]
pub struct Env {
    pub globals: Rc<RefCell<HashMap<String, Rc<Expr>>>>,
    pub locals: im::HashMap<String, Rc<Expr>>,
}

impl Env {
    pub fn with(self, name: String, expr: Rc<Expr>) -> Self {
        Self {
            globals: self.globals,
            locals: self.locals.update(name, expr),
        }
    }

    pub fn add_global(&self, name: String, expr: Rc<Expr>) {
        self.globals.borrow_mut().insert(name, expr);
    }

    pub fn add_prim(&mut self, name: &str, expr: Prim) {
        self.add_global(name.to_string(), Rc::new(Expr::Prim(expr)))
    }

    pub fn eval_fun(&self, fun: Rc<Expr>, args: Vec<Rc<Expr>>) -> Rc<Expr> {
        match &*fun {
            Expr::Prim(fun) => fun(self.clone(), args),
            Expr::Closure(closure) => {
                let args = required_args(args, closure.params.len());

                let mut ctx = closure.env.clone();

                for (name, arg) in closure.params.iter().zip(args) {
                    ctx = ctx.with(name.clone(), arg);
                }

                ctx.eval(closure.body.clone())
            }
            _ => panic!("not a function"),
        }
    }

    pub fn eval(&self, term: Rc<Expr>) -> Rc<Expr> {
        match &*term {
            Expr::Cons(fun, args) => {
                let function_value = self.eval(fun.clone());
                let (args, _) = args.accumulate_cons();

                self.eval_fun(function_value, args)
            }
            Expr::Atom(name) => {
                if let Some(value) = self.locals.get(name) {
                    value.clone()
                } else {
                    let globals = self.globals.borrow_mut();
                    globals
                        .get(name)
                        .unwrap_or_else(|| panic!("undefined variable `{}`", name))
                        .clone()
                }
            }
            _ => term.clone(),
        }
    }
}

pub fn required_args(args: Vec<Rc<Expr>>, n: usize) -> Vec<Rc<Expr>> {
    if args.len() != n {
        panic!("expected {} arguments, got {}", n, args.len());
    }
    args
}

fn main() {
    let code = "
        (set ata (lambda (x) (lambda (y) (print x y))))
        (set atb (ata 1))
        (atb 3)
    ";

    let parsed = parse(code);

    let mut env = Env {
        globals: Rc::new(RefCell::new(HashMap::new())),
        locals: im::HashMap::new(),
    };

    env.add_prim("print", |env, args| {
        for arg in args {
            print!("{}", env.eval(arg.clone()));
        }
        Rc::new(Expr::Nil)
    });

    env.add_prim("set", |env, args| {
        let args = required_args(args, 2);
        let name = args[0].as_atom();
        let val = env.eval(args[1].clone());
        env.add_global(name, val);

        Rc::new(Expr::Nil)
    });

    // (lambda (x y z) 2)

    env.add_prim("lambda", |env, args| {
        let args = required_args(args, 2);
        args[0].as_cons();

        let (params_raw, _) = args[0].accumulate_cons();

        let mut params = Vec::new();

        for param in &params_raw {
            params.push(param.as_atom());
        }

        let body = args[1].clone();

        Rc::new(Expr::Closure(Closure { env, params, body }))
    });

    for expr in &parsed {
        println!("{}", expr);
    }

    println!("----------------------------------");

    for expr in parsed {
        let rc = Rc::new(expr);
        env.eval(rc);
    }
}