From 8cba98d2e207d6572b398cff084f969447f82c4d Mon Sep 17 00:00:00 2001
From: axp <645124937@qq.com>
Date: Mon, 1 Apr 2024 21:48:55 +0800
Subject: [PATCH] switchStmt support single caseStmt and defaultStmt

---
 clang/lib/CIR/CodeGen/CIRGenStmt.cpp | 18 +++++++++++++--
 clang/test/CIR/CodeGen/switch.cpp    | 33 ++++++++++++++++++++++++++++
 2 files changed, 49 insertions(+), 2 deletions(-)

diff --git a/clang/lib/CIR/CodeGen/CIRGenStmt.cpp b/clang/lib/CIR/CodeGen/CIRGenStmt.cpp
index 0dd5bd345ab6..57b331a6d028 100644
--- a/clang/lib/CIR/CodeGen/CIRGenStmt.cpp
+++ b/clang/lib/CIR/CodeGen/CIRGenStmt.cpp
@@ -980,11 +980,25 @@ mlir::LogicalResult CIRGenFunction::buildSwitchStmt(const SwitchStmt &S) {
         getLoc(S.getBeginLoc()), condV,
         /*switchBuilder=*/
         [&](mlir::OpBuilder &b, mlir::Location loc, mlir::OperationState &os) {
+          SmallVector<mlir::Attribute, 4> caseAttrs;
+          currLexScope->setAsSwitch();
+
+          if (auto *caseStmt = dyn_cast<CaseStmt>(S.getBody())) {
+            res = buildCaseStmt(*caseStmt, condV.getType(), caseAttrs, os);
+            os.addAttribute("cases", builder.getArrayAttr(caseAttrs));
+            return;
+          }
+
+          if (auto *defaultStmt = dyn_cast<DefaultStmt>(S.getBody())) {
+            res =
+                buildDefaultStmt(*defaultStmt, condV.getType(), caseAttrs, os);
+            os.addAttribute("cases", builder.getArrayAttr(caseAttrs));
+            return;
+          }
+
           auto *cs = dyn_cast<CompoundStmt>(S.getBody());
           assert(cs && "expected compound stmt");
-          SmallVector<mlir::Attribute, 4> caseAttrs;
 
-          currLexScope->setAsSwitch();
           mlir::Block *lastCaseBlock = nullptr;
           for (auto *c : cs->body()) {
             bool caseLike = isa<CaseStmt, DefaultStmt>(c);
diff --git a/clang/test/CIR/CodeGen/switch.cpp b/clang/test/CIR/CodeGen/switch.cpp
index 2e3ec86c2ca8..f42e0633e0a1 100644
--- a/clang/test/CIR/CodeGen/switch.cpp
+++ b/clang/test/CIR/CodeGen/switch.cpp
@@ -266,6 +266,7 @@ void sw12(int a) {
     break;
   }
 }
+
 //      CHECK: cir.func @_Z4sw12i
 //      CHECK:   cir.scope {
 //      CHECK:     cir.switch
@@ -275,6 +276,38 @@ void sw12(int a) {
 // CHECK-NEXT:       cir.break
 // CHECK-NEXT:     }
 
+int sw13(int a) {
+  switch (a)
+  case 1:
+    return -1;
+  return a;
+}
+
+//      CHECK: cir.func @_Z4sw13i
+//      CHECK:   cir.scope {
+//      CHECK:     cir.switch
+// CHECK-NEXT:     case (equal, 1) {
+//      CHECK:       cir.return
+//      CHECK:     }
+//      CHECK:   }
+//      CHECK: cir.return
+
+int sw14(int a) {
+  switch (a)
+  default:
+    return -1;
+  return a;
+}
+
+//      CHECK: cir.func @_Z4sw14i
+//      CHECK:   cir.scope {
+//      CHECK:     cir.switch
+// CHECK-NEXT:     case (default) {
+//      CHECK:       cir.return
+//      CHECK:     }
+//      CHECK:   }
+//      CHECK: cir.return
+
 void fallthrough(int x) {
   switch (x) {
     case 1: