[CIR] [CodeGen] handle the case of 'case' after label statement after 'case'

Author: ChuanqiXu9

clang/lib/CIR/CodeGen/CIRGenFunction.h

@@ -478,6 +478,13 @@ class CIRGenFunction : public CIRGenTypeCache {
   // applies to. nullptr if there is no 'musttail' on the current statement.
   const clang::CallExpr *MustTailCall = nullptr;
 
+  /// The attributes of cases collected during emitting the body of a switch
+  /// stmt.
+  llvm::SmallVector<llvm::SmallVector<mlir::Attribute, 4>, 2> caseAttrsStack;
+
+  /// The type of the condition for the emitting switch statement.
+  llvm::SmallVector<mlir::Type, 2> condTypeStack;
+
   clang::ASTContext &getContext() const;
 
   CIRGenBuilderTy &getBuilder() { return builder; }
@@ -1210,13 +1217,9 @@ class CIRGenFunction : public CIRGenTypeCache {
   buildDefaultStmt(const clang::DefaultStmt &S, mlir::Type condType,
                    SmallVector<mlir::Attribute, 4> &caseAttrs);
 
-  mlir::LogicalResult
-  buildSwitchCase(const clang::SwitchCase &S, mlir::Type condType,
-                  SmallVector<mlir::Attribute, 4> &caseAttrs);
+  mlir::LogicalResult buildSwitchCase(const clang::SwitchCase &S);
 
-  mlir::LogicalResult
-  buildSwitchBody(const clang::Stmt *S, mlir::Type condType,
-                  SmallVector<mlir::Attribute, 4> &caseAttrs);
+  mlir::LogicalResult buildSwitchBody(const clang::Stmt *S);
 
   mlir::cir::FuncOp generateCode(clang::GlobalDecl GD, mlir::cir::FuncOp Fn,
                                  const CIRGenFunctionInfo &FnInfo);

clang/lib/CIR/CodeGen/CIRGenStmt.cpp

@@ -303,8 +303,7 @@ mlir::LogicalResult CIRGenFunction::buildSimpleStmt(const Stmt *S,
 
   case Stmt::CaseStmtClass:
   case Stmt::DefaultStmtClass:
-    assert(0 &&
-           "Should not get here, currently handled directly from SwitchStmt");
+    return buildSwitchCase(cast<SwitchCase>(*S));
     break;
 
   case Stmt::BreakStmtClass:
@@ -715,14 +714,19 @@ CIRGenFunction::buildDefaultStmt(const DefaultStmt &S, mlir::Type condType,
   return buildCaseDefaultCascade(&S, condType, caseAttrs);
 }
 
-mlir::LogicalResult
-CIRGenFunction::buildSwitchCase(const SwitchCase &S, mlir::Type condType,
-                                SmallVector<mlir::Attribute, 4> &caseAttrs) {
+mlir::LogicalResult CIRGenFunction::buildSwitchCase(const SwitchCase &S) {
+  assert(!caseAttrsStack.empty() &&
+         "build switch case without seeting case attrs");
+  assert(!condTypeStack.empty() &&
+         "build switch case without specifying the type of the condition");
+
   if (S.getStmtClass() == Stmt::CaseStmtClass)
-    return buildCaseStmt(cast<CaseStmt>(S), condType, caseAttrs);
+    return buildCaseStmt(cast<CaseStmt>(S), condTypeStack.back(),
+                         caseAttrsStack.back());
 
   if (S.getStmtClass() == Stmt::DefaultStmtClass)
-    return buildDefaultStmt(cast<DefaultStmt>(S), condType, caseAttrs);
+    return buildDefaultStmt(cast<DefaultStmt>(S), condTypeStack.back(),
+                            caseAttrsStack.back());
 
   llvm_unreachable("expect case or default stmt");
 }
@@ -987,15 +991,13 @@ mlir::LogicalResult CIRGenFunction::buildWhileStmt(const WhileStmt &S) {
   return mlir::success();
 }
 
-mlir::LogicalResult CIRGenFunction::buildSwitchBody(
-    const Stmt *S, mlir::Type condType,
-    llvm::SmallVector<mlir::Attribute, 4> &caseAttrs) {
+mlir::LogicalResult CIRGenFunction::buildSwitchBody(const Stmt *S) {
   if (auto *compoundStmt = dyn_cast<CompoundStmt>(S)) {
     mlir::Block *lastCaseBlock = nullptr;
     auto res = mlir::success();
     for (auto *c : compoundStmt->body()) {
       if (auto *switchCase = dyn_cast<SwitchCase>(c)) {
-        res = buildSwitchCase(*switchCase, condType, caseAttrs);
+        res = buildSwitchCase(*switchCase);
       } else if (lastCaseBlock) {
         // This means it's a random stmt following up a case, just
         // emit it as part of previous known case.
@@ -1045,12 +1047,16 @@ mlir::LogicalResult CIRGenFunction::buildSwitchStmt(const SwitchStmt &S) {
         [&](mlir::OpBuilder &b, mlir::Location loc, mlir::OperationState &os) {
           currLexScope->setAsSwitch();
 
-          llvm::SmallVector<mlir::Attribute, 4> caseAttrs;
+          caseAttrsStack.push_back({});
+          condTypeStack.push_back(condV.getType());
 
-          res = buildSwitchBody(S.getBody(), condV.getType(), caseAttrs);
+          res = buildSwitchBody(S.getBody());
 
           os.addRegions(currLexScope->getSwitchRegions());
-          os.addAttribute("cases", builder.getArrayAttr(caseAttrs));
+          os.addAttribute("cases", builder.getArrayAttr(caseAttrsStack.back()));
+
+          caseAttrsStack.pop_back();
+          condTypeStack.pop_back();
         });
 
     if (res.failed())

clang/test/CIR/CodeGen/goto.cpp

@@ -288,3 +288,31 @@ 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 case_follow_label(int v) {
+  switch (v) {
+    case 1:
+    label:
+    case 2:
+      action1();
+      break;
+    default:
+      action2();
+      goto label;
+  }
+}
+
+// NOFLAT: cir.func  @case_follow_label
+// NOFLAT: cir.switch
+// NOFLAT: case (equal, 1)
+// NOFLAT: cir.label "label"
+// NOFLAT: cir.yield
+// NOFLAT: case (equal, 2)
+// NOFLAT: cir.call @action1()
+// NOFLAT: cir.break
+// NOFLAT: case (default)
+// NOFLAT: cir.call @action2()
+// NOFLAT: cir.goto "label"