Add "return" statement

Functions may now be exited by returning a value or no value.

Signed-off-by: Alek Ratzloff <alekratz@gmail.com>

src/compile/locals.rs

@@ -33,6 +33,7 @@ impl Visit for CollectLocals<'_> {
     fn visit_body(&mut self, body: &Body) -> Self::Out { DefaultAccept::default_accept(body, self); }
     fn visit_stmt(&mut self, stmt: &Stmt) -> Self::Out { DefaultAccept::default_accept(stmt, self); }
     fn visit_assign_stmt(&mut self, assign: &AssignStmt) -> Self::Out { DefaultAccept::default_accept(assign, self); }
+    fn visit_return_stmt(&mut self, ret: &ReturnStmt) -> Self::Out { DefaultAccept::default_accept(ret, self); }
     fn visit_lhs_expr(&mut self, lhs_expr: &LhsExpr) -> Self::Out {
         match lhs_expr {
             LhsExpr::Name(name) => {

src/compile/thunk.rs

@@ -251,6 +251,17 @@ impl Visit for CompileBody<'_> {
         Ok(thunk)
     }
 
+    fn visit_return_stmt(&mut self, stmt: &ReturnStmt) -> Self::Out {
+        let mut thunk = if let Some(expr) = stmt.expr.as_ref() {
+            self.visit_expr(expr)?
+        } else {
+            Inst::PushSym(crate::obj::reserved::NIL_NAME.sym).into()
+        };
+        thunk.push(Inst::Return);
+
+        Ok(thunk)
+    }
+
     fn visit_lhs_expr(&mut self, lhs_expr: &LhsExpr) -> Self::Out {
         // Do different things depending on the LHS
         let mut thunk;
@@ -365,12 +376,26 @@ impl Visit for CompileBody<'_> {
             let sym = global_sym(param.to_string());
             self.compile.create_local(sym);
         }
+
+        // Compile function body
+        let mut code = self.visit_body(&expr.body)?
+            .flatten()
+            .to_vec();
+
+        // If the last instruction is not a return, or if there are no instructions, then return
+        // :nil value.
+        if !matches!(code.last(), Some(Inst::Return)) {
+            code.push(Inst::PushSym(crate::obj::reserved::NIL_NAME.sym));
+            code.push(Inst::Return);
+        }
+
         // remap (Sym -> Name) to be (Name -> Sym) and make sure it's all in order.
         let scope_locals: BTreeMap<_, _> = self.compile.pop_scope_layer()
             .unwrap()
             .into_iter()
             .map(|(sym, name)| (name, sym))
             .collect();
+
         // this should be in numeric order since:
         // 1. locals are created exactly once or looked up
         // 2. scope_locals is a btreemap, keyed by names, which are in order from 0..N
@@ -382,13 +407,6 @@ impl Visit for CompileBody<'_> {
             })
             .collect();
 
-        let mut code = self.visit_body(&expr.body)?
-            .flatten()
-            .to_vec();
-        // XXX TODO(compile)
-        // remove this when we get returns implemented
-        code.push(Inst::PushSym(crate::obj::reserved::NIL_NAME.sym));
-        code.push(Inst::Return);
         let (hdl, _fun) = self.compile.push_const(UserFun::new_obj(code, locals));
 
         // TODO(compile) : determine return value at the end of the body (preferably at parse-time)