From cd4d012b016958e451df948108d896e1335717fe Mon Sep 17 00:00:00 2001
From: Vurv <56230599+Vurv78@users.noreply.github.com>
Date: Thu, 5 Oct 2023 17:13:10 -0700
Subject: [PATCH] Fix return inside switch case

Fixes #2783
---
 data/expression2/tests/regressions/2783.txt   | 20 +++++++++++++++++++
 .../gmod_wire_expression2/base/compiler.lua   | 13 +++++++-----
 2 files changed, 28 insertions(+), 5 deletions(-)
 create mode 100644 data/expression2/tests/regressions/2783.txt

diff --git a/data/expression2/tests/regressions/2783.txt b/data/expression2/tests/regressions/2783.txt
new file mode 100644
index 0000000000..ca87673619
--- /dev/null
+++ b/data/expression2/tests/regressions/2783.txt
@@ -0,0 +1,20 @@
+## SHOULD_PASS:EXECUTE
+
+function number test(N:number) {
+    switch (N) {
+        case 1,
+            return 5
+        case 3,
+            error("unreachable")
+        case 2,
+            return 2
+        case 5,
+            error("unreachable")
+    }
+
+    return 1
+}
+
+assert( test(1) == 5 )
+assert( test(2) == 2 )
+assert( test(4) == 1 )
\ No newline at end of file
diff --git a/lua/entities/gmod_wire_expression2/base/compiler.lua b/lua/entities/gmod_wire_expression2/base/compiler.lua
index af9c0be094..dfa48646a6 100644
--- a/lua/entities/gmod_wire_expression2/base/compiler.lua
+++ b/lua/entities/gmod_wire_expression2/base/compiler.lua
@@ -536,8 +536,8 @@ local CompileVisitors = {
 
 		return function(state) ---@param state RuntimeContext
 			local expr = expr(state)
-
 			state:PushScope()
+
 			for i = 1, ncases do
 				local case = cases[i]
 				if case[1](state, expr) ~= 0 then
@@ -545,15 +545,17 @@ local CompileVisitors = {
 
 					if state.__break__ then
 						state.__break__ = false
-						state:PopScope()
-						return
+						goto exit
+					elseif state.__return__ then -- Yes this should only be checked if the switch is inside a function, but I don't care enough about the performance of switch case to add another duplicated 30 lines to the file
+						goto exit
 					else -- Fallthrough, run every case until break found.
 						for j = i + 1, ncases do
 							cases[j][2](state)
 							if state.__break__ then
 								state.__break__ = false
-								state:PopScope()
-								return
+								goto exit
+							elseif state.__return__ then
+								goto exit
 							end
 						end
 					end
@@ -564,6 +566,7 @@ local CompileVisitors = {
 				default(state)
 			end
 
+			::exit::
 			state:PopScope()
 		end
 	end,