Implement register reuse

src/compile.rs

@@ -30,10 +30,18 @@ pub struct Scope {
 }
 
 #[derive(Clone, Debug, Default)]
-pub struct LocalCounter(u16);
+pub struct LocalCounter(u16, Vec<u16>);
 
 impl LocalCounter {
     pub fn next(&mut self) -> u16 {
+        if let Some(reg) = self.1.pop() {
+            reg
+        } else {
+            self.new()
+        }
+    }
+
+    pub fn new(&mut self) -> u16 {
         let temp = self.0;
         self.0 += 1;
         temp
@@ -238,6 +246,18 @@ impl Statement {
             }
         }
 
+        for reg in 0..scope.register_counter.0 {
+            if scope.locals.values().find(|r| **r == reg).is_some() {
+                // Register is still in use
+                continue;
+            }
+            if scope.register_counter.1.contains(&reg) {
+                // Register is already in the list of unused registers
+                continue;
+            }
+            scope.register_counter.1.push(reg);
+        }
+
         instructions
     }
 }
@@ -433,10 +453,10 @@ impl Expression {
                     original_param_regs.extend(registers);
                 }
 
-                let function_reg = scope.register_counter.next();
+                let function_reg = scope.register_counter.new();
                 let mut param_regs = Vec::new();
                 for _ in &original_param_regs {
-                    param_regs.push(scope.register_counter.next());
+                    param_regs.push(scope.register_counter.new());
                 }
 
                 for (i, param_reg) in original_param_regs.iter().enumerate().rev() {

src/vm.rs

@@ -104,6 +104,8 @@ pub enum RuntimeError {
     InvalidOperands(BinaryOperator, Value, Value),
     #[error("Unable to perform {0:?} operator to {1:?}")]
     InvalidOperand(UnaryOperator, Value),
+    #[error("Tried calling a non-function: {0:?}")]
+    TriedCallingNonFunction(Value),
     #[error("{0}")]
     Custom(String),
 }
@@ -234,7 +236,9 @@ impl Debug for Value {
                 .field(&LuaNumber::from_bits(*arg0))
                 .finish(),
             Self::RustFunction(arg0) => f.debug_tuple("RustFunction").field(arg0).finish(),
-            Self::Function(_) => f.debug_tuple("Function").finish(),
+            Self::Function(closure) => f
+                .debug_tuple(&format!("Function({})", closure.prototype))
+                .finish(),
             Self::Nil => write!(f, "Nil"),
         }
     }
@@ -384,22 +388,28 @@ impl ClosureRunner {
                     );
                 }
                 Instruction::GetGlobal(reg, global) => {
-                    if let Some(global) = self
+                    let glob = self
                         .closure
                         .environment
                         .borrow()
                         .globals
                         .get(constants.get(*global as usize).unwrap())
-                    {
-                        self.stack
-                            .insert(*reg, Rc::new(RefCell::new(global.clone())));
+                        .cloned();
+                    if let Some(global) = glob {
+                        self.set_stack(*reg, global.clone());
                     } else {
                         todo!("Global not found: {:?}", constants.get(*global as usize))
                     }
                 }
                 Instruction::GetUpVal(reg, upvalreg) => {
-                    self.stack
-                        .insert(*reg, self.closure.upvalues.get(upvalreg).unwrap().clone());
+                    self.set_stack(
+                        *reg,
+                        self.closure
+                            .upvalues
+                            .get(upvalreg)
+                            .map(|v| v.borrow().clone())
+                            .unwrap(),
+                    );
                 }
                 Instruction::SetUpVal(upvalreg, reg) => {
                     *self.closure.upvalues.get(upvalreg).unwrap().borrow_mut() = self
@@ -489,13 +499,7 @@ impl ClosureRunner {
                             }
                             self.inner = Some(Box::new(closure.run(params)));
                         }
-                        _ => {
-                            if *ret_len > 0 {
-                                for i in 0..=(*ret_len - 2) {
-                                    self.set_stack(*func_reg + i, Value::Nil);
-                                }
-                            }
-                        }
+                        _ => return Err(RuntimeError::TriedCallingNonFunction(value.clone())),
                     }
                 }
                 Instruction::Close(_) => {}