Define lambda shorthand

src/grammar.lalrpop

@@ -1,4 +1,5 @@
-use crate::{Tree, Literal};
+use crate::{Tree, Literal, Lambda};
+use std::rc::Rc;
 
 grammar<'s>();
 
@@ -6,7 +7,9 @@ grammar<'s>();
 pub(crate) Tree: Tree = {
     "(" <Tree+> ")" => Tree::Branch(<>),
     "(" "define" <Sym> <BTree> ")" => Tree::Define(<>),
+    "(" "define" "(" <name:Sym> <args:RcSlice<Sym>> ")" <body:Rc<Tree>> ")" => Tree::Define(name, Box::new(Tree::Lambda(Lambda(args, body)))),
     "(" "if" <Tree> <Tree> <Tree> ")" => Tree::If(Box::new([<>])),
+    "(" "lambda (" <RcSlice<Sym>> ")" <Rc<Tree>> ")" => Tree::Lambda(Lambda(<>)),
     Literal => Tree::Leaf(<>),
 }
 
@@ -19,6 +22,8 @@ Literal: Literal = {
 }
 
 Box<T>: Box<T> = { T => Box::new(<>) }
+RcSlice<T>: Rc<[T]> = { T* => <>.into() }
+Rc<T>: Rc<T> = { T => Rc::new(<>) }
 
 Sym: String = { r"[A-Za-z!$%&*+\-./:<=>?@^_~][A-Za-z!$%&*+\-./:<=>?@^_~0-9]*" => <>.to_owned() }
 Num: f64 = { r"[0-9]+(\.[0-9]+)?" => <>.parse().unwrap() }

src/main.rs

@@ -11,15 +11,20 @@ use rustyline::validate::{
 use rustyline::Editor;
 use rustyline_derive::{Completer, Helper, Highlighter, Hinter};
 use std::collections::HashMap;
+use std::rc::Rc;
 
 #[derive(Debug, debug2::Debug, PartialEq)]
 enum Tree {
     Leaf(Literal),
     Define(String, Box<Tree>),
     If(Box<[Tree; 3]>),
+    Lambda(Lambda),
     Branch(Vec<Tree>),
 }
 
+#[derive(Debug, debug2::Debug, PartialEq, Clone)]
+struct Lambda(Rc<[String]>, Rc<Tree>);
+
 #[derive(Debug, debug2::Debug, PartialEq)]
 enum Literal {
     Sym(String),
@@ -36,10 +41,11 @@ fn err<T>(s: String) -> Result<T, RTError> {
 
 type Func = fn(&[Value]) -> Result<Value, RTError>;
 
-#[derive(Clone, Copy)]
+#[derive(Clone)]
 enum Value {
     Num(f64),
     Func(Func),
+    Lambda(Lambda),
     Bool(bool),
     NotAValue,
 }
@@ -49,11 +55,15 @@ impl PartialEq for Value {
         use Value::*;
         match (self, other) {
             (Num(l), Num(r)) => l == r,
-            (Num(_), _) => false,
             (Func(l), Func(r)) => *l as usize == *r as usize,
-            (Func(_), _) => false,
             (Bool(l), Bool(r)) => l == r,
+            (Lambda(l), Lambda(r)) => Rc::ptr_eq(&l.0, &r.0) && Rc::ptr_eq(&l.1, &r.1),
+
+            (Num(_), _) => false,
+            (Func(_), _) => false,
             (Bool(_), _) => false,
+            (Lambda(_), _) => false,
+
             (NotAValue, _) => panic!("Trap value"),
         }
     }
@@ -64,6 +74,7 @@ impl std::fmt::Debug for Value {
         match self {
             Self::Num(n) => n.fmt(f),
             Self::Func(_) => f.write_str("#<procedure>"),
+            Self::Lambda(_) => f.write_str("#<procedure>"),
             Self::Bool(b) => f.write_str(if *b { "#t" } else { "#f" }),
             Self::NotAValue => Ok(()),
         }
@@ -79,11 +90,11 @@ impl Value {
         }
     }
 
-    fn as_func(&self) -> Result<Func, RTError> {
-        if let Self::Func(f) = self {
-            Ok(*f)
-        } else {
-            Err(RTError("Expected a function".to_owned()))
+    fn as_func(&self) -> Result<Callable, RTError> {
+        match self {
+            Self::Func(f) => Ok(Callable::Func(*f)),
+            Self::Lambda(l) => Ok(Callable::Lambda(l)),
+            _ => Err(RTError("Expected a function".to_owned())),
         }
     }
 
@@ -96,8 +107,37 @@ impl Value {
     }
 }
 
-struct Env {
+enum Callable<'a> {
+    Func(Func),
+    Lambda(&'a Lambda),
+}
+
+struct Env<'a> {
     vars: HashMap<String, Value>,
+    enclosing: Option<&'a Env<'a>>,
+}
+
+impl<'a> Env<'a> {
+    fn child(&'a self) -> Env<'a> {
+        Env {
+            vars: HashMap::new(),
+            enclosing: Some(self),
+        }
+    }
+
+    fn lookup(&self, s: &str) -> Option<Value> {
+        if let Some(v) = self.vars.get(s) {
+            Some(v.clone())
+        } else if let Some(parent) = self.enclosing {
+            parent.lookup(s)
+        } else {
+            None
+        }
+    }
+
+    fn define(&mut self, name: String, val: Value) {
+        self.vars.insert(name, val);
+    }
 }
 
 fn main() {
@@ -112,7 +152,7 @@ fn main() {
         rl.add_history_entry(&line);
 
         let tree = grammar::TreeParser::new().parse(&line).unwrap();
-        dbg!(&tree);
+        // dbg!(&tree);
         println!("< {:?}", eval(&tree, &mut env))
     }
 }
@@ -120,26 +160,42 @@ fn main() {
 fn eval(t: &Tree, env: &mut Env) -> Result<Value, RTError> {
     Ok(match t {
         Tree::Leaf(l) => match l {
-            Literal::Sym(s) => match env.vars.get(s) {
-                Some(&v) => v,
+            Literal::Sym(s) => match env.lookup(s) {
+                Some(v) => v.clone(),
                 None => return err(format!("Undefined variable `{}`", s)),
             },
             Literal::Num(v) => Value::Num(*v),
             Literal::Bool(b) => Value::Bool(*b),
         },
+        Tree::Lambda(l) => Value::Lambda(l.clone()),
         Tree::Branch(args) => {
             let Some(fun) = args.get(0) else { return err("No argument given".to_owned()) };
-            let fun = eval(fun, env)?.as_func()?;
+            let fun = eval(fun, env)?;
+            let fun = fun.as_func()?;
             let args = args
                 .iter()
                 .skip(1)
                 .map(|a| eval(a, env))
                 .collect::<Result<Vec<_>, _>>()?;
-            return fun(&args);
+            // return fun(&args);
+            match fun {
+                Callable::Func(f) => f(&args)?,
+                Callable::Lambda(l) => {
+                    if l.0.len() == args.len() {
+                        let mut new_env = env.child();
+                        for (x, y) in l.0.iter().zip(args) {
+                            new_env.define(x.clone(), y)
+                        }
+                        eval(&l.1, &mut new_env)?
+                    } else {
+                        return err(format!("Need {} args, got {}", l.0.len(), args.len()));
+                    }
+                }
+            }
         }
         Tree::Define(name, to) => {
             let val = eval(to, env)?;
-            env.vars.insert(name.to_owned(), val);
+            env.define(name.to_owned(), val);
             Value::NotAValue
         }
         Tree::If(box [cond, tcase, fcase]) => {
@@ -150,7 +206,7 @@ fn eval(t: &Tree, env: &mut Env) -> Result<Value, RTError> {
     })
 }
 
-fn default_env() -> Env {
+fn default_env() -> Env<'static> {
     let mut vars = HashMap::new();
 
     for (name, fun) in [
@@ -163,7 +219,10 @@ fn default_env() -> Env {
         vars.insert(name.to_owned(), Value::Func(fun));
     }
 
-    Env { vars }
+    Env {
+        vars,
+        enclosing: None,
+    }
 }
 
 mod prims {