Make returns work in functions

sample.lox

@@ -1,9 +1,11 @@
 {
-fun c() {
-  c("too", "many");
-}
-fun b() { c(); }
-fun a() { b(); }
+  fun fib(n) {
+    if (n < 2) {
+      print n;
+    } else {
+      fib(n - 1);
+    }
+  }
 
-a();
+  fib(40);
 }

src/compile.rs

@@ -17,34 +17,20 @@ pub struct Compiler<'a> {
     scanner: Scanner<'a>,
     parser: Parser<'a>,
     strings: &'a mut StringPool,
-    locals: Vec<Local<'a>>,
 }
 
 impl<'a> Compiler<'a> {
     pub fn new(source: &'a str, strings: &'a mut StringPool) -> Self {
-        let mut locals = Vec::new();
-        locals.push(Local {
-            name: Token::new(),
-            depth: 0,
-        });
-
         Self {
             scanner: Scanner::new(source.as_bytes()),
             parser: Parser::new(),
             strings,
-            locals,
         }
     }
 
     pub fn compile(mut self, kind: FunctionKind) -> crate::Result<Function> {
-        let mut session = CompileSession::new(
-            kind,
-            &mut self.scanner,
-            &mut self.parser,
-            &mut self.locals,
-            self.strings,
-            0,
-        );
+        let mut session =
+            CompileSession::new(kind, &mut self.scanner, &mut self.parser, self.strings);
         session.advance();
         session.declaration();
         session.end_compile();
@@ -63,7 +49,7 @@ struct CompileSession<'a, 'b> {
     function: Function,
     kind: FunctionKind,
     strings: &'b mut StringPool,
-    locals: &'b mut Vec<Local<'a>>,
+    locals: Vec<Local<'a>>,
     scope_depth: isize,
 }
 
@@ -84,35 +70,31 @@ impl<'a, 'b> CompileSession<'a, 'b> {
         kind: FunctionKind,
         scanner: &'b mut Scanner<'a>,
         parser: &'b mut Parser<'a>,
-        locals: &'b mut Vec<Local<'a>>,
         strings: &'b mut StringPool,
-        scope_depth: isize,
     ) -> Self {
         let mut function = Function::new();
         if kind != FunctionKind::Script {
             function.name = parser.previous.lexeme_str().to_string();
         }
 
+        let locals = vec![Local {
+            name: Token::new(),
+            depth: 0,
+        }];
+
         Self {
             scanner,
             parser,
             function,
             kind,
             strings,
             locals,
-            scope_depth,
+            scope_depth: 0,
         }
     }
 
     fn new_inner(&mut self, kind: FunctionKind) -> CompileSession<'a, '_> {
-        CompileSession::new(
-            kind,
-            self.scanner,
-            self.parser,
-            self.locals,
-            self.strings,
-            self.scope_depth,
-        )
+        CompileSession::new(kind, self.scanner, self.parser, self.strings)
     }
 
     fn declaration(&mut self) {

src/vm.rs

@@ -19,6 +19,7 @@ pub struct Vm {
 struct CallFrame {
     function: Rc<Function>,
     ip: usize,
+    stack_top: usize,
 }
 
 impl Vm {
@@ -141,7 +142,13 @@ impl Vm {
                 }
                 GetLocal => {
                     let slot = self.read_byte().unwrap();
-                    push!(self, self.stack[slot as usize].clone());
+                    let top = frame!(self).stack_top;
+                    push!(self, self.stack[top..][slot as usize].clone());
+                }
+                SetLocal => {
+                    let slot = self.read_byte().unwrap();
+                    let top = frame!(self).stack_top;
+                    self.stack[top..][slot as usize] = peek!(self).clone();
                 }
                 GetGlobal => {
                     let name = self.read_constant();
@@ -152,10 +159,6 @@ impl Vm {
                         runtime_error!(self, "Undefined variable {}", self.strings.lookup(name));
                     }
                 }
-                SetLocal => {
-                    let slot = self.read_byte().unwrap();
-                    self.stack[slot as usize] = peek!(self).clone();
-                }
                 SetGlobal => {
                     let name = self.read_constant();
                     let name = name.as_string().unwrap();
@@ -234,7 +237,21 @@ impl Vm {
                     let arg_count = self.read_byte().unwrap();
                     self.call_value(peek!(self, arg_count as usize).clone(), arg_count)?;
                 }
-                Return => return Ok(()),
+                Return => {
+                    let result = pop!(self);
+
+                    let f = self.frames.pop().unwrap();
+
+                    // Unwind the stack
+                    self.stack.truncate(f.stack_top);
+
+                    if self.frames.is_empty() {
+                        self.stack.pop();
+                        return Ok(());
+                    }
+
+                    self.stack.push(result);
+                }
             }
         }
 
@@ -264,7 +281,11 @@ impl Vm {
             runtime_error!(self, "Stack overflow",);
         }
 
-        let frame = CallFrame { function, ip: 0 };
+        let frame = CallFrame {
+            function,
+            ip: 0,
+            stack_top: self.stack.len() - arg_count as usize - 1,
+        };
         self.frames.push(frame);
         Ok(())
     }