From 70df6cfebe080e4ba233b1a409b38e96559dc92b Mon Sep 17 00:00:00 2001
From: Mustafa Quraish <mustafaq9@gmail.com>
Date: Fri, 26 Apr 2024 20:14:18 -0400
Subject: [PATCH] Closures!

---
 compiler/bytecode.oc      |  87 ++++++++----
 compiler/compiler.oc      | 286 +++++++++++++++++++++++++-------------
 compiler/lexer.oc         |   2 +-
 compiler/main.oc          |   9 +-
 compiler/parser.oc        |   2 +-
 compiler/vm/mod.oc        | 205 ++++++++++++++++++++++-----
 compiler/vm/value.oc      |  74 ++++++++--
 tests/bad/invalid_plus.td |   3 +
 8 files changed, 492 insertions(+), 176 deletions(-)
 create mode 100644 tests/bad/invalid_plus.td

diff --git a/compiler/bytecode.oc b/compiler/bytecode.oc
index 0c19ac2..ed52265 100644
--- a/compiler/bytecode.oc
+++ b/compiler/bytecode.oc
@@ -29,6 +29,15 @@ enum OpCode {
     GreaterThan
     Equal
     Halt
+    CloseFunction
+    SetUpvalue
+    GetUpvalue
+    CloseUpvalue
+}
+
+struct DebugLocRun {
+    span: Span
+    count: u32
 }
 
 struct Chunk {
@@ -39,7 +48,7 @@ struct Chunk {
 
     // Line number for each instruction
     // TODO: Run-length encoding or something
-    lines_numbers: &Vector<u16>
+    debug_locs: &Vector<DebugLocRun>
 }
 
 def Chunk::new(span: Span): &Chunk {
@@ -48,16 +57,50 @@ def Chunk::new(span: Span): &Chunk {
     chunk.code = Buffer::make()
     chunk.literals = Vector<Value>::new()
     chunk.literal_map = Map<Value, u32>::new()
-    chunk.lines_numbers = Vector<u16>::new()
+    chunk.debug_locs = Vector<DebugLocRun>::new()
     return chunk
 }
 
-def Chunk::push(&this, op: OpCode, span: Span) {
-    .code += op as u8
-    .lines_numbers += span.start.line as u16
+def Chunk::span_for_offset(&this, off: u32): Span {
+    let cur = 0
+    for it in .debug_locs.iter() {
+        cur += it.count
+        if cur > off return it.span
+    }
+    return .start_span
+}
+
+def Chunk::push_u8(&this, val: u8, span: Span) {
+    .code += val
+    if .debug_locs.size > 0 and .debug_locs.back().span == span {
+        .debug_locs.back_ptr().count += 1
+    } else {
+        .debug_locs.push(DebugLocRun(span, count: 1))
+    }
+}
+
+def Chunk::push_u16(&this, val: u16, span: Span) {
+    .code.write_u16(val)
+    if .debug_locs.size > 0 and .debug_locs.back().span == span {
+        .debug_locs.back_ptr().count += 2
+    } else {
+        .debug_locs.push(DebugLocRun(span, count: 2))
+    }
+}
+
+def Chunk::push_op(&this, op: OpCode, span: Span) => .push_u8(op as u8, span)
+
+def Chunk::push_with_arg_u16(&this, op: OpCode, arg: u16, span: Span) {
+    .push_op(op, span)
+    .push_u16(arg, span)
+}
+
+def Chunk::push_with_arg_u8(&this, op: OpCode, arg: u8, span: Span) {
+    .push_op(op, span)
+    .push_u8(arg, span)
 }
 
-def Chunk::push_with_value(&this, op: OpCode, value: Value, span: Span) {
+def Chunk::push_with_literal(&this, op: OpCode, value: Value, span: Span) {
     let it = .literal_map.get_item(value)
     let idx = match it? {
         true => it.value
@@ -72,20 +115,6 @@ def Chunk::push_with_value(&this, op: OpCode, value: Value, span: Span) {
     .push_with_arg_u16(op, idx as u16, span)
 }
 
-def Chunk::push_with_arg_u16(&this, op: OpCode, arg: u16, span: Span) {
-    .push(op, span)
-    .code.write_u16(arg)
-    // Push line number twice to keep the same size
-    .lines_numbers += span.start.line as u16
-    .lines_numbers += span.start.line as u16
-}
-
-def Chunk::push_with_arg_u8(&this, op: OpCode, arg: u8, span: Span) {
-    .push(op, span)
-    .code.write_u8(arg)
-    .lines_numbers += span.start.line as u16
-}
-
 def Chunk::literal(&this, idx: u16): Value {
     return .literals[idx as u32]
 }
@@ -116,8 +145,8 @@ def Chunk::disassemble_inst(&this, off: u32, found_chunks: &Vector<&Chunk> = nul
             constant.print();
             print("\n");
 
-            if constant.is_function() and found_chunks? {
-                let f = constant.as_function()
+            if constant.is_function_code() and found_chunks? {
+                let f = constant.as_function_code()
                 found_chunks.push(f.chunk)
             }
         }
@@ -147,14 +176,24 @@ def Chunk::disassemble_inst(&this, off: u32, found_chunks: &Vector<&Chunk> = nul
             let val = .read_literal(&off).as_obj() as &String
             println(f"{op}: {val.data}")
         }
-        GetLocal | SetLocal => {
+        GetLocal | SetLocal | SetUpvalue | GetUpvalue => {
             let idx = .read_u16(&off)
             println(f"{op}: {idx}")
         }
+        CloseFunction => {
+            let count = .read_u16(&off)
+            print(f"CloseFunction [{count}]: ")
+            for let i = 0u16; i < count; i += 1 {
+                if i > 0 print(" ")
+                print(`({.read_u8(&off) as bool}, {.read_u16(&off)})`)
+            }
+            print("\n")
+        }
         Null => println("Null")
         Pop => println("Pop")
         True => println("True")
         False => println("False")
+        CloseUpvalue => println("CloseUpvalue")
     }
     return off
 }
@@ -178,7 +217,7 @@ def Chunk::dump(&this) {
 def Chunk::free(&this) {
     .code.free()
     .literals.free()
-    .lines_numbers.free()
+    .debug_locs.free()
     .literal_map.free()
     mem::free(this)
 }
\ No newline at end of file
diff --git a/compiler/compiler.oc b/compiler/compiler.oc
index ad969f7..a070c61 100644
--- a/compiler/compiler.oc
+++ b/compiler/compiler.oc
@@ -6,31 +6,49 @@ import @ast::nodes::{ AST, Variable, Symbol }
 import @errors::{ Error }
 import @bytecode::{ Chunk, OpCode  }
 import @vm::{ VM, allocate_object }
-import @vm::value::{ Value, String, Function }
+import @vm::value::{ Value, String, FunctionCode }
 
 struct LocalVar {
     sym: &Symbol
     depth: i32
+    captured: bool
+}
+
+struct UpVar {
+    idx: u16
+    is_local: bool
 }
 
 struct Compiler {
     vm: &VM
+    func: &FunctionCode
     chunk: &Chunk
+    upvars: &Vector<UpVar>
     locals: &Vector<LocalVar>
     scope_depth: u32
+    enclosing: &Compiler
 }
 
-def Compiler::make(vm: &VM, span: Span): Compiler {
+def Compiler::make(func: &FunctionCode, vm: &VM, span: Span, enclosing: &Compiler = null): Compiler {
     let chunk = Chunk::new(span)
-    let locals = Vector<LocalVar>::new(capacity: 256)
+    let locals = Vector<LocalVar>::new()
+    let upvars = Vector<UpVar>::new(capacity: 256)
     return Compiler(
         vm,
-        chunk,
+        func,
+        func.chunk,
+        upvars,
         locals,
         scope_depth: 0,
+        enclosing
     )
 }
 
+def Compiler::free(&this) {
+    .locals.free()
+    .upvars.free()
+}
+
 def Compiler::make_str(&this, text: SV): Value {
     return .vm.copy_string(text.data, text.len)
 }
@@ -39,8 +57,12 @@ def Compiler::begin_scope(&this) => .scope_depth++
 def Compiler::end_scope(&this, span: Span) {
     .scope_depth--
     while .locals.size > 0 and .locals.back().depth as u32 > .scope_depth {
-        .chunk.push(Pop, Span(span.end, span.end))
-        .locals.pop()
+        let var = .locals.pop()
+        if var.captured {
+            .chunk.push_op(CloseUpvalue, span)
+        } else {
+            .chunk.push_op(Pop, span)
+        }
     }
 }
 
@@ -60,6 +82,44 @@ def Compiler::find_local(&this, name: SV, span: Span): i32 {
     return -1
 }
 
+def Compiler::add_upvar(&this, idx: u16, is_local: bool, span: Span): i32 {
+    for let i = 0; i < .upvars.size; i++ {
+        let up = .upvars[i]
+        if (up.idx == idx and up.is_local == is_local) {
+            return i as i32
+        }
+    }
+
+    let res = .upvars.size
+
+    if res > 0xffff {
+        Error::new(
+            span, "Compiler: Too many upvalues in function."
+        ).panic()
+    }
+
+    .upvars.push(UpVar(idx, is_local))
+    return res as i32
+}
+
+def Compiler::find_upvar(&this, name: SV, span: Span): i32 {
+    if not .enclosing? return -1
+
+    let local_idx = .enclosing.find_local(name, span)
+    if local_idx >= 0 {
+        let var = .enclosing.locals.at_ptr(local_idx as u32)
+        var.captured = true
+        return .add_upvar(local_idx as u16, true, span)
+    }
+
+    let up_idx = .enclosing.find_upvar(name, span)
+    if up_idx >= 0 {
+        return .add_upvar(up_idx as u16, false, span)
+    }
+
+    return -1
+}
+
 //! Create a new variable in the current scope
 def Compiler::create_variable(&this, sym: &Symbol) {
     // Global variable
@@ -68,7 +128,7 @@ def Compiler::create_variable(&this, sym: &Symbol) {
     }
 
     // Local variable
-    let local_var = LocalVar(sym, -1) // Uninitialized...
+    let local_var = LocalVar(sym, -1, false) // Uninitialized...
     .locals.push(local_var)
 }
 
@@ -88,42 +148,96 @@ def Compiler::define_variable(&this, sym: &Symbol) {
     if .scope_depth > 0 {
         return
     }
+
     // Global variable
-    .chunk.push_with_value(SetGlobal, .make_str(sym.name), sym.span)
-    .chunk.push(Pop, sym.span)
+    .chunk.push_with_literal(SetGlobal, .make_str(sym.name), sym.span)
+    .chunk.push_op(Pop, sym.span)
 }
 
-def Compiler::find_variable(&this, name: SV, span: Span) {
+def Compiler::compile_variable(&this, name: SV, span: Span) {
     let local_idx = .find_local(name, span)
-    // Global Variable
-    if local_idx < 0 {
-        .chunk.push_with_value(GetGlobal, .make_str(name), span)
-    
-    // Local Variable
-    } else {
+    if local_idx >= 0 {
+        // Local Variable
         .chunk.push_with_arg_u16(GetLocal, local_idx as u16, span)
+        return
+    }
+
+    let up_idx = .find_upvar(name, span)
+    if up_idx >= 0 {
+        // Upvalue
+        .chunk.push_with_arg_u16(GetUpvalue, up_idx as u16, span)
+        return
+    }
+
+    // Global Variable
+    .chunk.push_with_literal(GetGlobal, .make_str(name), span)
+}
+
+def Compiler::make_jump(&this, op: OpCode, span: Span): u32 {
+    .chunk.push_with_arg_u16(op, 0xbeef, span) // Placeholder
+    let patch_off = .chunk.code.size - 2
+    return patch_off
+}
+
+def Compiler::make_loop(&this, target: u32, span: Span) {
+    let offset = .chunk.code.size - target + 3
+    if offset > 0xffff {
+        Error::new(
+            span, "Compiler: Loop offset too large."
+        ).panic()
+    }
+
+    .chunk.push_with_arg_u16(Loop, offset as u16, span)
+}
+
+def Compiler::patch_jump(&this, from: u32, span: Span) {
+    let offset = .chunk.code.size - from - 2
+    if offset > 0xffff {
+        Error::new(
+            span, "Compiler: Jump offset too large."
+        ).panic()
+    }
+
+    .chunk.code.data[from + 0] = (offset >> 8) as u8 & 0xff
+    .chunk.code.data[from + 1] = offset as u8
+}
+
+def Compiler::compile_if(&this, node: &AST, is_expression: bool = false) {
+    let ifs = &node.u.if_stmt
+    .compile_expression(ifs.cond)
+
+    let false_jump = .make_jump(JumpIfFalse, ifs.cond.span)
+    .chunk.push_op(Pop, ifs.cond.span)
+    .compile_statement(ifs.body)
+    let end_jump = .make_jump(Jump, ifs.body.span)
+
+    .patch_jump(false_jump, ifs.cond.span)
+    .chunk.push_op(Pop, ifs.cond.span)
+    if ifs.els? {
+        .compile_statement(ifs.els)
     }
+    .patch_jump(end_jump, ifs.body.span)
 }
 
 def Compiler::compile_expression(&this, node: &AST) {
     match node.type {
         IntLiteral => {
             let value = Value::Int(node.u.num_literal.text.to_i32())
-            .chunk.push_with_value(Constant, value, node.span)
+            .chunk.push_with_literal(Constant, value, node.span)
         }
         FloatLiteral => {
             let value = Value::Float(std::libc::strtod(node.u.num_literal.text, null))
-            .chunk.push_with_value(Constant, value, node.span)
+            .chunk.push_with_literal(Constant, value, node.span)
         }
         StringLiteral => {
             let text = node.u.string_literal
-            .chunk.push_with_value(Constant, .make_str(text), node.span)
+            .chunk.push_with_literal(Constant, .make_str(text), node.span)
         }
         BoolLiteral => {
             let value = Value::Bool(node.u.bool_literal)
-            .chunk.push_with_value(Constant, value, node.span)
+            .chunk.push_with_literal(Constant, value, node.span)
         }
-        Identifier => .find_variable(node.u.ident.name, node.span)
+        Identifier => .compile_variable(node.u.ident.name, node.span)
         Call => {
             let callee = node.u.call.callee
 
@@ -145,14 +259,20 @@ def Compiler::compile_expression(&this, node: &AST) {
                     Identifier => {
                         let name = lhs.u.ident.name
                         let idx = .find_local(lhs.u.ident.name, lhs.span)
+                        let up_idx = .find_upvar(name, node.span)
                         
                         // Global variable
-                        if idx < 0 {
-                            .chunk.push_with_value(SetGlobal, .make_str(name), node.span)
+                        if idx >= 0 {
+                            .chunk.push_with_arg_u16(SetLocal, idx as u16, node.span)
+
+                        // Upvalue
+                        } else if up_idx >= 0 {
+                            .chunk.push_with_arg_u16(SetUpvalue, up_idx as u16, node.span)
 
                         // Local variable
                         } else {
-                            .chunk.push_with_arg_u16(SetLocal, idx as u16, node.span)
+                            .chunk.push_with_literal(SetGlobal, .make_str(name), node.span)
+
                         }
                     }
                     else => {
@@ -165,7 +285,7 @@ def Compiler::compile_expression(&this, node: &AST) {
             And => {
                 .compile_expression(node.u.binary.lhs)
                 let false_jump = .make_jump(JumpIfFalse, node.span)
-                .chunk.push(Pop, node.span)
+                .chunk.push_op(Pop, node.span)
                 .compile_expression(node.u.binary.rhs)
                 .patch_jump(false_jump, node.span)
             }
@@ -174,7 +294,7 @@ def Compiler::compile_expression(&this, node: &AST) {
                 let false_jump = .make_jump(JumpIfFalse, node.span)
                 let true_jump = .make_jump(Jump, node.span)
                 .patch_jump(false_jump, node.span)
-                .chunk.push(Pop, node.span)
+                .chunk.push_op(Pop, node.span)
                 .compile_expression(node.u.binary.rhs)
                 .patch_jump(true_jump, node.span)
             }
@@ -183,13 +303,13 @@ def Compiler::compile_expression(&this, node: &AST) {
                 .compile_expression(node.u.binary.lhs)
                 .compile_expression(node.u.binary.rhs)
                 match node.u.binary.op {
-                    Plus => .chunk.push(Add, node.u.binary.op_span),
-                    Minus => .chunk.push(Sub, node.u.binary.op_span),
-                    Multiply => .chunk.push(Mul, node.u.binary.op_span),
-                    Divide => .chunk.push(Div, node.u.binary.op_span),
-                    LessThan => .chunk.push(LessThan, node.u.binary.op_span),
-                    GreaterThan => .chunk.push(GreaterThan, node.u.binary.op_span),
-                    Equals => .chunk.push(Equal, node.u.binary.op_span),
+                    Plus => .chunk.push_op(Add, node.span),
+                    Minus => .chunk.push_op(Sub, node.span),
+                    Multiply => .chunk.push_op(Mul, node.span),
+                    Divide => .chunk.push_op(Div, node.span),
+                    LessThan => .chunk.push_op(LessThan, node.span),
+                    GreaterThan => .chunk.push_op(GreaterThan, node.span),
+                    Equals => .chunk.push_op(Equal, node.span),
                     else => std::panic(`Unimplemented binary operator: {node.u.binary.op}`)
                 }
             }
@@ -202,52 +322,6 @@ def Compiler::compile_expression(&this, node: &AST) {
     }
 }
 
-def Compiler::make_jump(&this, op: OpCode, span: Span): u32 {
-    .chunk.push_with_arg_u16(op, 0xbeef, span) // Placeholder
-    let patch_off = .chunk.code.size - 2
-    return patch_off
-}
-
-def Compiler::make_loop(&this, target: u32, span: Span) {
-    let offset = .chunk.code.size - target + 3
-    if offset > 0xffff {
-        Error::new(
-            span, "Compiler: Loop offset too large."
-        ).panic()
-    }
-
-    .chunk.push_with_arg_u16(Loop, offset as u16, span)
-}
-
-def Compiler::patch_jump(&this, from: u32, span: Span) {
-    let offset = .chunk.code.size - from - 2
-    if offset > 0xffff {
-        Error::new(
-            span, "Compiler: Jump offset too large."
-        ).panic()
-    }
-
-    .chunk.code.data[from + 0] = (offset >> 8) as u8 & 0xff
-    .chunk.code.data[from + 1] = offset as u8
-}
-
-def Compiler::compile_if(&this, node: &AST, is_expression: bool = false) {
-    let ifs = &node.u.if_stmt
-    .compile_expression(ifs.cond)
-
-    let false_jump = .make_jump(JumpIfFalse, ifs.cond.span)
-    .chunk.push(Pop, ifs.cond.span)
-    .compile_statement(ifs.body)
-    let end_jump = .make_jump(Jump, ifs.body.span)
-
-    .patch_jump(false_jump, ifs.cond.span)
-    .chunk.push(Pop, ifs.cond.span)
-    if ifs.els? {
-        .compile_statement(ifs.els)
-    }
-    .patch_jump(end_jump, ifs.body.span)
-}
-
 def Compiler::compile_statement(&this, node: &AST) {
     match node.type {
         Block => {
@@ -265,7 +339,7 @@ def Compiler::compile_statement(&this, node: &AST) {
             if expr? {
                 .compile_expression(expr)
             } else {
-                .chunk.push(Null, node.span)
+                .chunk.push_op(Null, node.span)
             }
             .create_variable(sym)
             .define_variable(sym)
@@ -278,9 +352,11 @@ def Compiler::compile_statement(&this, node: &AST) {
             .mark_variable_initialized()
 
             let name_val = .make_str(ast_func.sym.name)
-            let func_obj = allocate_object<Function>(.vm, Function)
+            let func_obj = allocate_object<FunctionCode>(.vm, FunctionCode)
+            let chunk = Chunk::new(ast_func.span)
+            func_obj.init(name_val.as_string(), chunk, ast_func.params.size as u8)
 
-            let cc = Compiler::make(.vm, node.span)
+            let cc = Compiler::make(func_obj, .vm, node.span, enclosing: this)
             cc.begin_scope()
             for param in ast_func.params.iter() {
                 cc.create_variable(param.sym)
@@ -290,16 +366,21 @@ def Compiler::compile_statement(&this, node: &AST) {
             cc.compile_statement(ast_func.body)
 
             // Add a placeholder return statement if none was provided
-            cc.chunk.push_with_value(Constant, Value::Null(), node.span)
-            cc.chunk.push(Return, node.span)
+            cc.chunk.push_with_literal(Constant, Value::Null(), node.span)
+            cc.chunk.push_op(Return, node.span)
             cc.end_scope(ast_func.sym.span)
 
-            func_obj.init(name_val.as_string(), cc.chunk, ast_func.params.size as u8)
-
             let func_val = Value::Object(&func_obj.obj)
-            .chunk.push_with_value(Constant, func_val, node.span)
+            .chunk.push_with_literal(Constant, func_val, node.span)
+
+            .chunk.push_with_arg_u16(CloseFunction, cc.upvars.size as u16, node.span)
+            for up in cc.upvars.iter() {
+                .chunk.push_u8(up.is_local as u8, node.span)
+                .chunk.push_u16(up.idx, node.span)
+            }
 
             .define_variable(ast_func.sym)
+            cc.free()
         }
         If => .compile_if(node, is_expression: false)
         While => {
@@ -307,13 +388,13 @@ def Compiler::compile_statement(&this, node: &AST) {
             let start = .chunk.code.size
             .compile_expression(loop.cond)
             let false_jump = .make_jump(JumpIfFalse, loop.cond.span)
-            .chunk.push(Pop, loop.cond.span)
+            .chunk.push_op(Pop, loop.cond.span)
 
             .compile_statement(loop.body)
 
             .make_loop(start, loop.cond.span)
             .patch_jump(false_jump, loop.cond.span)
-            .chunk.push(Pop, loop.cond.span)
+            .chunk.push_op(Pop, loop.cond.span)
         }
         For => {
             let loop = &node.u.loop
@@ -327,11 +408,11 @@ def Compiler::compile_statement(&this, node: &AST) {
             if loop.cond? {
                 .compile_expression(loop.cond)
             } else {
-                .chunk.push_with_value(Constant, Value::True(), node.span)
+                .chunk.push_with_literal(Constant, Value::True(), node.span)
             }
 
             let false_jump = .make_jump(JumpIfFalse, node.span)
-            .chunk.push(Pop, node.span)
+            .chunk.push_op(Pop, node.span)
 
             .compile_statement(loop.body)
 
@@ -341,7 +422,7 @@ def Compiler::compile_statement(&this, node: &AST) {
 
             .make_loop(start, node.span)
             .patch_jump(false_jump, node.span)
-            .chunk.push(Pop, node.span)
+            .chunk.push_op(Pop, node.span)
 
             .end_scope(node.span)
         }
@@ -349,13 +430,13 @@ def Compiler::compile_statement(&this, node: &AST) {
             if node.u.child? {
                 .compile_expression(node.u.child)
             } else {
-                .chunk.push(Null, node.span)
+                .chunk.push_op(Null, node.span)
             }
-            .chunk.push(Return, node.span)
+            .chunk.push_op(Return, node.span)
         }
         else => {
             .compile_expression(node)
-            .chunk.push(Pop, node.span)
+            .chunk.push_op(Pop, node.span)
         }
     }
 }
@@ -368,9 +449,16 @@ def Compiler::compile_ns(&this, root: &AST) {
 }
 
 // TODO: More complex structures than just expressions?
-def compile_program(vm: &VM, root: &AST): &Chunk {
-    let compiler = Compiler::make(vm, root.span)
+def compile_program(vm: &VM, root: &AST): &FunctionCode {
+    let func = allocate_object<FunctionCode>(vm, FunctionCode)
+    let func_name = vm.copy_string("", 0)  // Empty string for script
+    let chunk = Chunk::new(root.span)
+    func.init(func_name.as_string(), chunk, 0)
+
+    let compiler = Compiler::make(func, vm, root.span)
     compiler.compile_ns(root)
-    compiler.chunk.push(Halt, root.span)
-    return compiler.chunk
+    compiler.chunk.push_op(Halt, root.span)
+    compiler.free()
+
+    return func
 }
diff --git a/compiler/lexer.oc b/compiler/lexer.oc
index ad0d523..e6b3828 100644
--- a/compiler/lexer.oc
+++ b/compiler/lexer.oc
@@ -117,7 +117,7 @@ def Lexer::lex_string_literal(&this, has_seen_f: bool): Token {
     let text = SV(.source + start, len)
     .inc()
 
-    if .i >= .source_len {
+    if .i > .source_len {
         .errors.push(Error::new(Span(.loc, .loc), "Unterminated string literal"))
     }
 
diff --git a/compiler/main.oc b/compiler/main.oc
index d20f679..1cc811a 100644
--- a/compiler/main.oc
+++ b/compiler/main.oc
@@ -94,18 +94,15 @@ def main(argc: i32, argv: &str) {
     mem::reset_to_default_allocator()
 
     let vm = VM::make()
-    let chunk = compile_program(&vm, ast)
+    let script = compile_program(&vm, ast)
 
     if debug {
-        chunk.dump()
+        script.chunk.dump()
         println("============================================")
         log(Info, "Running VM:")
     }
 
-    let status = vm.run(chunk)
-
-    chunk.free()
+    let status = vm.run(script)
     vm.free()
-
     std::exit(status)
 }
diff --git a/compiler/parser.oc b/compiler/parser.oc
index eec9b9a..e8bfd82 100644
--- a/compiler/parser.oc
+++ b/compiler/parser.oc
@@ -1029,6 +1029,7 @@ def Parser::parse_statement(&this): &AST {
 
     match .token().type {
         OpenCurly => node = .parse_block()
+        Def => node = .parse_function()
         Return => {
             .consume(Return)
             let expr = null as &AST
@@ -1444,7 +1445,6 @@ def Parser::parse_namespace_until(&this, end_type: TokenType) {
                 still_parsing_attributes = true
                 yield null as &AST
             }
-            Def => .parse_function()
             // Import => {
             //     let import_ = .parse_import()
             //     if import_? then .ns.imports.push(import_)   // For typechecker...
diff --git a/compiler/vm/mod.oc b/compiler/vm/mod.oc
index 135b8c0..e8a6590 100644
--- a/compiler/vm/mod.oc
+++ b/compiler/vm/mod.oc
@@ -8,7 +8,11 @@ import std::logging::{ log }
 import @errors::{ Error }
 import @bytecode::{ Chunk, OpCode }
 import @main::debug
-import .value::{ Value, String, Object, ObjectType, NativeFunctionType, NativeFunction }
+import .value::{ 
+    Value, String, Object, ObjectType, 
+    NativeFunctionType, NativeFunction, FunctionCode,
+    UpValue, Function
+}
 import .native
 
 //! Wrapper around a string that forces comparison by value, not
@@ -28,6 +32,7 @@ def ValueCompareString::eq(this, other: ValueCompareString): bool {
 }
 
 struct CallFrame {
+    func: &Function
     chunk: &Chunk
     ip: &u8
     stack_base: u32
@@ -37,13 +42,17 @@ struct VM {
     frames: &Vector<CallFrame>
     
     // For the current frame
+    func: &Function
     chunk: &Chunk
     ip: &u8
     stack_base: u32
 
+    cur_inst_ip: &u8
+
     stack: &Vector<Value>
     globals: &Map<&String, Value>
     strings: &Map<ValueCompareString, Value>
+    open_upvalues: &UpValue  // Linked list of open upvalues
 
     //! Linked list of all allocated objects
     objects: &Object
@@ -52,12 +61,15 @@ struct VM {
 def VM::make(): VM {
     return VM(
         frames: Vector<CallFrame>::new(),
+        func: null,
         chunk: null,
         ip: null,
         stack_base: 0,
+        cur_inst_ip: null,
         stack: Vector<Value>::new(),
         globals: Map<&String, Value>::new(),
         strings: Map<ValueCompareString, Value>::new(),
+        open_upvalues: null,
         objects: null
     )
 }
@@ -119,6 +131,12 @@ def VM::copy_string(&this, data: str, len: u32): Value {
     return value
 }
 
+def VM::make_function(&this, code: &FunctionCode): &Function {
+    let func = allocate_object<Function>(this, ObjectType::Function)
+    func.init(code)
+    return func
+}
+
 def VM::add(&this, a: Value, b: Value): Value => if {
     a.is_float() and b.is_float() => Value::Float(a.as_float() + b.as_float())
     a.is_float() and b.is_int() => Value::Float(a.as_float() + b.as_int() as f64)
@@ -188,8 +206,30 @@ def VM::print_value(&this, val: Value) => if {
     val.is_null() => print("null")
     val.is_obj() => match val.as_obj().type {
         String => print(`{val.as_string().data}`)
-        Function => print(`<function {val.as_function().name.data}>`)
+        Function => {
+            let clos = val.as_function()
+            let name = val.as_function_code().name.data
+
+            if {
+                clos.upvalues? => print(`<func "{name}, {clos.upvalues.size} captures">`)
+                name == "" => print("<script>")
+                else => print(`<function "{name}">`)
+            }
+        }
+        FunctionCode => {
+            let name = val.as_function_code().name.data
+            if name == "" {
+                print("<script>")
+            } else {
+                print(`<function "{name}">`)
+            }
+        }
         NativeFunction => print(`<native {val.as_native_function().name.data}>`)
+        UpValue => {
+            let slot = val.as_upvalue().slot
+            let upval = .stack.data[slot]
+            .print_value(upval)
+        }
     }
     else => .error(`Unhandled value in VM::print(): {val.u.uint:08x}`)
 }
@@ -223,12 +263,16 @@ def VM::setup_native_functions(&this) {
     .add_native_function("print", native::print)
 }
 
-def VM::run(&this, chunk: &Chunk): i32 {
+def VM::run(&this, code: &FunctionCode): i32 {
     .setup_native_functions()
-    .stack.push(Value::Null())
-    .chunk = chunk
-    .ip = chunk.code.data
-    .stack_base = 1
+
+    let func = .make_function(code)
+
+    .stack.push(Value::Object(&func.obj))
+    .func = func
+    .chunk = code.chunk
+    .ip = .chunk.code.data
+    .stack_base = .stack.size
     return .main_loop()
 }
 
@@ -240,11 +284,12 @@ def VM::is_falsey(&this, val: Value): bool => if {
     else => false  // TODO: Custom falsey values?
 }
 
-def VM::push_frame(&this, chunk: &Chunk, ip: &u8, stack_base: u32) {
+def VM::push_frame(&this, func: &Function, stack_base: u32) {
     // Push current frame, and set up the new one
-    .frames.push(CallFrame(.chunk, .ip, .stack_base))
-    .chunk = chunk
-    .ip = ip
+    .frames.push(CallFrame(.func, .chunk, .ip, .stack_base))
+    .func = func
+    .chunk = func.code.chunk
+    .ip = .chunk.code.data
     .stack_base = stack_base
 }
 
@@ -253,40 +298,73 @@ def VM::pop_frame(&this) {
     .stack.size = .stack_base - 1
 
     let frame = .frames.pop()
+    .func = frame.func
     .chunk = frame.chunk
     .ip = frame.ip
     .stack_base = frame.stack_base
 }
 
 [exits]
-def VM::error(&this, msg: str) {
-    let off = (.ip - .chunk.code.data) as u32 - 1  // -1 because we read the opcode
-    let line = .chunk.lines_numbers[off]
-    let err_loc = .chunk.start_span.start
-    err_loc.line = line as u32
-    err_loc.col = 1
-    let span = Span(err_loc, err_loc)
+def VM::error_at(&this, span: Span, msg: str) {
     Error::new(span, msg).panic()
 }
 
+[exits]
+def VM::error(&this, msg: str) {
+    let inst_off = (.cur_inst_ip - .chunk.code.data) as u32
+    let span = .chunk.span_for_offset(inst_off)
+    .error_at(span, msg)
+}
+
+def VM::debug_dump_inst(&this, print_stack: bool) {
+    if print_stack {
+        print("\nStack: ()")
+        for let i = 1; i < .stack.size; i++ {
+            let val = .stack[i]
+            if i > 0 then print(" ")
+            val.print()
+        }
+        println("")
+    }
+
+    let off = (.ip - .chunk.code.data) as u32
+    let span = .chunk.span_for_offset(off)
+    print(`[VM] (L#{span.start.line:2d}) {off:4d}: `)
+    .chunk.disassemble_inst(off)
+}
+
+def VM::capture_upvalue(&this, slot: u32): &UpValue {
+    let cur = &.open_upvalues
+    while (*cur)? and (*cur).slot > slot {
+        (*cur) = (*cur).next
+    }
+
+    if (*cur)? and (*cur).slot == slot {
+        return *cur
+    }
+
+    let upval = allocate_object<UpValue>(this, ObjectType::UpValue)
+    upval.init(slot)
+
+    (*cur) = upval
+    return upval
+}
+
+def VM::close_upvalue(&this, slot: u32) {
+    while .open_upvalues? and .open_upvalues.slot >= slot {
+        let upval = .open_upvalues
+        let val = .stack.data[upval.slot]
+        upval.closed = val
+        upval.slot = 0
+        .open_upvalues = upval.next
+    }
+}
+
 def VM::main_loop(&this): i32 {
-    // TODO: branching instructions?
     while true {
-        if debug {
-            print("\nStack: ")
-            for let i = 0; i < .stack.size; i++ {
-                let val = .stack[i]
-                if i > 0 then print(", ")
-                .print_value(val)
-            }
-            println("")
+        if debug then .debug_dump_inst(print_stack: true)
 
-            let off = (.ip - .chunk.code.data) as u32
-            let line = .chunk.lines_numbers[off]
-            print(`[VM] (L#{line:2d}) {off:4d}: `)
-            .chunk.disassemble_inst(off)
-
-        }
+        .cur_inst_ip = .ip
         let op = .read_u8() as OpCode
         match op {
             Halt => return 0
@@ -353,6 +431,7 @@ def VM::main_loop(&this): i32 {
                     return value.as_int()
 
                 } else {
+                    .close_upvalue(.stack_base)
                     .pop_frame()
                     .stack.push(value)
                 }
@@ -361,14 +440,14 @@ def VM::main_loop(&this): i32 {
                 let arity = .read_u8()
                 let val = .stack.back(arity as u32)
 
-                // TODO: Location?
                 if {
                     val.is_function() => {
                         let func = val.as_function()
-                        if func.arity != arity {
-                            .error(f"Function {func.name.data} expected {func.arity} arguments, got {arity}")
+                        let code = func.code
+                        if code.arity != arity {
+                            .error(f"Function {code.name.data} expected {code.arity} arguments, got {arity}")
                         }
-                        .push_frame(func.chunk, func.chunk.code.data, .stack.size - arity as u32)
+                        .push_frame(func, .stack.size - arity as u32)
                     }
                     val.is_native_function() => {
                         let func = val.as_native_function()
@@ -376,6 +455,7 @@ def VM::main_loop(&this): i32 {
                         .stack.size -= arity as u32 + 1  // Pop the arguments and the function
                         .stack.push(res)
                     }
+                    val.is_function_code() => .error("Function should have been a func")
                     else => .error(f"Can only call functions, got {val.type_str()}")
                 }
             }
@@ -403,6 +483,57 @@ def VM::main_loop(&this): i32 {
             Null => .stack.push(Value::Null())
             True => .stack.push(Value::True())
             False => .stack.push(Value::False())
+            CloseFunction => {
+                let arg = .stack.pop()
+                if not arg.is_function_code() {
+                    .error(f"Expected function, got {arg.type_str()}")
+                }
+                let code = arg.as_function_code()
+                let func = .make_function(code)
+
+                let num_upvals = .read_u16() as u32
+                if num_upvals > 0 {
+                    func.upvalues = Vector<&UpValue>::new()
+                }
+                for let i = 0; i < num_upvals; i++ {
+                    let is_local = .read_u8() == 1
+                    let index = .read_u16() as u32
+                    if is_local {
+                        let upval = .capture_upvalue(.stack_base + index)
+                        func.upvalues += upval
+                        // print(`  [VM] Upval {index} is local: {upval.slot} ({upval})\n`)
+
+                    } else {
+                        let upval = .func.upvalues[index]
+                        func.upvalues += upval
+                        // print(`  [VM] Upval {index} is outside: {upval.slot}\n`)
+                    }
+                }
+                .stack.push(Value::Object(&func.obj))
+            }
+            GetUpvalue => {
+                let idx = .read_u16() as u32
+                let upval = .func.upvalues[idx]
+                // print(`  [VM] Getting upval {idx} ({upval})\n`)
+                if upval.slot == 0 {
+                    .stack.push(upval.closed)
+                } else {
+                    .stack.push(.stack.data[upval.slot])
+                }
+            }
+            SetUpvalue => {
+                let idx = .read_u16() as u32
+                let upval = .func.upvalues[idx]
+                if upval.slot == 0 {
+                    upval.closed = .stack.back()
+                } else {
+                    .stack.data[upval.slot] = .stack.back()
+                }
+            }
+            CloseUpvalue => {
+                .close_upvalue(.stack.size - 1)
+                .stack.pop()
+            }
         }
     }
     return 0
diff --git a/compiler/vm/value.oc b/compiler/vm/value.oc
index 2d3f7de..5ebfb24 100644
--- a/compiler/vm/value.oc
+++ b/compiler/vm/value.oc
@@ -60,12 +60,18 @@ def Value::is_obj_type(this, type: ObjectType): bool {
 def Value::is_string(this): bool => .is_obj_type(String)
 def Value::as_string(this): &String => .as_obj() as &String
 
-def Value::is_function(this): bool => .is_obj_type(Function)
-def Value::as_function(this): &Function => .as_obj() as &Function
+def Value::is_function_code(this): bool => .is_obj_type(FunctionCode)
+def Value::as_function_code(this): &FunctionCode => .as_obj() as &FunctionCode
 
 def Value::is_native_function(this): bool => .is_obj_type(NativeFunction)
 def Value::as_native_function(this): &NativeFunction => .as_obj() as &NativeFunction
 
+def Value::is_upvalue(this): bool => .is_obj_type(UpValue)
+def Value::as_upvalue(this): &UpValue => .as_obj() as &UpValue
+
+def Value::is_function(this): bool => .is_obj_type(Function)
+def Value::as_function(this): &Function => .as_obj() as &Function
+
 def Value::Float(f: f64): Value => Value::from_float(f)
 def Value::True(): Value => Value::from_uint(BOOL_MASK | 3)
 def Value::False(): Value => Value::from_uint(BOOL_MASK | 2)
@@ -87,7 +93,9 @@ def Value::type_str(this): str => if {
         yield match obj.type {
             String => "string"
             Function => "function"
+            FunctionCode => "code"
             NativeFunction => "<native function>"
+            UpValue => "upvalue"
         }
     }
     else => "<unknown>"
@@ -107,8 +115,12 @@ def Value::print(this) => if {
 
 enum ObjectType {
     String
-    Function
     NativeFunction
+    UpValue
+    //! Representing the function code itself, not a closure
+    FunctionCode
+    //! Representing a function with captured variables
+    Function
 }
 
 //! Generic Heap-allocated object. Should only be
@@ -125,11 +137,17 @@ def Object::free(&this) {
             s.data.free()
         }
         Function => {
-            let f = this as &Function
+            let c = this as &Function
+            if c.upvalues? {
+                c.upvalues.free()
+            }
+        }
+        FunctionCode => {
+            let f = this as &FunctionCode
             f.chunk.free()
             // NOTE: A function does not own its name
         }
-        NativeFunction => {}
+        NativeFunction | UpValue => {}
     }
     mem::free(this)
 }
@@ -141,13 +159,26 @@ def Object::print(&this) {
             print(`"{s.data}"`)
         }
         Function => {
-            let f = this as &Function
-            print(`<function {f.name.data}>`)
+            let c = this as &Function
+            let name = c.code.name.data
+            if {
+                name == "" => print("<script>")
+                not c.upvalues? => print(f"<function {name}>")
+                else => print(f"<function {name}, {c.upvalues.size} captures>")
+            }
+        }
+        FunctionCode => {
+            let f = this as &FunctionCode
+            print(`<code "{f.name.data}">`)
         }
         NativeFunction => {
             let nf = this as &NativeFunction
             print(`<native function {nf.name.data}>`)
         }
+        UpValue => {
+            let uv = this as &UpValue
+            print(f"<upvalue {uv.slot}>")
+        }
     }
 }
 
@@ -169,7 +200,32 @@ def String::init(&this, data: str, len: u32, shash: u32) {
     .shash = shash
 }
 
+struct UpValue {
+    obj: Object
+    slot: u32
+    next: &UpValue
+    closed: Value
+}
+
+def UpValue::init(&this, slot: u32) {
+    .slot = slot
+    .next = null
+    .closed = Value::Null()
+}
+
 struct Function {
+    obj: Object
+    code: &FunctionCode
+    upvalues: &Vector<&UpValue>
+}
+
+def Function::init(&this, code: &FunctionCode) {
+    .code = code
+    .upvalues = null
+}
+
+//! Representing the function itself, not a closure
+struct FunctionCode {
     obj: Object
     name: &String
     chunk: &Chunk
@@ -177,10 +233,12 @@ struct Function {
 }
 
 //! Expects `obj` to already be initialized
-def Function::init(&this, name: &String, chunk: &Chunk, arity: u8) {
+def FunctionCode::init(&this, name: &String, chunk: &Chunk, arity: u8) {
     .name = name
     .chunk = chunk
     .arity = arity
+    // this is owned by the function, so the GC doesn't need to
+    // know about it
 }
 
 typedef NativeFunctionType = fn(&VM, u32, &Value): Value
diff --git a/tests/bad/invalid_plus.td b/tests/bad/invalid_plus.td
new file mode 100644
index 0000000..000d6c9
--- /dev/null
+++ b/tests/bad/invalid_plus.td
@@ -0,0 +1,3 @@
+/// fail: Cant add int and string
+
+1 + "Hello"
\ No newline at end of file