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
Change-Id: I6310840996a99fa61c4e5099e72d7fc9f23b4912
diff --git a/third_party/mlir/include/mlir/Support/STLExtras.h b/third_party/mlir/include/mlir/Support/STLExtras.h
index 2ab890e..24e2ac6 100644
--- a/third_party/mlir/include/mlir/Support/STLExtras.h
+++ b/third_party/mlir/include/mlir/Support/STLExtras.h
@@ -193,6 +193,11 @@
});
}
+/// 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/third_party/mlir/lib/Dialect/AffineOps/AffineOps.cpp b/third_party/mlir/lib/Dialect/AffineOps/AffineOps.cpp
index 6d90464..76417bb 100644
--- a/third_party/mlir/lib/Dialect/AffineOps/AffineOps.cpp
+++ b/third_party/mlir/lib/Dialect/AffineOps/AffineOps.cpp
@@ -1337,8 +1337,7 @@
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/third_party/mlir/lib/Dialect/SPIRV/SPIRVOps.cpp b/third_party/mlir/lib/Dialect/SPIRV/SPIRVOps.cpp
index 85e22a5..c47d4b8 100644
--- a/third_party/mlir/lib/Dialect/SPIRV/SPIRVOps.cpp
+++ b/third_party/mlir/lib/Dialect/SPIRV/SPIRVOps.cpp
@@ -1418,7 +1418,7 @@
/// 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/third_party/mlir/lib/IR/Module.cpp b/third_party/mlir/lib/IR/Module.cpp
index 9f3bdaa..f55f960 100644
--- a/third_party/mlir/lib/IR/Module.cpp
+++ b/third_party/mlir/lib/IR/Module.cpp
@@ -87,7 +87,7 @@
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.
