From 58154b9e6ee92357d1eb6ea97366282ccf7ac866 Mon Sep 17 00:00:00 2001
From: lisyarus <lisyarus@gmail.com>
Date: Mon, 22 Dec 2025 18:33:37 +0300
Subject: [PATCH] Handle return statement outside function scope in AST
 finilizator

---
 libs/parser/source/finilize.cpp | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

diff --git a/libs/parser/source/finilize.cpp b/libs/parser/source/finilize.cpp
index 4eab83b..6ae85dd 100644
--- a/libs/parser/source/finilize.cpp
+++ b/libs/parser/source/finilize.cpp
@@ -16,6 +16,7 @@ namespace pslang::parser
 		std::vector<stack_entry> stack;
 		stack.push_back(result.get());
 		std::size_t current_indent = 0;
+		std::size_t in_function_scope = 0;
 
 		auto current_statement_list = [&](ast::location const & location) -> ast::statement_list *
 		{
@@ -50,11 +51,14 @@ namespace pslang::parser
 			{
 				stack.pop_back();
 				--current_indent;
+				if (in_function_scope > 0)
+					--in_function_scope;
 			}
 
 			 // Now statement.indentation == current_indent
 
 			ast::statement_list * list = nullptr;
+			bool is_function_definition = false;
 
 			if (auto if_block = std::get_if<ast::if_block>(statement.statement.get()))
 			{
@@ -96,6 +100,7 @@ namespace pslang::parser
 				function_definition->statements = std::make_unique<ast::statement_list>();
 				list = function_definition->statements.get();
 				current_statement_list(location)->statements.push_back(std::move(statement.statement));
+				is_function_definition = true;
 			}
 			else if (auto field_definition = std::get_if<ast::field_definition>(statement.statement.get()))
 			{
@@ -110,6 +115,14 @@ namespace pslang::parser
 				current_statement_list(location)->statements.push_back(std::move(statement.statement));
 				stack.push_back(std::get_if<ast::struct_definition>(current_statement_list(location)->statements.back().get()));
 				++current_indent;
+				if (in_function_scope > 0)
+					++in_function_scope;
+			}
+			else if (auto return_statement = std::get_if<ast::return_statement>(statement.statement.get()))
+			{
+				if (in_function_scope == 0)
+					throw parse_error("Return statement outside of function scope", return_statement->location);
+				current_statement_list(location)->statements.push_back(std::move(statement.statement));
 			}
 			else
 			{
@@ -120,6 +133,8 @@ namespace pslang::parser
 			{
 				stack.push_back(list);
 				++current_indent;
+				if (in_function_scope > 0 || is_function_definition)
+					++in_function_scope;
 			}
 		}