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)

src/syn/ast.rs

@@ -24,6 +24,7 @@ impl<V: Visit<Out=()>> DefaultAccept<V> for Body {
 pub enum Stmt {
     Expr(Expr),
     Assign(AssignStmt),
+    Return(ReturnStmt),
 }
 
 //
@@ -43,6 +44,7 @@ impl<V: Visit> DefaultAccept<V> for Stmt {
         match self {
             Stmt::Expr(e) => e.accept(visitor),
             Stmt::Assign(a) => a.accept(visitor),
+            Stmt::Return(r) => r.accept(visitor),
         }
     }
 }
@@ -99,6 +101,29 @@ impl<V: Visit<Out=()>> DefaultAccept<V> for LhsExpr {
     }
 }
 
+//
+// struct ReturnStmt
+//
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct ReturnStmt {
+    pub expr: Option<Expr>,
+}
+
+impl<V: Visit> Accept<V> for ReturnStmt {
+    fn accept(&self, visitor: &mut V) -> V::Out {
+        visitor.visit_return_stmt(self)
+    }
+}
+
+impl<V: Visit<Out=()>> DefaultAccept<V> for ReturnStmt {
+    fn default_accept(&self, visitor: &mut V) -> V::Out {
+        if let Some(expr) = self.expr.as_ref() {
+            expr.accept(visitor)
+        }
+    }
+}
+
+
 //
 // struct Expr
 //

src/syn/lexer.l

@@ -1,5 +1,6 @@
 %%
 fn                          "fn"
+return                      "return"
 
 [\r\n;]+                    "EOL"
 [a-zA-Z_][a-zA-Z0-9_]*      "IDENT"

src/syn/parser.y

@@ -22,6 +22,7 @@ Body -> Result<Vec<Stmt>>:
 Stmt -> Result<Stmt>:
         Expr { Ok(Stmt::Expr($1?)) }
     |   Assign { Ok(Stmt::Assign($1?)) }
+    |   Return { Ok(Stmt::Return($1?)) }
     ;
 
 Assign -> Result<AssignStmt>:
@@ -41,6 +42,11 @@ LhsExpr -> Result<LhsExpr>:
         }
     ;
 
+Return -> Result<ReturnStmt>:
+        'return' Expr { Ok(ReturnStmt { expr: Some($2?), }) }
+    |   'return'      { Ok(ReturnStmt { expr: None, }) }
+    ;
+
 Expr -> Result<Expr>: BinExpr { $1 };
 
 BinExpr -> Result<Expr>:

src/syn/visit.rs

@@ -14,6 +14,7 @@ pub trait Visit {
     fn visit_body(&mut self, body: &Body) -> Self::Out;
     fn visit_stmt(&mut self, stmt: &Stmt) -> Self::Out;
     fn visit_assign_stmt(&mut self, assign: &AssignStmt) -> Self::Out;
+    fn visit_return_stmt(&mut self, ret: &ReturnStmt) -> Self::Out;
     fn visit_lhs_expr(&mut self, lhs_expr: &LhsExpr) -> Self::Out;
     fn visit_expr(&mut self, expr: &Expr) -> Self::Out;
     fn visit_bin_expr(&mut self, expr: &BinExpr) -> Self::Out;
@@ -31,6 +32,7 @@ copy/paste of default_accepts
     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 { DefaultAccept::default_accept(lhs_expr, self); }
     fn visit_expr(&mut self, expr: &Expr) -> Self::Out { DefaultAccept::default_accept(expr, self); }
     fn visit_bin_expr(&mut self, expr: &BinExpr) -> Self::Out { DefaultAccept::default_accept(expr, self); }