Calling of functions, in functions

2 files changed, 46 insertions, 13 deletions

diff --git a/src/compiler.rs b/src/compiler.rs
index b7dad75..ad97783 100644
--- a/src/compiler.rs
+++ b/src/compiler.rs
@@ -125,7 +125,6 @@ struct Frame {
 
 impl Frame {
     fn find_local(&self, name: &str) -> Option<Variable> {
-        println!("LOCAL!");
         for var in self.stack.iter().rev() {
             if var.name == name && var.active {
                 return Some(var.clone());
@@ -135,7 +134,6 @@ impl Frame {
     }
 
     fn find_upvalue(&self, name: &str) -> Option<Variable> {
-        println!("UPVALUE!");
         for var in self.upvalues.iter().rev() {
             if var.name == name && var.active {
                 return Some(var.clone());
@@ -202,6 +200,7 @@ macro_rules! push_scope {
         $code;
 
         $compiler.frame_mut().scope -= 1;
+
         for var in $compiler.frame().stack[ss..$compiler.stack().len()].iter().rev() {
             if var.captured {
                 $block.add(Op::PopUpvalue, $compiler.line());
@@ -496,17 +495,15 @@ impl Compiler {
         }
 
         block.add(Op::Call(arity), self.line());
-
-        for _ in 0..arity {
-            block.add(Op::Pop, self.line());
-        }
     }
 
     fn function(&mut self, block: &mut Block) {
         expect!(self, Token::Fn, "Expected 'fn' at start of function.");
 
-        let name = if !self.stack()[self.stack().len() - 1].active {
-            &self.stack()[self.stack().len() - 1].name
+        let top = self.stack().len() - 1;
+        let name = if !self.stack()[top].active {
+            self.stack_mut()[top].active = true;
+            &self.stack()[top].name
         } else {
             "anonumus function"
         };
@@ -562,9 +559,18 @@ impl Compiler {
             // so we know from where to copy them.
         });
 
-        if !matches!(function_block.last_op(), Some(&Op::Return)) {
-            function_block.add(Op::Constant(Value::Nil), self.line());
-            function_block.add(Op::Return, self.line());
+        let mut prev = function_block.ops.len() - 1;
+        loop {
+            match function_block.ops[prev] {
+                Op::Pop | Op::PopUpvalue => {}
+                Op::Return => { break; } ,
+                _ => {
+                    function_block.add(Op::Constant(Value::Nil), self.line());
+                    function_block.add(Op::Return, self.line());
+                    break;
+                }
+            }
+            prev -= 1;
         }
 
         function_block.ty = Type::Function(args, Box::new(return_type));
diff --git a/src/vm.rs b/src/vm.rs
index abe7874..cb2285f 100644
--- a/src/vm.rs
+++ b/src/vm.rs
@@ -369,7 +369,29 @@ impl VM {
             }
 
             Op::Constant(value) => {
-                self.stack.push(value.clone());
+                let offset = self.frame().stack_offset;
+                let value = match value {
+                    Value::Function(_, block) => {
+                        let mut ups = Vec::new();
+                        println!("UPS: {:?}", block.ups);
+                        for (slot, is_up, _) in block.ups.iter() {
+                            let up = if *is_up {
+                                if let Value::Function(local_ups, _) = &self.stack[offset] {
+                                    Rc::clone(&local_ups[*slot])
+                                } else {
+                                    unreachable!()
+                                }
+                            } else {
+                                let slot = self.frame().stack_offset + slot;
+                                Rc::clone(self.find_upvalue(slot))
+                            };
+                            ups.push(up);
+                        }
+                        Value::Function(ups, block)
+                    },
+                    _ => value.clone(),
+                };
+                self.stack.push(value);
             }
 
             Op::Neg => {
@@ -521,7 +543,7 @@ impl VM {
             Op::Call(num_args) => {
                 let new_base = self.stack.len() - 1 - num_args;
                 match &self.stack[new_base] {
-                    Value::Function(ups, block) => {
+                    Value::Function(_, block) => {
                         let args = block.args();
                         if args.len() != num_args {
                             error!(self,
@@ -555,6 +577,7 @@ impl VM {
                     return Ok(OpResult::Done);
                 } else {
                     self.stack[last.stack_offset] = self.stack.pop().unwrap();
+                    self.stack.truncate(last.stack_offset + 1);
                 }
             }
         }
@@ -613,6 +636,10 @@ impl VM {
 
             Op::Jmp(_line) => {}
 
+            Op::Constant(value) => {
+                self.stack.push(value.clone());
+            }
+
             Op::ReadUpvalue(slot) => {
                 self.stack.push(self.frame().block.ups[slot].2.as_value());
             }