Add a utility accessor 'has_single_element' for ranges.

This provides an easy way to check if a range has a single element.

PiperOrigin-RevId: 277544647

4 files changed, 8 insertions, 4 deletions

diff --git a/mlir/include/mlir/Support/STLExtras.h b/mlir/include/mlir/Support/STLExtras.h
index 2ab890eb2e6..24e2ac6a5a3 100644
--- a/mlir/include/mlir/Support/STLExtras.h
+++ b/mlir/include/mlir/Support/STLExtras.h
@@ -193,6 +193,11 @@ template <typename ContainerTy> auto make_second_range(ContainerTy &&c) {
       });
 }
 
+/// Returns true of the given range only contains a single element.
+template <typename ContainerTy> bool has_single_element(ContainerTy &&c) {
+  auto it = std::begin(c), e = std::end(c);
+  return it != e && std::next(it) == e;
+}
 } // end namespace mlir
 
 // Allow tuples to be usable as DenseMap keys.
diff --git a/mlir/lib/Dialect/AffineOps/AffineOps.cpp b/mlir/lib/Dialect/AffineOps/AffineOps.cpp
index 6d90464e632..76417bb3289 100644
--- a/mlir/lib/Dialect/AffineOps/AffineOps.cpp
+++ b/mlir/lib/Dialect/AffineOps/AffineOps.cpp
@@ -1337,8 +1337,7 @@ struct AffineForEmptyLoopFolder : public OpRewritePattern<AffineForOp> {
   PatternMatchResult matchAndRewrite(AffineForOp forOp,
                                      PatternRewriter &rewriter) const override {
     // Check that the body only contains a terminator.
-    auto *body = forOp.getBody();
-    if (std::next(body->begin()) != body->end())
+    if (!has_single_element(*forOp.getBody()))
       return matchFailure();
     rewriter.eraseOp(forOp);
     return matchSuccess();
diff --git a/mlir/lib/Dialect/SPIRV/SPIRVOps.cpp b/mlir/lib/Dialect/SPIRV/SPIRVOps.cpp
index 85e22a5c391..c47d4b87027 100644
--- a/mlir/lib/Dialect/SPIRV/SPIRVOps.cpp
+++ b/mlir/lib/Dialect/SPIRV/SPIRVOps.cpp
@@ -1418,7 +1418,7 @@ static void print(spirv::LoopOp loopOp, OpAsmPrinter &printer) {
 /// given `dstBlock`.
 static inline bool hasOneBranchOpTo(Block &srcBlock, Block &dstBlock) {
   // Check that there is only one op in the `srcBlock`.
-  if (srcBlock.empty() || std::next(srcBlock.begin()) != srcBlock.end())
+  if (!has_single_element(srcBlock))
     return false;
 
   auto branchOp = dyn_cast<spirv::BranchOp>(srcBlock.back());
diff --git a/mlir/lib/IR/Module.cpp b/mlir/lib/IR/Module.cpp
index 9f3bdaa854c..f55f960ffa2 100644
--- a/mlir/lib/IR/Module.cpp
+++ b/mlir/lib/IR/Module.cpp
@@ -87,7 +87,7 @@ LogicalResult ModuleOp::verify() {
   auto &bodyRegion = getOperation()->getRegion(0);
 
   // The body must contain a single basic block.
-  if (bodyRegion.empty() || std::next(bodyRegion.begin()) != bodyRegion.end())
+  if (!has_single_element(bodyRegion))
     return emitOpError("expected body region to have a single block");
 
   // Check that the body has no block arguments.