From ee53c2fba4adbf10dbf2bb6d7e43802caadbb677 Mon Sep 17 00:00:00 2001
From: Emmanuel Benoit <tseeker@nocternity.net>
Date: Wed, 11 Jan 2023 15:33:21 +0100
Subject: [PATCH] Resolver - Improved handling of 'return'

---
 src/resolver.rs | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/src/resolver.rs b/src/resolver.rs
index 94453f3..de8d466 100644
--- a/src/resolver.rs
+++ b/src/resolver.rs
@@ -383,13 +383,13 @@ impl VarResolver for StmtNode {
             StmtNode::Return {
                 token,
                 value: Some(expr),
-            } => {
-                if !matches!(rs.current_type(), ScopeType::Method | ScopeType::Function) {
+            } => match rs.current_type() {
+                ScopeType::TopLevel => rs.error(token, "'return' not allowed here"),
+                ScopeType::Initializer => {
                     rs.error(token, "'return' with value is not allowed here")
-                } else {
-                    expr.resolve(rs)
                 }
-            }
+                _ => expr.resolve(rs),
+            },
 
             StmtNode::Expression(expr) => expr.resolve(rs),
             StmtNode::Print(expr) => expr.resolve(rs),