From 224533ea8af3f7f5edf3eecfac55ff70a268bf94 Mon Sep 17 00:00:00 2001
From: Alek Ratzloff <alekratz@gmail.com>
Date: Mon, 21 Oct 2024 15:59:36 -0700
Subject: [PATCH] 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               |  4 ++--
 src/compiler.rs          | 13 +++++++++----
 src/compiler/visitors.rs | 14 ++++++++++----
 src/parser.rs            |  2 +-
 tools/genast.py          |  2 +-
 5 files changed, 23 insertions(+), 12 deletions(-)

diff --git a/src/ast.rs b/src/ast.rs
index 47ca9fb..9086abf 100644
--- a/src/ast.rs
+++ b/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,
 }
 
diff --git a/src/compiler.rs b/src/compiler.rs
index 5f6e9d4..7893ea9 100644
--- a/src/compiler.rs
+++ b/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"
-        let nil = self.insert_constant(Nil::create())?;
-        self.emit(end_line, Op::PushConstant(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);
diff --git a/src/compiler/visitors.rs b/src/compiler/visitors.rs
index 987ff57..413101c 100644
--- a/src/compiler/visitors.rs
+++ b/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
     }
 }
diff --git a/src/parser.rs b/src/parser.rs
index 534eee7..599118c 100644
--- a/src/parser.rs
+++ b/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 {
diff --git a/tools/genast.py b/tools/genast.py
index 7eff05e..a6d1dd9 100755
--- a/tools/genast.py
+++ b/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",
         ],
     )