From 00129f3bc09cbf80061ecb2448374296438b65c8 Mon Sep 17 00:00:00 2001
From: Chuanqi Xu <yedeng.yd@linux.alibaba.com>
Date: Tue, 24 Sep 2024 14:00:53 +0800
Subject: [PATCH] [CIR] [CodeGen] Fix the case there are multiple stmt followed
 in case

Close https://github.com/llvm/clangir/issues/876
---
 clang/lib/CIR/CodeGen/CIRGenStmt.cpp |  4 ++--
 clang/test/CIR/CodeGen/goto.cpp      | 22 ++++++++++++++++++++++
 2 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/clang/lib/CIR/CodeGen/CIRGenStmt.cpp b/clang/lib/CIR/CodeGen/CIRGenStmt.cpp
index 543b744cc5d0..426da35b5238 100644
--- a/clang/lib/CIR/CodeGen/CIRGenStmt.cpp
+++ b/clang/lib/CIR/CodeGen/CIRGenStmt.cpp
@@ -996,18 +996,18 @@ mlir::LogicalResult CIRGenFunction::buildSwitchBody(
     for (auto *c : compoundStmt->body()) {
       if (auto *switchCase = dyn_cast<SwitchCase>(c)) {
         res = buildSwitchCase(*switchCase, condType, caseAttrs);
+        lastCaseBlock = builder.getBlock();
       } else if (lastCaseBlock) {
         // This means it's a random stmt following up a case, just
         // emit it as part of previous known case.
         mlir::OpBuilder::InsertionGuard guardCase(builder);
         builder.setInsertionPointToEnd(lastCaseBlock);
         res = buildStmt(c, /*useCurrentScope=*/!isa<CompoundStmt>(c));
+        lastCaseBlock = builder.getBlock();
       } else {
         llvm_unreachable("statement doesn't belong to any case region, NYI");
       }
 
-      lastCaseBlock = builder.getBlock();
-
       if (res.failed())
         break;
     }
diff --git a/clang/test/CIR/CodeGen/goto.cpp b/clang/test/CIR/CodeGen/goto.cpp
index 5a8d598d95cd..81eb4ec43e65 100644
--- a/clang/test/CIR/CodeGen/goto.cpp
+++ b/clang/test/CIR/CodeGen/goto.cpp
@@ -288,3 +288,25 @@ void foo() {
 // NOFLAT:   cir.scope {
 // NOFLAT:     cir.label "label"
 // NOFLAT:     %0 = cir.alloca !ty_S, !cir.ptr<!ty_S>, ["agg.tmp0"]
+
+extern "C" void action1();
+extern "C" void action2();
+extern "C" void multiple_non_case(int v) {
+  switch (v) {
+    default:
+        action1();
+      l2:
+        action2();
+        break;
+  }
+}
+
+// NOFLAT: cir.func @multiple_non_case
+// NOFLAT: cir.switch
+// NOFLAT: case (default)
+// NOFLAT: cir.call @action1()
+// NOFLAT: cir.br ^[[BB1:[a-zA-Z0-9]+]]
+// NOFLAT: ^[[BB1]]:
+// NOFLAT: cir.label
+// NOFLAT: cir.call @action2()
+// NOFLAT: cir.break