From a39d879c40820f87c30ddc321e804dd78e8e1246 Mon Sep 17 00:00:00 2001
From: Jeff <dev@jeffa.io>
Date: Sun, 31 Dec 2023 14:09:03 -0500
Subject: [PATCH] Implement return statements at root level

---
 src/abstract_tree/mod.rs |  6 +++++-
 tests/interpret.rs       | 14 ++++++++++++++
 2 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/src/abstract_tree/mod.rs b/src/abstract_tree/mod.rs
index 455319a..50002f8 100644
--- a/src/abstract_tree/mod.rs
+++ b/src/abstract_tree/mod.rs
@@ -60,7 +60,11 @@ impl AbstractTree for Root {
         let mut value = Value::none();
 
         for statement in &self.statements {
-            value = statement.run(source, context)?;
+            if let Statement::Return(inner_statement) = statement {
+                return inner_statement.run(source, context);
+            } else {
+                value = statement.run(source, context)?;
+            }
         }
 
         Ok(value)
diff --git a/tests/interpret.rs b/tests/interpret.rs
index 32076ea..1415f45 100644
--- a/tests/interpret.rs
+++ b/tests/interpret.rs
@@ -476,4 +476,18 @@ mod blocks {
             Ok(Value::Integer(1))
         );
     }
+
+    #[test]
+    fn root_returns_like_block() {
+        assert_eq!(
+            interpret(
+                "
+                return 1
+                1 + 1
+                3
+                "
+            ),
+            Ok(Value::Integer(1))
+        );
+    }
 }