Handle multiple breaks out of switches.

Use a switch stack instead.

4 files changed, 171 insertions, 13 deletions

diff --git a/reference/shaders-no-opt/asm/comp/multi-break-switch-out-of-loop.asm.comp b/reference/shaders-no-opt/asm/comp/multi-break-switch-out-of-loop.asm.comp
new file mode 100644
index 00000000..7de95ae6
--- /dev/null
+++ b/reference/shaders-no-opt/asm/comp/multi-break-switch-out-of-loop.asm.comp
@@ -0,0 +1,52 @@
+#version 450
+layout(local_size_x = 1, local_size_y = 1, local_size_z = 1) in;
+
+layout(binding = 0, std140) uniform UBO
+{
+    int v;
+} _6;
+
+void main()
+{
+    uint count = 0u;
+    for (int i = 0; i < 4; i++)
+    {
+        bool _31_ladder_break = false;
+        do
+        {
+            bool _33_ladder_break = false;
+            do
+            {
+                bool _35_ladder_break = false;
+                do
+                {
+                    if (_6.v == 20)
+                    {
+                        _35_ladder_break = true;
+                        _33_ladder_break = true;
+                        _31_ladder_break = true;
+                        break;
+                    }
+                    break;
+                } while(false);
+                if (_35_ladder_break)
+                {
+                    break;
+                }
+                break;
+            } while(false);
+            if (_33_ladder_break)
+            {
+                break;
+            }
+            count++;
+            break;
+        } while(false);
+        if (_31_ladder_break)
+        {
+            break;
+        }
+        count++;
+    }
+}
+
diff --git a/shaders-no-opt/asm/comp/multi-break-switch-out-of-loop.asm.comp b/shaders-no-opt/asm/comp/multi-break-switch-out-of-loop.asm.comp
new file mode 100644
index 00000000..82137037
--- /dev/null
+++ b/shaders-no-opt/asm/comp/multi-break-switch-out-of-loop.asm.comp
@@ -0,0 +1,94 @@
+; SPIR-V
+; Version: 1.0
+; Generator: Khronos Glslang Reference Front End; 10
+; Bound: 53
+; Schema: 0
+               OpCapability Shader
+          %1 = OpExtInstImport "GLSL.std.450"
+               OpMemoryModel Logical GLSL450
+               OpEntryPoint GLCompute %main "main"
+               OpExecutionMode %main LocalSize 1 1 1
+               OpSource GLSL 450
+               OpName %main "main"
+               OpName %count "count"
+               OpName %i "i"
+               OpName %UBO "UBO"
+               OpMemberName %UBO 0 "v"
+               OpName %_ ""
+               OpMemberDecorate %UBO 0 Offset 0
+               OpDecorate %UBO Block
+               OpDecorate %_ DescriptorSet 0
+               OpDecorate %_ Binding 0
+               OpDecorate %gl_WorkGroupSize BuiltIn WorkgroupSize
+       %void = OpTypeVoid
+          %3 = OpTypeFunction %void
+       %uint = OpTypeInt 32 0
+%_ptr_Function_uint = OpTypePointer Function %uint
+     %uint_0 = OpConstant %uint 0
+        %int = OpTypeInt 32 1
+%_ptr_Function_int = OpTypePointer Function %int
+      %int_0 = OpConstant %int 0
+      %int_4 = OpConstant %int 4
+       %bool = OpTypeBool
+        %UBO = OpTypeStruct %int
+%_ptr_Uniform_UBO = OpTypePointer Uniform %UBO
+          %_ = OpVariable %_ptr_Uniform_UBO Uniform
+%_ptr_Uniform_int = OpTypePointer Uniform %int
+     %int_20 = OpConstant %int 20
+      %int_1 = OpConstant %int 1
+     %v3uint = OpTypeVector %uint 3
+     %uint_1 = OpConstant %uint 1
+%gl_WorkGroupSize = OpConstantComposite %v3uint %uint_1 %uint_1 %uint_1
+       %main = OpFunction %void None %3
+          %5 = OpLabel
+      %count = OpVariable %_ptr_Function_uint Function
+          %i = OpVariable %_ptr_Function_int Function
+               OpStore %count %uint_0
+               OpStore %i %int_0
+               OpBranch %14
+         %14 = OpLabel
+               OpLoopMerge %16 %17 None
+               OpBranch %18
+         %18 = OpLabel
+         %19 = OpLoad %int %i
+         %22 = OpSLessThan %bool %19 %int_4
+               OpBranchConditional %22 %15 %16
+         %15 = OpLabel
+               OpSelectionMerge %24 None
+               OpSwitch %int_0 %23
+         %23 = OpLabel
+               OpSelectionMerge %26 None
+               OpSwitch %int_0 %25
+         %25 = OpLabel
+               OpSelectionMerge %28 None
+               OpSwitch %int_0 %27
+         %27 = OpLabel
+         %33 = OpAccessChain %_ptr_Uniform_int %_ %int_0
+         %34 = OpLoad %int %33
+         %36 = OpIEqual %bool %34 %int_20
+               OpSelectionMerge %38 None
+               OpBranchConditional %36 %37 %38
+         %37 = OpLabel
+               OpBranch %16
+         %38 = OpLabel
+               OpBranch %28
+         %28 = OpLabel
+               OpBranch %26
+         %26 = OpLabel
+         %42 = OpLoad %uint %count
+         %44 = OpIAdd %uint %42 %int_1
+               OpStore %count %44
+               OpBranch %24
+         %24 = OpLabel
+         %46 = OpLoad %uint %count
+         %47 = OpIAdd %uint %46 %int_1
+               OpStore %count %47
+               OpBranch %17
+         %17 = OpLabel
+         %48 = OpLoad %int %i
+         %49 = OpIAdd %int %48 %int_1
+               OpStore %i %49
+               OpBranch %14
+         %16 = OpLabel
+               OpReturn
+               OpFunctionEnd
diff --git a/spirv_glsl.cpp b/spirv_glsl.cpp
index 9936c404..f47ac62a 100644
--- a/spirv_glsl.cpp
+++ b/spirv_glsl.cpp
@@ -327,6 +327,8 @@ void CompilerGLSL::reset(uint32_t iteration_count)
 	// Ensure that we declare phi-variable copies even if the original declaration isn't deferred
 	flushed_phi_variables.clear();
 
+	current_emitting_switch_stack.clear();
+
 	reset_name_caches();
 
 	ir.for_each_typed_id<SPIRFunction>([&](uint32_t, SPIRFunction &func) {
@@ -14895,21 +14897,32 @@ void CompilerGLSL::branch(BlockID from, BlockID to)
 		// - Break merge target all at once ...
 
 		// Very dirty workaround.
-		// Switch constructs are able to break, but they cannot break out of a loop at the same time.
+		// Switch constructs are able to break, but they cannot break out of a loop at the same time,
+		// yet SPIR-V allows it.
 		// Only sensible solution is to make a ladder variable, which we declare at the top of the switch block,
 		// write to the ladder here, and defer the break.
 		// The loop we're breaking out of must dominate the switch block, or there is no ladder breaking case.
-		if (current_emitting_switch && is_loop_break(to) &&
-		    current_emitting_switch->loop_dominator != BlockID(SPIRBlock::NoDominator) &&
-		    get<SPIRBlock>(current_emitting_switch->loop_dominator).merge_block == to)
+		if (is_loop_break(to))
 		{
-			if (!current_emitting_switch->need_ladder_break)
+			for (size_t n = current_emitting_switch_stack.size(); n; n--)
 			{
-				force_recompile();
-				current_emitting_switch->need_ladder_break = true;
-			}
+				auto *current_emitting_switch = current_emitting_switch_stack[n - 1];
+
+				if (current_emitting_switch &&
+				    current_emitting_switch->loop_dominator != BlockID(SPIRBlock::NoDominator) &&
+				    get<SPIRBlock>(current_emitting_switch->loop_dominator).merge_block == to)
+				{
+					if (!current_emitting_switch->need_ladder_break)
+					{
+						force_recompile();
+						current_emitting_switch->need_ladder_break = true;
+					}
 
-			statement("_", current_emitting_switch->self, "_ladder_break = true;");
+					statement("_", current_emitting_switch->self, "_ladder_break = true;");
+				}
+				else
+					break;
+			}
 		}
 		statement("break;");
 	}
@@ -15594,8 +15607,7 @@ void CompilerGLSL::emit_block_chain(SPIRBlock &block)
 		else if (type.basetype == SPIRType::Short)
 			label_suffix = backend.int16_t_literal_suffix;
 
-		SPIRBlock *old_emitting_switch = current_emitting_switch;
-		current_emitting_switch = &block;
+		current_emitting_switch_stack.push_back(&block);
 
 		if (block.need_ladder_break)
 			statement("bool _", block.self, "_ladder_break = false;");
@@ -15880,7 +15892,7 @@ void CompilerGLSL::emit_block_chain(SPIRBlock &block)
 			end_scope();
 		}
 
-		current_emitting_switch = old_emitting_switch;
+		current_emitting_switch_stack.pop_back();
 		break;
 	}
 
diff --git a/spirv_glsl.hpp b/spirv_glsl.hpp
index dc693787..a798737c 100644
--- a/spirv_glsl.hpp
+++ b/spirv_glsl.hpp
@@ -364,7 +364,7 @@ protected:
 	virtual void emit_function_prototype(SPIRFunction &func, const Bitset &return_flags);
 
 	SPIRBlock *current_emitting_block = nullptr;
-	SPIRBlock *current_emitting_switch = nullptr;
+	SmallVector<SPIRBlock *> current_emitting_switch_stack;
 	bool current_emitting_switch_fallthrough = false;
 
 	virtual void emit_instruction(const Instruction &instr);