From b937207dccd05f700855e9aeb3f7c375b595b827 Mon Sep 17 00:00:00 2001
From: Yong He <yonghe@outlook.com>
Date: Fri, 12 Apr 2024 14:40:30 -0700
Subject: Fix IR lowering bug of do-while loops. (#3941)

---
 source/slang/slang-ir-validate.cpp |  7 ++++++-
 source/slang/slang-lower-to-ir.cpp | 14 ++++++++++----
 tests/bugs/gh-3930.slang           | 26 ++++++++++++++++++++++++++
 3 files changed, 42 insertions(+), 5 deletions(-)
 create mode 100644 tests/bugs/gh-3930.slang

diff --git a/source/slang/slang-ir-validate.cpp b/source/slang/slang-ir-validate.cpp
index 2868c9929..a4a2921fe 100644
--- a/source/slang/slang-ir-validate.cpp
+++ b/source/slang/slang-ir-validate.cpp
@@ -66,6 +66,7 @@ namespace Slang
         State state = kState_Initial;
 
         IRInst* prevChild = nullptr;
+        bool hasSeenTerminatorInst = false;
         for(auto child : parent->getDecorationsAndChildren() )
         {
             // We need to check the integrity of the parent/next/prev links of
@@ -105,7 +106,11 @@ namespace Slang
                 validate(context, !as<IRTerminatorInst>(child), child, "terminator must be last instruction in a block");
             }
 
-
+            if (as<IRTerminatorInst>(child))
+            {
+                validate(context, !hasSeenTerminatorInst, child, "block must not contain more than one terminator");
+                hasSeenTerminatorInst = true;
+            }
             prevChild = child;
         }
     }
diff --git a/source/slang/slang-lower-to-ir.cpp b/source/slang/slang-lower-to-ir.cpp
index 3f776c1a1..855bde6a6 100644
--- a/source/slang/slang-lower-to-ir.cpp
+++ b/source/slang/slang-lower-to-ir.cpp
@@ -5890,6 +5890,12 @@ struct StmtLoweringVisitor : StmtVisitor<StmtLoweringVisitor>
         {
             auto irCondition = getSimpleVal(context,
                 lowerRValueExpr(context, condExpr));
+
+            // One thing to be careful here is that lowering irCondition
+            // may create additional blocks due to short circuiting, so
+            // the block we are current inserting into is not necessarily
+            // the same as `testLabel`.
+            // 
             auto invCondition = builder->emitNot(irCondition->getDataType(), irCondition);
 
             // Now we want to `break` if the loop condition is false,
@@ -5907,15 +5913,15 @@ struct StmtLoweringVisitor : StmtVisitor<StmtLoweringVisitor>
             // 
             // mergeBlock:
             //   goto breakLabel;
-            auto mergeBlock = builder->emitBlock();
-            builder->emitBranch(loopHead);
-
-            builder->setInsertInto(testLabel);
+            auto mergeBlock = builder->createBlock();
             builder->emitIfElse(
                 invCondition,
                 breakLabel,
                 mergeBlock,
                 mergeBlock);
+
+            insertBlock(mergeBlock);
+            builder->emitBranch(loopHead);
         }
 
         // Finally we insert the label that a `break` will jump to
diff --git a/tests/bugs/gh-3930.slang b/tests/bugs/gh-3930.slang
new file mode 100644
index 000000000..6ba860ef3
--- /dev/null
+++ b/tests/bugs/gh-3930.slang
@@ -0,0 +1,26 @@
+// A regression: the do-while loop ir lowering logic is making incorrect assumptions and is
+// broken after short circuiting expr change, leading to invalid IR being generated from lowering
+// (a block has two terminator insts).
+
+//TEST:SIMPLE(filecheck=CHECK): -target spirv -emit-spirv-directly
+
+// CHECK: OpEntryPoint
+
+struct Constants {
+  uint *buffer;
+}
+[[vk::push_constant]] Constants push_constants;
+
+[shader("compute")] [numthreads(1, 1, 1)] void compute(void) {
+  uint index = push_constants.buffer[0];
+
+  [[unroll]] for (uint i = 0; i < 2; i++) {
+    GroupMemoryBarrierWithWaveSync();
+
+    if (index > 0 && i < push_constants.buffer[index - 1]) {
+      do {
+        index--;
+      } while (index > 0 && i < push_constants.buffer[index - 1]);
+    }
+  }
+}
-- 
cgit v1.2.3