[CIR][NFC] Add getBreakTarget method to cir::BreakOp

clang/include/clang/CIR/Dialect/IR/CIRDialect.h

@@ -30,6 +30,7 @@
 #include "clang/CIR/Dialect/IR/CIROpsDialect.h.inc"
 #include "clang/CIR/Dialect/IR/CIROpsEnums.h"
 #include "clang/CIR/Dialect/IR/CIROpsStructs.h.inc"
+#include "clang/CIR/Dialect/IR/CIRTraits.h"
 #include "clang/CIR/Dialect/IR/CIRTypes.h"
 
 #include "clang/CIR/Interfaces/ASTAttrInterfaces.h"

clang/include/clang/CIR/Dialect/IR/CIROps.td

@@ -1044,6 +1044,8 @@ def CIR_YieldOp : CIR_Op<"yield", [
 // BreakOp
 //===----------------------------------------------------------------------===//
 
+def CIR_Breakable : NativeOpTrait<"cir::Breakable">;
+
 def CIR_BreakOp : CIR_Op<"break", [Terminator]> {
   let summary = "C/C++ `break` statement equivalent";
   let description = [{
@@ -1053,6 +1055,10 @@ def CIR_BreakOp : CIR_Op<"break", [Terminator]> {
   }];
   let assemblyFormat = "attr-dict";
   let hasVerifier = 1;
+  let extraClassDeclaration = [{
+    /// Find the parent operation corresponding to this break.
+    mlir::Operation *getBreakTarget();
+    }];
 }
 
 //===----------------------------------------------------------------------===//
@@ -1993,7 +1999,8 @@ def CIR_CaseOp : CIR_Op<"case", [
 def CIR_SwitchOp : CIR_Op<"switch", [
   SameVariadicOperandSize,
   DeclareOpInterfaceMethods<RegionBranchOpInterface>,
-  RecursivelySpeculatable, AutomaticAllocationScope, NoRegionArguments
+  RecursivelySpeculatable, AutomaticAllocationScope, NoRegionArguments,
+  CIR_Breakable
 ]> {
   let summary = "Switch operation";
   let description = [{
@@ -2251,7 +2258,7 @@ def CIR_BrCondOp : CIR_Op<"brcond", [
 //===----------------------------------------------------------------------===//
 
 class CIR_WhileOpBase<string mnemonic> : CIR_Op<mnemonic, [
-  CIR_LoopOpInterface, NoRegionArguments
+  CIR_LoopOpInterface, CIR_Breakable, NoRegionArguments
 ]> {
   defvar isWhile = !eq(mnemonic, "while");
   let summary = "C/C++ " # !if(isWhile, "while", "do-while") # " loop";
@@ -2327,7 +2334,7 @@ def CIR_DoWhileOp : CIR_WhileOpBase<"do"> {
 // ForOp
 //===----------------------------------------------------------------------===//
 
-def CIR_ForOp : CIR_Op<"for", [CIR_LoopOpInterface, NoRegionArguments]> {
+def CIR_ForOp : CIR_Op<"for", [CIR_LoopOpInterface, CIR_Breakable, NoRegionArguments]> {
   let summary = "C/C++ for loop counterpart";
   let description = [{
     Represents a C/C++ for loop. It consists of three regions:

clang/include/clang/CIR/Dialect/IR/CIRTraits.h

@@ -0,0 +1,30 @@
+//===- CIRTraits.h - MLIR CIR Traits ----------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file declares the traits in the CIR dialect.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_CLANG_CIR_DIALECT_IR_CIRTRAITS_H
+#define LLVM_CLANG_CIR_DIALECT_IR_CIRTRAITS_H
+
+#include "mlir/Dialect/Traits.h"
+
+namespace mlir {
+namespace OpTrait {
+namespace cir {
+
+/// This class provides the API for ops that are known to be terminators.
+template <typename ConcreteType>
+class Breakable : public TraitBase<ConcreteType, Breakable> {};
+
+} // namespace cir
+} // namespace OpTrait
+} // namespace mlir
+
+#endif // LLVM_CLANG_CIR_DIALECT_IR_CIRTRAITS_H

clang/lib/CIR/Dialect/IR/CIRDialect.cpp

@@ -324,10 +324,13 @@ void cir::AllocaOp::build(::mlir::OpBuilder &odsBuilder,
 // BreakOp
 //===----------------------------------------------------------------------===//
 
+mlir::Operation *cir::BreakOp::getBreakTarget() {
+  return getOperation()->getParentWithTrait<mlir::OpTrait::cir::Breakable>();
+}
+
 LogicalResult cir::BreakOp::verify() {
-  if (!getOperation()->getParentOfType<LoopOpInterface>() &&
-      !getOperation()->getParentOfType<SwitchOp>())
-    return emitOpError("must be within a loop or switch");
+  if (!getOperation()->getParentWithTrait<mlir::OpTrait::cir::Breakable>())
+    return emitOpError("must be within a breakable operation");
   return success();
 }
 

clang/test/CIR/IR/invalid.cir

@@ -53,7 +53,7 @@ cir.func @yield0() {
 cir.func @yieldbreak() {
   %0 = cir.const #true
   cir.if %0 {
-    cir.break // expected-error {{op must be within a loop or switch}}
+    cir.break // expected-error {{op must be within a breakable operation}}
   }
   cir.return
 }