closures work?!

src/base.rs

@@ -53,16 +53,7 @@ pub fn store(b: &Value, val: &Value) -> Value {
 // name, patterns, AND docstring
 pub fn doc(f: &Value) -> Value {
     match f {
-        Value::Fn(f) => {
-            let lfn = f.as_ref().get().unwrap();
-            let name = lfn.name;
-            let doc = lfn.doc;
-            if let Some(docstr) = doc {
-                Value::String(Rc::new(format!("{name}: {docstr}")))
-            } else {
-                Value::String(Rc::new(format!("{name}: no documentation found")))
-            }
-        }
+        Value::Fn(f) => f.as_ref().doc(),
         _ => Value::Interned("no documentation found"),
     }
 }

src/compiler.rs

@@ -5,7 +5,7 @@ use crate::value::*;
 use chumsky::prelude::SimpleSpan;
 use num_derive::{FromPrimitive, ToPrimitive};
 use num_traits::FromPrimitive;
-use std::cell::OnceCell;
+use std::cell::RefCell;
 use std::collections::HashMap;
 use std::rc::Rc;
 
@@ -76,6 +76,9 @@ pub enum Op {
     At,
 
     Not,
+    Print,
+    SetUpvalue,
+    GetUpvalue,
     // Inc,
     // Dec,
     // Gt,
@@ -169,6 +172,7 @@ impl std::fmt::Display for Op {
             EmptyString => "empty_string",
             ConcatStrings => "concat_strings",
             Stringify => "stringify",
+            Print => "print",
 
             Eq => "eq",
             Add => "add",
@@ -185,6 +189,9 @@ impl std::fmt::Display for Op {
 
             Call => "call",
             Return => "return",
+
+            SetUpvalue => "set_upvalue",
+            GetUpvalue => "get_upvalue",
         };
         write!(f, "{rep}")
     }
@@ -197,6 +204,12 @@ pub struct Binding {
     stack_pos: usize,
 }
 
+#[derive(Clone, Debug, PartialEq)]
+pub struct Upvalue {
+    name: &'static str,
+    stack_pos: usize,
+}
+
 #[derive(Clone, Debug, PartialEq)]
 pub struct Chunk {
     pub constants: Vec<Value>,
@@ -214,7 +227,7 @@ impl Chunk {
             | PanicIfNoMatch | ResetMatch | GetKey | PanicNoWhen | PanicNoMatch | TypeOf
             | Duplicate | Decrement | Truncate | Noop | LoadTuple | LoadList | Eq | Add | Sub
             | Mult | Div | Unbox | BoxStore | Assert | Get | At | Not | Panic | EmptyString
-            | ConcatStrings | Stringify | MatchType | Return | Match => {
+            | ConcatStrings | Stringify | MatchType | Return | Match | Print => {
                 println!("{i:04}: {op}")
             }
             Constant | MatchConstant => {
@@ -225,7 +238,8 @@ impl Chunk {
             }
             PushBinding | MatchTuple | MatchList | MatchDict | LoadDictValue | PushTuple
             | PushDict | PushList | PushBox | Jump | JumpIfFalse | JumpIfTrue | JumpIfNoMatch
-            | JumpIfMatch | JumpBack | JumpIfZero | MatchDepth | PopN | StoreAt | Call => {
+            | JumpIfMatch | JumpBack | JumpIfZero | MatchDepth | PopN | StoreAt | Call
+            | SetUpvalue | GetUpvalue => {
                 let next = self.bytecode[*i + 1];
                 println!("{i:04}: {:16} {next:04}", op.to_string());
                 *i += 1;
@@ -246,9 +260,9 @@ impl Chunk {
     //     self.kw_index_from(kw).map(Value::Keyword)
     // }
 
-    pub fn kw_index_from(&self, kw: &str) -> Option<usize> {
-        self.keywords.iter().position(|s| *s == kw)
-    }
+    // pub fn kw_index_from(&self, kw: &str) -> Option<usize> {
+    //     self.keywords.iter().position(|s| *s == kw)
+    // }
 }
 
 #[derive(Debug, Clone, PartialEq)]
@@ -277,13 +291,14 @@ fn get_builtin(name: &str, arity: usize) -> Option<Op> {
         ("get", 2) => Some(Op::Get),
         ("at", 2) => Some(Op::At),
         ("not", 1) => Some(Op::Not),
+        ("print!", 1) => Some(Op::Print),
 
         _ => None,
     }
 }
 
 #[derive(Debug, Clone, PartialEq)]
-pub struct Compiler {
+pub struct Compiler<'a> {
     pub chunk: Chunk,
     pub bindings: Vec<Binding>,
     pub scope_depth: isize,
@@ -295,6 +310,8 @@ pub struct Compiler {
     pub span: SimpleSpan,
     pub src: &'static str,
     pub name: &'static str,
+    pub enclosing: Option<&'a Compiler<'a>>,
+    pub upvalues: Vec<Upvalue>,
     loop_info: Vec<LoopInfo>,
 }
 
@@ -303,13 +320,18 @@ fn is_binding(expr: &Spanned<Ast>) -> bool {
     use Ast::*;
     match ast {
         Let(..) | LBox(..) => true,
-        Fn(name, ..) => !name.is_empty(),
+        // Fn(name, ..) => !name.is_empty(),
         _ => false,
     }
 }
 
-impl Compiler {
-    pub fn new(ast: &'static Spanned<Ast>, name: &'static str, src: &'static str) -> Compiler {
+impl<'a> Compiler<'a> {
+    pub fn new(
+        ast: &'static Spanned<Ast>,
+        name: &'static str,
+        src: &'static str,
+        enclosing: Option<&'a Compiler>,
+    ) -> Compiler<'a> {
         let chunk = Chunk {
             constants: vec![],
             bytecode: vec![],
@@ -329,6 +351,8 @@ impl Compiler {
             ast: &ast.0,
             span: ast.1,
             loop_info: vec![],
+            enclosing,
+            upvalues: vec![],
             src,
             name,
         }
@@ -431,6 +455,40 @@ impl Compiler {
         self.bindings.push(binding);
     }
 
+    fn resolve_local(&self, name: &'static str) -> Option<usize> {
+        for binding in self.bindings.iter() {
+            if binding.name == name {
+                return Some(binding.stack_pos);
+            }
+        }
+        None
+    }
+
+    fn resolve_upvalue(&self, name: &'static str) -> Option<usize> {
+        self.upvalues.iter().position(|uv| uv.name == name)
+    }
+
+    fn get_upvalue(&self, name: &'static str) -> Upvalue {
+        let local = self.bindings.iter().find(|b| b.name == name);
+        match local {
+            Some(binding) => Upvalue {
+                name,
+                stack_pos: binding.stack_pos,
+            },
+            None => self.enclosing.unwrap().get_upvalue(name),
+        }
+    }
+
+    fn resolve_binding(&self, name: &'static str) -> usize {
+        if let Some(pos) = self.resolve_local(name) {
+            return pos;
+        }
+        match self.enclosing {
+            Some(compiler) => compiler.resolve_binding(name),
+            None => unreachable!(),
+        }
+    }
+
     fn pop(&mut self) {
         self.emit_op(Op::Pop);
         self.stack_depth -= 1;
@@ -554,17 +612,30 @@ impl Compiler {
                 self.emit_op(Op::Match);
                 self.bind(name);
             }
-            Word(name) => {
+            Word(name) => match self.resolve_local(name) {
+                Some(position) => {
                     self.emit_op(Op::PushBinding);
+                    self.emit_byte(position);
                     self.stack_depth += 1;
-                let biter = self.bindings.iter().rev();
-                for binding in biter {
-                    if binding.name == *name {
-                        self.emit_byte(binding.stack_pos);
-                        break;
                 }
+                None => match self.resolve_upvalue(name) {
+                    Some(position) => {
+                        println!("resolved upvalue: {name}");
+                        self.emit_op(Op::GetUpvalue);
+                        self.emit_byte(position);
+                        self.stack_depth += 1;
                     }
+                    None => {
+                        println!("setting upvalue: {name}");
+                        let upvalue = self.get_upvalue(name);
+                        self.emit_op(Op::GetUpvalue);
+                        self.emit_byte(upvalue.stack_pos);
+                        self.upvalues.push(upvalue);
+                        dbg!(&self.upvalues);
+                        self.stack_depth += 1;
                     }
+                },
+            },
             PlaceholderPattern => {
                 self.emit_op(Op::Match);
             }
@@ -859,7 +930,6 @@ impl Compiler {
                     },
                     _ => unreachable!(),
                 }
-                // TODO: implement longer synthetic expressions
                 for (term, _) in rest {
                     match term {
                         Keyword(str) => {
@@ -980,6 +1050,7 @@ impl Compiler {
             }
             MatchClause(..) => unreachable!(),
             Fn(name, body, doc) => {
+                let name = if name.is_empty() { "anonymous" } else { name };
                 // first, declare the function
                 // TODO: or, check if the function has already been declared!
                 let FnBody(fn_body) = &body.as_ref().0 else {
@@ -988,6 +1059,8 @@ impl Compiler {
 
                 let mut compilers: HashMap<usize, Compiler> = HashMap::new();
 
+                let mut upvalues = vec![];
+
                 for clause in fn_body {
                     let MatchClause(pattern, _, clause_body) = &clause.0 else {
                         unreachable!()
@@ -1001,8 +1074,8 @@ impl Compiler {
                     let compiler = match compilers.get_mut(&arity) {
                         Some(compiler) => compiler,
                         None => {
-                            let mut compiler = Compiler::new(clause, self.name, self.src);
-                            // compiler.emit_op(Op::Load);
+                            let mut compiler =
+                                Compiler::new(clause, self.name, self.src, Some(self));
                             compiler.emit_op(Op::ResetMatch);
                             compilers.insert(arity, compiler);
                             compilers.get_mut(&arity).unwrap()
@@ -1013,6 +1086,8 @@ impl Compiler {
                     compiler.scope_depth += 1;
                     compiler.match_depth = arity;
 
+                    std::mem::swap(&mut upvalues, &mut compiler.upvalues);
+
                     let mut tup_jump_idxes = vec![];
                     for member in pattern {
                         compiler.match_depth -= 1;
@@ -1056,6 +1131,8 @@ impl Compiler {
                     compiler.emit_op(Op::Return);
                     compiler.chunk.bytecode[jnm_idx] = compiler.len() as u8 - jnm_idx as u8 - 1;
                     compiler.scope_depth -= 1;
+
+                    std::mem::swap(&mut compiler.upvalues, &mut upvalues);
                 }
 
                 let mut the_chunks = vec![];
@@ -1070,25 +1147,26 @@ impl Compiler {
                     the_chunks.push((arity as u8, chunk));
                 }
 
-                let lfn = crate::value::LFn {
+                let lfn = crate::value::LFn::Defined {
                     name,
                     doc: *doc,
                     chunks: the_chunks,
-                    closed: vec![],
+                    closed: RefCell::new(vec![]),
                 };
 
-                let cell = OnceCell::new();
-                let _ = cell.set(lfn);
-                let init_val = Value::Fn(Rc::new(cell));
+                // TODO: check if the function is already declared, and pull out the relevant OnceCell if need be
+                let init_val = Value::Fn(Rc::new(lfn));
                 self.emit_constant(init_val);
                 self.bind(name);
 
                 // TODO: close over everything accessed in the function
-
-                // TODO: pull the function off the stack, and set the OnceCell.
+                for upvalue in upvalues {
+                    self.emit_op(Op::SetUpvalue);
+                    self.emit_byte(upvalue.stack_pos);
+                }
             }
             FnDeclaration(name) => {
-                let lfn = Value::Fn(Rc::new(OnceCell::new()));
+                let lfn = Value::Fn(Rc::new(LFn::Declared { name }));
                 self.emit_constant(lfn);
                 self.bind(name);
             }

src/main.rs

@@ -48,7 +48,7 @@ pub fn run(src: &'static str) {
     // in any event, the AST should live forever
     let parsed: &'static Spanned<Ast> = Box::leak(Box::new(parse_result.unwrap()));
 
-    let mut compiler = Compiler::new(parsed, "test", src);
+    let mut compiler = Compiler::new(parsed, "test", src, None);
     compiler.compile();
     if DEBUG_COMPILE {
         println!("=== source code ===");
@@ -76,9 +76,12 @@ pub fn run(src: &'static str) {
 pub fn main() {
     env::set_var("RUST_BACKTRACE", "1");
     let src = "
-fn foo () -> fn () -> fn () -> #{:foo :bar}
+let foo = {
+    let bar = :bar
+    fn baz () -> bar
+}
 
-foo () () () :foo
+foo ()
     ";
     run(src);
 }

src/value.rs

@@ -3,24 +3,84 @@ use crate::compiler::Chunk;
 use crate::parser::Ast;
 use crate::spans::Spanned;
 use imbl::{HashMap, Vector};
-use std::cell::{OnceCell, RefCell};
+use std::cell::RefCell;
 use std::rc::Rc;
 
+// #[derive(Clone, Debug, PartialEq)]
+// pub struct Defined {
+//     pub name: &'static str,
+//     pub doc: Option<&'static str>,
+//     // pub enclosing: Vec<(usize, Value)>,
+//     // pub has_run: bool,
+//     // pub input: &'static str,
+//     // pub src: &'static str,
+//     pub chunks: Vec<(u8, Chunk)>,
+//     pub closed: Vec<Value>,
+// }
+
+// impl Defined {
+//     pub fn close(&mut self, val: Value) {
+//         self.closed.push(val);
+//     }
+// }
+
 #[derive(Clone, Debug, PartialEq)]
-pub struct LFn {
-    pub name: &'static str,
-    pub doc: Option<&'static str>,
-    // pub enclosing: Vec<(usize, Value)>,
-    // pub has_run: bool,
-    // pub input: &'static str,
-    // pub src: &'static str,
-    pub chunks: Vec<(u8, Chunk)>,
-    pub closed: Vec<Value>,
+pub enum LFn {
+    Declared {
+        name: &'static str,
+    },
+    Defined {
+        name: &'static str,
+        doc: Option<&'static str>,
+        chunks: Vec<(u8, Chunk)>,
+        closed: RefCell<Vec<Value>>,
+    },
 }
 
 impl LFn {
-    pub fn close(&mut self, val: Value) {
-        self.closed.push(val);
+    pub fn close(&self, value: Value) {
+        match self {
+            LFn::Declared { .. } => unreachable!(),
+            LFn::Defined { closed, .. } => {
+                closed.borrow_mut().push(value);
+            }
+        }
+    }
+
+    pub fn doc(&self) -> Value {
+        match self {
+            LFn::Declared { name } => {
+                Value::String(Rc::new(format!("fn {name}: undefined function")))
+            }
+            LFn::Defined {
+                name,
+                doc: Some(doc),
+                ..
+            } => Value::String(Rc::new(format!("fn {name}\n{doc}"))),
+            LFn::Defined { name, .. } => {
+                Value::String(Rc::new(format!("fn {name}: no documentation")))
+            }
+        }
+    }
+
+    pub fn name(&self) -> &'static str {
+        match self {
+            LFn::Declared { name } | LFn::Defined { name, .. } => name,
+        }
+    }
+
+    pub fn chunk(&self, arity: u8) -> &Chunk {
+        match self {
+            LFn::Declared { .. } => unreachable!(),
+            LFn::Defined { chunks, .. } => &chunks.iter().find(|(a, _)| *a == arity).unwrap().1,
+        }
+    }
+
+    pub fn upvalue(&self, idx: u8) -> Value {
+        match self {
+            LFn::Declared { .. } => unreachable!(),
+            LFn::Defined { closed, .. } => closed.borrow()[idx as usize].clone(),
+        }
     }
 }
 
@@ -38,7 +98,7 @@ pub enum Value {
     List(Box<Vector<Value>>),
     Dict(Box<HashMap<&'static str, Value>>),
     Box(Rc<RefCell<Value>>),
-    Fn(Rc<OnceCell<LFn>>),
+    Fn(Rc<LFn>),
     BaseFn(BaseFn),
 }
 
@@ -103,7 +163,7 @@ impl std::fmt::Display for Value {
                     .join(", ")
             ),
             Box(value) => write!(f, "box {}", value.as_ref().borrow()),
-            Fn(lfn) => write!(f, "fn {}", lfn.get().unwrap().name),
+            Fn(lfn) => write!(f, "fn {}", lfn.name()),
             BaseFn(_) => write!(f, "base fn"),
         }
     }
@@ -144,7 +204,7 @@ impl Value {
                 format!("#{{{members}}}")
             }
             Box(x) => format!("box {{ {} }}", x.as_ref().borrow().show()),
-            Fn(lfn) => format!("fn {}", lfn.get().unwrap().name),
+            Fn(lfn) => format!("fn {}", lfn.name()),
             BaseFn(_) => "base fn".to_string(),
             Nothing => unreachable!(),
         }
@@ -189,7 +249,7 @@ impl Value {
             }
             String(s) => s.as_ref().clone(),
             Box(x) => x.as_ref().borrow().stringify(),
-            Fn(lfn) => lfn.get().unwrap().name.to_string(),
+            Fn(lfn) => format!("fn {}", lfn.name()),
             _ => todo!(),
         }
     }

src/vm.rs

@@ -57,14 +57,7 @@ impl CallFrame {
         let Value::Fn(ref function) = self.function else {
             unreachable!()
         };
-        &function
-            .get()
-            .unwrap()
-            .chunks
-            .iter()
-            .find(|(arity, _)| *arity == self.arity)
-            .unwrap()
-            .1
+        &function.chunk(self.arity)
     }
 }
 
@@ -73,9 +66,13 @@ impl fmt::Display for CallFrame {
         let Value::Fn(ref function) = self.function else {
             unreachable!()
         };
-        let inner_fn = function.get().unwrap();
-        let name = inner_fn.name;
-        write!(f, "CallFrame: {name}/{} @ {}", self.arity, self.ip)
+        write!(
+            f,
+            "CallFrame: {}/{} @ {}",
+            function.name(),
+            self.arity,
+            self.ip
+        )
     }
 }
 
@@ -92,15 +89,13 @@ pub struct Vm {
 
 impl Vm {
     pub fn new(chunk: Chunk) -> Vm {
-        let cell = OnceCell::new();
-        let lfn = LFn {
+        let lfn = LFn::Defined {
             name: "user script",
             doc: None,
             chunks: vec![(0, chunk)],
-            closed: vec![],
+            closed: RefCell::new(vec![]),
         };
-        let _ = cell.set(lfn);
-        let base_fn = Value::Fn(Rc::new(cell));
+        let base_fn = Value::Fn(Rc::new(lfn));
         let base_frame = CallFrame {
             function: base_fn.clone(),
             stack_base: 0,
@@ -706,6 +701,28 @@ impl Vm {
                     swap(&mut self.return_register[0], &mut value);
                     self.push(value);
                 }
+                Print => {
+                    println!("{}", self.pop().show());
+                    self.push(Value::Keyword("ok"));
+                    self.ip += 1;
+                }
+                SetUpvalue => {
+                    let idx = self.chunk().bytecode[self.ip + 1];
+                    self.ip += 2;
+                    let closed_over = self.stack[idx as usize].clone();
+                    if let Value::Fn(lfn) = self.peek() {
+                        lfn.close(closed_over);
+                    }
+                }
+                GetUpvalue => {
+                    let idx = self.chunk().bytecode[self.ip + 1];
+                    self.ip += 2;
+                    if let Value::Fn(ref inner) = self.frame.function {
+                        self.push(inner.as_ref().upvalue(idx));
+                    } else {
+                        unreachable!();
+                    }
+                }
             }
         }
     }