Add last-expression-return-value to functions as well

We did it for if and block statements, now functions too support the
last expression being its return value. If there isn't an expression
(e.g. an assign statement) we just return `nil`.

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

src/ast.rs

@@ -1,5 +1,5 @@
 // This is an auto-generated file. Any changes made to this file may be overwritten.
-// This file was created at: 2024-10-21 15:09:43
+// This file was created at: 2024-10-21 15:50:12
 #![allow(dead_code)]
 use std::fmt::Debug;
 use std::any::Any;
@@ -127,7 +127,7 @@ pub struct FunctionExpr {
     pub lparen: Token,
     pub params: Vec<(Token, Option<ExprP>)>,
     pub return_type: Option<ExprP>,
-    pub body: Vec<StmtP>,
+    pub body: BlockExpr,
     pub rbrace: Token,
 }
 

src/compiler.rs

@@ -943,13 +943,18 @@ impl ExprVisitor for Compiler<'_> {
         }
 
         // compile body
-        for stmt in &expr.body {
+        for stmt in &expr.body.stmts {
             self.compile_stmt(stmt)?;
         }
 
-        // always end with a "return nil"
+        // compile the last expression in the block and compile that as the return value
+        if let Some(expr) = expr.body.return_expr.as_ref() {
+            self.compile_expr(expr)?;
+        } else {
+            // otherwise end with a "return nil"
             let nil = self.insert_constant(Nil::create())?;
             self.emit(end_line, Op::PushConstant(nil));
+        }
         self.emit(end_line, Op::Return);
 
         self.end_scope(end_line);

src/compiler/visitors.rs

@@ -165,11 +165,14 @@ pub(super) struct LocalAssignCollector {
 }
 
 impl LocalAssignCollector {
-    pub fn collect(body: &Vec<StmtP>) -> HashSet<String> {
+    pub fn collect(block: &BlockExpr) -> HashSet<String> {
         let mut collector = Self::default();
-        for stmt in body {
+        for stmt in &block.stmts {
             stmt.accept(&mut collector).unwrap();
         }
+        if let Some(expr) = block.return_expr.as_ref() {
+            expr.accept(&mut collector).unwrap();
+        }
         collector.names
     }
 }
@@ -319,11 +322,14 @@ pub(super) struct LocalNameCollector {
 }
 
 impl LocalNameCollector {
-    pub fn collect(body: &Vec<StmtP>) -> HashSet<String> {
+    pub fn collect(block: &BlockExpr) -> HashSet<String> {
         let mut collector = Self::default();
-        for stmt in body {
+        for stmt in &block.stmts {
             stmt.accept(&mut collector).unwrap();
         }
+        if let Some(expr) = block.return_expr.as_ref() {
+            expr.accept(&mut collector).unwrap();
+        }
         collector.names
     }
 }

src/parser.rs

@@ -995,7 +995,7 @@ impl Parser {
         }
 
         self.expect("expect '{' after function signature", TokenKind::LBrace)?;
-        let body = self.block()?;
+        let body = self.block_expr()?;
         let rbrace = self.prev.clone().unwrap();
 
         Ok(Box::new(FunctionExpr {

tools/genast.py

@@ -243,7 +243,7 @@ GENERATE = [
             "lparen: Token",
             "params: Vec<(Token, Option<ExprP>)>",
             "return_type: Option<ExprP>",
-            "body: Vec<StmtP>",
+            "body: BlockExpr",
             "rbrace: Token",
         ],
     )