simple scheme interpreter in rust

scheme.rs

use std::collections::HashMap;
use std::fmt;

#[derive(Debug, Clone, PartialEq)]
pub enum Value {
    Number(f64),
    String(String),
    Symbol(String),
    List(Vec<Value>),
    Function(fn(&[Value]) -> Result<Value, String>),
    Lambda {
        params: Vec<String>,
        body: Box<Value>,
        closure: HashMap<String, Value>,
    },
    Bool(bool),
    Nil,
}

impl fmt::Display for Value {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            Value::Number(n) => write!(f, "{}", n),
            Value::String(s) => write!(f, "\"{}\"", s),
            Value::Symbol(s) => write!(f, "{}", s),
            Value::List(items) => {
                write!(f, "(")?;
                for (i, item) in items.iter().enumerate() {
                    if i > 0 { write!(f, " ")?; }
                    write!(f, "{}", item)?;
                }
                write!(f, ")")
            }
            Value::Function(_) => write!(f, "#<function>"),
            Value::Lambda { .. } => write!(f, "#<lambda>"),
            Value::Bool(b) => write!(f, "{}", if *b { "#t" } else { "#f" }),
            Value::Nil => write!(f, "()"),
        }
    }
}

#[derive(Clone)]
pub struct Environment {
    vars: HashMap<String, Value>,
    parent: Option<Box<Environment>>,
}

impl Environment {
    pub fn new() -> Self {
        let mut env = Environment {
            vars: HashMap::new(),
            parent: None,
        };
        env.define_builtins();
        env
    }

    pub fn with_parent(parent: Environment) -> Self {
        Environment {
            vars: HashMap::new(),
            parent: Some(Box::new(parent)),
        }
    }

    pub fn define(&mut self, name: String, value: Value) {
        self.vars.insert(name, value);
    }

    pub fn get(&self, name: &str) -> Option<Value> {
        self.vars.get(name).cloned()
            .or_else(|| self.parent.as_ref().and_then(|p| p.get(name)))
    }

    fn define_builtins(&mut self) {
        self.define("+".to_string(), Value::Function(|args| {
            let sum = args.iter().try_fold(0.0, |acc, arg| {
                match arg {
                    Value::Number(n) => Ok(acc + n),
                    _ => Err("+ requires numbers".to_string()),
                }
            })?;
            Ok(Value::Number(sum))
        }));

        self.define("-".to_string(), Value::Function(|args| {
            if args.is_empty() {
                return Err("- requires at least one argument".to_string());
            }
            match &args[0] {
                Value::Number(first) => {
                    if args.len() == 1 {
                        Ok(Value::Number(-first))
                    } else {
                        let result = args[1..].iter().try_fold(*first, |acc, arg| {
                            match arg {
                                Value::Number(n) => Ok(acc - n),
                                _ => Err("- requires numbers".to_string()),
                            }
                        })?;
                        Ok(Value::Number(result))
                    }
                }
                _ => Err("- requires numbers".to_string()),
            }
        }));

        self.define("*".to_string(), Value::Function(|args| {
            let product = args.iter().try_fold(1.0, |acc, arg| {
                match arg {
                    Value::Number(n) => Ok(acc * n),
                    _ => Err("* requires numbers".to_string()),
                }
            })?;
            Ok(Value::Number(product))
        }));

        self.define("/".to_string(), Value::Function(|args| {
            if args.is_empty() {
                return Err("/ requires at least one argument".to_string());
            }
            match &args[0] {
                Value::Number(first) => {
                    let result = args[1..].iter().try_fold(*first, |acc, arg| {
                        match arg {
                            Value::Number(n) => {
                                if *n == 0.0 {
                                    Err("Division by zero".to_string())
                                } else {
                                    Ok(acc / n)
                                }
                            }
                            _ => Err("/ requires numbers".to_string()),
                        }
                    })?;
                    Ok(Value::Number(result))
                }
                _ => Err("/ requires numbers".to_string()),
            }
        }));

        self.define("=".to_string(), Value::Function(|args| {
            if args.len() != 2 {
                return Err("= requires exactly 2 arguments".to_string());
            }
            Ok(Value::Bool(args[0] == args[1]))
        }));

        self.define("<".to_string(), Value::Function(|args| {
            if args.len() != 2 {
                return Err("< requires exactly 2 arguments".to_string());
            }
            match (&args[0], &args[1]) {
                (Value::Number(a), Value::Number(b)) => Ok(Value::Bool(a < b)),
                _ => Err("< requires numbers".to_string()),
            }
        }));

        self.define("car".to_string(), Value::Function(|args| {
            if args.len() != 1 {
                return Err("car requires exactly 1 argument".to_string());
            }
            match &args[0] {
                Value::List(items) if !items.is_empty() => Ok(items[0].clone()),
                Value::List(_) => Err("car: empty list".to_string()),
                _ => Err("car requires a list".to_string()),
            }
        }));

        self.define("cdr".to_string(), Value::Function(|args| {
            if args.len() != 1 {
                return Err("cdr requires exactly 1 argument".to_string());
            }
            match &args[0] {
                Value::List(items) if !items.is_empty() => {
                    Ok(Value::List(items[1..].to_vec()))
                }
                Value::List(_) => Ok(Value::Nil),
                _ => Err("cdr requires a list".to_string()),
            }
        }));

        self.define("cons".to_string(), Value::Function(|args| {
            if args.len() != 2 {
                return Err("cons requires exactly 2 arguments".to_string());
            }
            match &args[1] {
                Value::List(items) => {
                    let mut new_list = vec![args[0].clone()];
                    new_list.extend(items.clone());
                    Ok(Value::List(new_list))
                }
                Value::Nil => Ok(Value::List(vec![args[0].clone()])),
                _ => Err("cons: second argument must be a list".to_string()),
            }
        }));

        self.define("list".to_string(), Value::Function(|args| {
            Ok(Value::List(args.to_vec()))
        }));
    }
}

pub struct Interpreter {
    env: Environment,
}

impl Interpreter {
    pub fn new() -> Self {
        Interpreter {
            env: Environment::new(),
        }
    }

    pub fn eval(&mut self, expr: &Value) -> Result<Value, String> {
        let env = &mut self.env;
        Self::eval_with_env(expr, env)
    }

    fn eval_with_env(expr: &Value, env: &mut Environment) -> Result<Value, String> {
        match expr {
            Value::Number(_) | Value::String(_) | Value::Bool(_) | Value::Nil => Ok(expr.clone()),
            
            Value::Symbol(name) => {
                env.get(name).ok_or_else(|| format!("Undefined symbol: {}", name))
            }
            
            Value::List(items) if items.is_empty() => Ok(Value::Nil),
            
            Value::List(items) => {
                let first = &items[0];
                
                match first {
                    Value::Symbol(name) => {
                        match name.as_str() {
                            "quote" => {
                                if items.len() != 2 {
                                    return Err("quote requires exactly 1 argument".to_string());
                                }
                                Ok(items[1].clone())
                            }
                            
                            "if" => {
                                if items.len() != 4 {
                                    return Err("if requires exactly 3 arguments".to_string());
                                }
                                let condition = Self::eval_with_env(&items[1], env)?;
                                match condition {
                                    Value::Bool(false) | Value::Nil => Self::eval_with_env(&items[3], env),
                                    _ => Self::eval_with_env(&items[2], env),
                                }
                            }
                            
                            "define" => {
                                if items.len() != 3 {
                                    return Err("define requires exactly 2 arguments".to_string());
                                }
                                match &items[1] {
                                    Value::Symbol(name) => {
                                        let value = Self::eval_with_env(&items[2], env)?;
                                        env.define(name.clone(), value.clone());
                                        Ok(value)
                                    }
                                    _ => Err("define: first argument must be a symbol".to_string()),
                                }
                            }
                            
                            "lambda" => {
                                if items.len() != 3 {
                                    return Err("lambda requires exactly 2 arguments".to_string());
                                }
                                match &items[1] {
                                    Value::List(param_list) => {
                                        let mut params = Vec::new();
                                        for param in param_list {
                                            match param {
                                                Value::Symbol(name) => params.push(name.clone()),
                                                _ => return Err("lambda: parameters must be symbols".to_string()),
                                            }
                                        }
                                        Ok(Value::Lambda {
                                            params,
                                            body: Box::new(items[2].clone()),
                                            closure: env.vars.clone(),
                                        })
                                    }
                                    _ => Err("lambda: first argument must be a parameter list".to_string()),
                                }
                            }
                            
                            _ => {
                                // Function call
                                let func = Self::eval_with_env(first, env)?;
                                let args: Result<Vec<_>, _> = items[1..].iter()
                                    .map(|arg| Self::eval_with_env(arg, env))
                                    .collect();
                                let args = args?;
                                
                                match func {
                                    Value::Function(f) => f(&args),
                                    Value::Lambda { params, body, closure } => {
                                        if params.len() != args.len() {
                                            return Err(format!("Expected {} arguments, got {}", params.len(), args.len()));
                                        }
                                        
                                        let mut new_env = Environment::with_parent((*env).clone());
                                        for (name, value) in closure {
                                            new_env.define(name, value);
                                        }
                                        for (param, arg) in params.iter().zip(args.iter()) {
                                            new_env.define(param.clone(), arg.clone());
                                        }
                                        
                                        Self::eval_with_env(&body, &mut new_env)
                                    }
                                    _ => Err(format!("Not a function: {}", func)),
                                }
                            }
                        }
                    }
                    _ => Err("First element of list must be a symbol".to_string()),
                }
            }
            
            _ => Err("Cannot evaluate this expression".to_string()),
        }
    }
}

// Parser
pub fn parse(input: &str) -> Result<Vec<Value>, String> {
    let tokens = tokenize(input)?;
    let mut parser = Parser::new(tokens);
    let mut exprs = Vec::new();
    
    while !parser.is_at_end() {
        exprs.push(parser.parse_expr()?);
    }
    
    Ok(exprs)
}

fn tokenize(input: &str) -> Result<Vec<String>, String> {
    let mut tokens = Vec::new();
    let mut chars = input.chars().peekable();
    
    while let Some(ch) = chars.next() {
        match ch {
            ' ' | '\t' | '\n' | '\r' => continue,
            '(' => tokens.push("(".to_string()),
            ')' => tokens.push(")".to_string()),
            '"' => {
                let mut string_val = String::new();
                while let Some(ch) = chars.next() {
                    if ch == '"' {
                        break;
                    }
                    string_val.push(ch);
                }
                tokens.push(format!("\"{}\"", string_val));
            }
            _ => {
                let mut token = String::new();
                token.push(ch);
                while let Some(&next_ch) = chars.peek() {
                    if next_ch.is_whitespace() || next_ch == '(' || next_ch == ')' {
                        break;
                    }
                    token.push(chars.next().unwrap());
                }
                tokens.push(token);
            }
        }
    }
    
    Ok(tokens)
}

struct Parser {
    tokens: Vec<String>,
    current: usize,
}

impl Parser {
    fn new(tokens: Vec<String>) -> Self {
        Parser { tokens, current: 0 }
    }
    
    fn is_at_end(&self) -> bool {
        self.current >= self.tokens.len()
    }
    
    fn current_token(&self) -> Option<&String> {
        self.tokens.get(self.current)
    }
    
    fn advance(&mut self) -> Option<&String> {
        if !self.is_at_end() {
            self.current += 1;
        }
        self.tokens.get(self.current - 1)
    }
    
    fn parse_expr(&mut self) -> Result<Value, String> {
        match self.current_token() {
            Some(token) => {
                if token == "(" {
                    self.parse_list()
                } else {
                    self.parse_atom()
                }
            }
            None => Err("Unexpected end of input".to_string()),
        }
    }
    
    fn parse_list(&mut self) -> Result<Value, String> {
        self.advance(); // consume '('
        let mut items = Vec::new();
        
        while let Some(token) = self.current_token() {
            if token == ")" {
                self.advance(); // consume ')'
                return Ok(Value::List(items));
            }
            items.push(self.parse_expr()?);
        }
        
        Err("Unterminated list".to_string())
    }
    
    fn parse_atom(&mut self) -> Result<Value, String> {
        match self.advance() {
            Some(token) => {
                if token.starts_with('"') && token.ends_with('"') {
                    Ok(Value::String(token[1..token.len()-1].to_string()))
                } else if let Ok(num) = token.parse::<f64>() {
                    Ok(Value::Number(num))
                } else if token == "#t" {
                    Ok(Value::Bool(true))
                } else if token == "#f" {
                    Ok(Value::Bool(false))
                } else {
                    Ok(Value::Symbol(token.clone()))
                }
            }
            None => Err("Unexpected end of input".to_string()),
        }
    }
}

// REPL
pub fn repl() {
    use std::io::{self, Write};
    
    let mut interpreter = Interpreter::new();
    println!("Scheme Interpreter in Rust");
    println!("Type (exit) to quit");
    
    loop {
        print!("scheme> ");
        io::stdout().flush().unwrap();
        
        let mut input = String::new();
        io::stdin().read_line(&mut input).unwrap();
        let input = input.trim();
        
        if input == "(exit)" {
            break;
        }
        
        if input.is_empty() {
            continue;
        }
        
        match parse(input) {
            Ok(exprs) => {
                for expr in exprs {
                    match interpreter.eval(&expr) {
                        Ok(result) => println!("{}", result),
                        Err(e) => println!("Error: {}", e),
                    }
                }
            }
            Err(e) => println!("Parse error: {}", e),
        }
    }
}

fn main() {
    repl();
}

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

    #[test]
    fn test_arithmetic() {
        let mut interp = Interpreter::new();
        
        let expr = parse("(+ 1 2 3)").unwrap()[0].clone();
        assert_eq!(interp.eval(&expr).unwrap(), Value::Number(6.0));
        
        let expr = parse("(* 2 3 4)").unwrap()[0].clone();
        assert_eq!(interp.eval(&expr).unwrap(), Value::Number(24.0));
    }
    
    #[test]
    fn test_define_and_lookup() {
        let mut interp = Interpreter::new();
        
        let expr = parse("(define x 42)").unwrap()[0].clone();
        interp.eval(&expr).unwrap();
        
        let expr = parse("x").unwrap()[0].clone();
        assert_eq!(interp.eval(&expr).unwrap(), Value::Number(42.0));
    }
    
    #[test]
    fn test_lambda() {
        let mut interp = Interpreter::new();
        
        let expr = parse("(define square (lambda (x) (* x x)))").unwrap()[0].clone();
        interp.eval(&expr).unwrap();
        
        let expr = parse("(square 5)").unwrap()[0].clone();
        assert_eq!(interp.eval(&expr).unwrap(), Value::Number(25.0));
    }
    
    #[test]
    fn test_if() {
        let mut interp = Interpreter::new();
        
        let expr = parse("(if #t 1 2)").unwrap()[0].clone();
        assert_eq!(interp.eval(&expr).unwrap(), Value::Number(1.0));
        
        let expr = parse("(if #f 1 2)").unwrap()[0].clone();
        assert_eq!(interp.eval(&expr).unwrap(), Value::Number(2.0));
    }
    
    #[test]
    fn test_list_operations() {
        let mut interp = Interpreter::new();
        
        let expr = parse("(car (list 1 2 3))").unwrap()[0].clone();
        assert_eq!(interp.eval(&expr).unwrap(), Value::Number(1.0));
        
        let expr = parse("(cdr (list 1 2 3))").unwrap()[0].clone();
        assert_eq!(interp.eval(&expr).unwrap(), Value::List(vec![Value::Number(2.0), Value::Number(3.0)]));
    }
}