From 9d7a63f25bd755d66689be50df7b1ff7f96e845c Mon Sep 17 00:00:00 2001
From: sofia <sofia@teascade.net>
Date: Wed, 2 Aug 2023 15:34:20 +0300
Subject: [PATCH] Make return special for block

---
 src/parser.rs | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/src/parser.rs b/src/parser.rs
index 10d0c7d..233f3a8 100644
--- a/src/parser.rs
+++ b/src/parser.rs
@@ -208,22 +208,23 @@ impl Parse for FunctionSignature {
 }
 
 #[derive(Debug)]
-pub struct Block(Vec<BlockLevelStatement>);
+pub struct Block(Vec<BlockLevelStatement>, Option<Expression>);
 
 impl Parse for Block {
     fn parse(mut stream: TokenStream) -> Result<Self, ()> {
         let mut statements = Vec::new();
+        let mut return_stmt = None;
         stream.expect(Token::BraceOpen)?;
         while !matches!(stream.peek(), Some(Token::BraceClose)) {
             let statement = stream.parse()?;
-            if let BlockLevelStatement::Return(_) = &statement {
-                statements.push(statement);
+            if let BlockLevelStatement::Return(e) = &statement {
+                return_stmt = Some(e.clone());
                 break; // Return has to be the last statement
             }
             statements.push(statement);
         }
         stream.expect(Token::BraceClose)?;
-        Ok(Block(statements))
+        Ok(Block(statements, return_stmt))
     }
 }