Implement upvalues

examples/test.lua

@@ -1,6 +1,8 @@
 function add(x)
-    return x
+    return function test() 
+        return x
+    end
 end
 
 global d = add(5)
-global c = print(d)
+global c = print(d())

src/compile.rs

@@ -25,6 +25,7 @@ impl State {
 pub struct Scope {
     pub locals: HashMap<String, u16>,
     pub register_counter: LocalCounter,
+    pub upvalues: HashMap<String, u16>,
 }
 
 #[derive(Clone, Debug, Default)]
@@ -165,6 +166,9 @@ impl Expression {
             Expression::ValueRef(name) => {
                 if let Some(local) = scope.locals.get(name) {
                     (Vec::new(), vec![*local])
+                } else if let Some(upvalue) = scope.upvalues.get(name) {
+                    let local = scope.register_counter.next();
+                    (vec![Instruction::GetUpVal(local, *upvalue)], vec![local])
                 } else {
                     let mut instructions = Vec::new();
                     let reg = scope.register_counter.next();
@@ -184,6 +188,13 @@ impl Expression {
                         .insert(param.kind.clone(), inner_scope.register_counter.next());
                 }
 
+                let highest_upvalue = scope.upvalues.iter().map(|(_, v)| *v).max().unwrap_or(0);
+                inner_scope.upvalues = scope.upvalues.clone();
+
+                for (name, reg) in &scope.locals {
+                    inner_scope.upvalues.insert(name.clone(), *reg + 1);
+                }
+
                 let instructions = block.compile(state, &mut inner_scope);
                 state.prototypes.push(instructions);
 

src/vm.rs

@@ -32,11 +32,16 @@ pub enum Instruction {
     SetGlobal(u16, u32),
     /// R(A) := G[K(Bx)]
     GetGlobal(u16, u32),
+    /// R(A) := U[B]
+    GetUpVal(u16, u16),
     /// [func] [params.len()] [ret_regs.len()]
     /// R(A), ... R(A+C-2) := R(A)(R(A+1), ... R(A+B-1))
     Call(u16, u16, u16),
+    /// return R(A), ... , R(B)
     Return(u16, u16),
+    /// close stack variables up to R(A)
     Close(u16),
+    /// R(A) := closure(KPROTO[Bx], R(A), ..., R(A+n))
     Closure(u16, u32),
 }
 
@@ -47,6 +52,7 @@ impl Debug for Instruction {
             Instruction::LoadK(arg0, arg1) => write!(f, "LOADK {} {}", arg0, arg1),
             Instruction::SetGlobal(arg0, arg1) => write!(f, "SETGLOBAL {} {}", arg0, arg1),
             Instruction::GetGlobal(arg0, arg1) => write!(f, "GETGLOBAL {} {}", arg0, arg1),
+            Instruction::GetUpVal(arg0, arg1) => write!(f, "GETUPVAL {} {}", arg0, arg1),
             Instruction::Call(arg0, arg1, arg2) => write!(f, "CALL {} {} {}", arg0, arg1, arg2),
             Instruction::Close(arg0) => write!(f, "CLOSE {}", arg0),
             Instruction::Closure(arg0, arg1) => write!(f, "CLOSURE {} {}", arg0, arg1),
@@ -240,13 +246,29 @@ impl ClosureRunner {
                 }
                 Instruction::Close(_) => {}
                 Instruction::Closure(reg, protok) => {
+                    let highest_upvalue = self
+                        .closure
+                        .upvalues
+                        .iter()
+                        .map(|(v, _)| *v)
+                        .max()
+                        .unwrap_or(0);
+
+                    let mut upvalues = self.closure.upvalues.clone();
+                    for (reg, value) in &self.stack {
+                        upvalues.insert(
+                            reg + highest_upvalue + 1,
+                            Rc::new(RefCell::new(value.clone())),
+                        );
+                    }
+
                     self.stack.insert(
                         *reg,
                         Value::Function(Closure {
                             vm: self.closure.vm.clone(),
                             prototype: *protok,
                             environment: self.closure.environment.clone(),
-                            upvalues: HashMap::new(),
+                            upvalues,
                         }),
                     );
                 }
@@ -256,8 +278,20 @@ impl ClosureRunner {
                     for i in *reg_start..=*reg_end {
                         ret_values.push(self.stack.get(&i).cloned().unwrap_or(Value::Nil));
                     }
+                    dbg!(&self.closure.upvalues);
                     return Some(ret_values);
                 }
+                Instruction::GetUpVal(reg, upvalreg) => {
+                    self.stack.insert(
+                        *reg,
+                        self.closure
+                            .upvalues
+                            .get(upvalreg)
+                            .unwrap()
+                            .borrow()
+                            .clone(),
+                    );
+                }
             };
 
             self.program_counter += 1;