Fix crash when processing nested switch. (#2624)

* Fix crash when processing nested switch.

* Clean up.

---------

Co-authored-by: Yong He <yhe@nvidia.com>

3 files changed, 48 insertions, 4 deletions

diff --git a/source/slang/slang-ir-eliminate-multilevel-break.cpp b/source/slang/slang-ir-eliminate-multilevel-break.cpp
index 2330991c5..35412cc07 100644
--- a/source/slang/slang-ir-eliminate-multilevel-break.cpp
+++ b/source/slang/slang-ir-eliminate-multilevel-break.cpp
@@ -61,15 +61,19 @@ struct EliminateMultiLevelBreakContext
 
         void collectBreakableRegionBlocks(BreakableRegionInfo& info)
         {
+            // Push break block to a stack so we can easily check if a block is a break block in its
+            // parent regions.
+            breakBlocks.Add(info.getBreakBlock());
+            
             auto successors = as<IRBlock>(info.headerInst->getParent())->getSuccessors();
             for (auto successor : successors)
             {
+                if (!breakBlocks.Add(successor))
+                    continue;
                 if (info.blockSet.Add(successor))
                     info.blocks.add(successor);
             }
-            // Push break block to a stack so we can easily check if a block is a break block in its
-            // parent regions.
-            breakBlocks.Add(info.getBreakBlock());
+
             for (Index i = 0; i < info.blocks.getCount(); i++)
             {
                 auto block = info.blocks[i];
@@ -133,7 +137,6 @@ struct EliminateMultiLevelBreakContext
                     break;
                 }
             }
-
             for (auto& l : regions)
             {
                 l->forEach(
diff --git a/tests/bugs/nested-switch.slang b/tests/bugs/nested-switch.slang
new file mode 100644
index 000000000..3cfc0758b
--- /dev/null
+++ b/tests/bugs/nested-switch.slang
@@ -0,0 +1,37 @@
+// nested-switch.slang
+
+//TEST(compute):COMPARE_COMPUTE: -shaderobj
+//TEST(compute):COMPARE_COMPUTE:-vk -shaderobj
+//TEST(compute):COMPARE_COMPUTE:-cpu -shaderobj
+
+int test(int t, int r)
+{
+    int result = 0;
+    switch (t)
+    {
+    case 0:
+        switch (r)
+        {
+        case 0:
+            return 1;
+        case 2:
+            return 2;
+        }
+    default:
+        result = 3;
+        break;
+    }
+    return result;
+}
+
+//TEST_INPUT:ubuffer(data=[0 0 0 0], stride=4):out,name=outputBuffer
+RWStructuredBuffer<int> outputBuffer;
+
+[numthreads(1, 1, 1)]
+void computeMain(int3 dispatchThreadID: SV_DispatchThreadID)
+{
+    outputBuffer[0] = test(dispatchThreadID.x, 0);
+    outputBuffer[1] = test(dispatchThreadID.x, 1);
+    outputBuffer[2] = test(dispatchThreadID.x, 2);
+    outputBuffer[3] = 0;
+}
diff --git a/tests/bugs/nested-switch.slang.expected.txt b/tests/bugs/nested-switch.slang.expected.txt
new file mode 100644
index 000000000..e5716bdfa
--- /dev/null
+++ b/tests/bugs/nested-switch.slang.expected.txt
@@ -0,0 +1,4 @@
+1
+3
+2
+0