- Add hasNItemsOrLess and container variants of hasNItems, hasNItemsOrMore, and hasNItemsOrLess

- Fixed a bug in hasNItems()
- Extend the STLExtras unit test to test hasSingleElement() and hasNItems() and friends.

Differential Revision: https://reviews.llvm.org/D82232

llvm/include/llvm/ADT/STLExtras.h

@@ -267,10 +267,10 @@ constexpr bool empty(const T &RangeOrContainer) {
   return adl_begin(RangeOrContainer) == adl_end(RangeOrContainer);
 }
 
-/// Returns true of the given range only contains a single element.
-template <typename ContainerTy> bool hasSingleElement(ContainerTy &&c) {
-  auto it = std::begin(c), e = std::end(c);
-  return it != e && std::next(it) == e;
+/// Returns true if the given container only contains a single element.
+template <typename ContainerTy> bool hasSingleElement(ContainerTy &&C) {
+  auto B = std::begin(C), E = std::end(C);
+  return B != E && std::next(B) == E;
 }
 
 /// Return a range covering \p RangeOrContainer with the first N elements
@@ -1919,12 +1919,15 @@ bool hasNItems(
       return false; // Too few.
     N -= ShouldBeCounted(*Begin);
   }
-  return Begin == End;
+  for (; Begin != End; ++Begin)
+    if (ShouldBeCounted(*Begin))
+      return false; // Too many.
+  return true;
 }
 
 /// Return true if the sequence [Begin, End) has N or more items. Runs in O(N)
 /// time. Not meant for use with random-access iterators.
-/// Can optionally take a predicate to filter lazily some items.
+/// Can optionally take a predicate to lazily filter some items.
 template<typename IterTy,
          typename Pred = bool (*)(const decltype(*std::declval<IterTy>()) &)>
 bool hasNItemsOrMore(
@@ -1944,6 +1947,36 @@ bool hasNItemsOrMore(
   return true;
 }
 
+/// Returns true if the sequence [Begin, End) has N or less items. Can
+/// optionally take a predicate to lazily filter some items.
+template <typename IterTy,
+          typename Pred = bool (*)(const decltype(*std::declval<IterTy>()) &)>
+bool hasNItemsOrLess(
+    IterTy &&Begin, IterTy &&End, unsigned N,
+    Pred &&ShouldBeCounted = [](const decltype(*std::declval<IterTy>()) &) {
+      return true;
+    }) {
+  assert(N != std::numeric_limits<unsigned>::max());
+  return !hasNItemsOrMore(Begin, End, N + 1, ShouldBeCounted);
+}
+
+/// Returns true if the given container has exactly N items
+template <typename ContainerTy> bool hasNItems(ContainerTy &&C, unsigned N) {
+  return hasNItems(std::begin(C), std::end(C), N);
+}
+
+/// Returns true if the given container has N or more items
+template <typename ContainerTy>
+bool hasNItemsOrMore(ContainerTy &&C, unsigned N) {
+  return hasNItemsOrMore(std::begin(C), std::end(C), N);
+}
+
+/// Returns true if the given container has N or less items
+template <typename ContainerTy>
+bool hasNItemsOrLess(ContainerTy &&C, unsigned N) {
+  return hasNItemsOrLess(std::begin(C), std::end(C), N);
+}
+
 /// Returns a raw pointer that represents the same address as the argument.
 ///
 /// This implementation can be removed once we move to C++20 where it's defined

llvm/unittests/ADT/STLExtrasTest.cpp

@@ -490,4 +490,82 @@ TEST(STLExtrasTest, partition_point) {
   EXPECT_EQ(V.end(), partition_point(V, [](unsigned X) { return X < 50; }));
 }
 
+TEST(STLExtrasTest, hasSingleElement) {
+  const std::vector<int> V0 = {}, V1 = {1}, V2 = {1, 2};
+  const std::vector<int> V10(10);
+
+  EXPECT_EQ(hasSingleElement(V0), false);
+  EXPECT_EQ(hasSingleElement(V1), true);
+  EXPECT_EQ(hasSingleElement(V2), false);
+  EXPECT_EQ(hasSingleElement(V10), false);
+}
+
+TEST(STLExtrasTest, hasNItems) {
+  const std::list<int> V0 = {}, V1 = {1}, V2 = {1, 2};
+  const std::list<int> V3 = {1, 3, 5};
+
+  EXPECT_TRUE(hasNItems(V0, 0));
+  EXPECT_FALSE(hasNItems(V0, 2));
+  EXPECT_TRUE(hasNItems(V1, 1));
+  EXPECT_FALSE(hasNItems(V1, 2));
+
+  EXPECT_TRUE(hasNItems(V3.begin(), V3.end(), 3, [](int x) { return x < 10; }));
+  EXPECT_TRUE(hasNItems(V3.begin(), V3.end(), 0, [](int x) { return x > 10; }));
+  EXPECT_TRUE(hasNItems(V3.begin(), V3.end(), 2, [](int x) { return x < 5; }));
+}
+
+TEST(STLExtras, hasNItemsOrMore) {
+  const std::list<int> V0 = {}, V1 = {1}, V2 = {1, 2};
+  const std::list<int> V3 = {1, 3, 5};
+
+  EXPECT_TRUE(hasNItemsOrMore(V1, 1));
+  EXPECT_FALSE(hasNItemsOrMore(V1, 2));
+
+  EXPECT_TRUE(hasNItemsOrMore(V2, 1));
+  EXPECT_TRUE(hasNItemsOrMore(V2, 2));
+  EXPECT_FALSE(hasNItemsOrMore(V2, 3));
+
+  EXPECT_TRUE(hasNItemsOrMore(V3, 3));
+  EXPECT_FALSE(hasNItemsOrMore(V3, 4));
+
+  EXPECT_TRUE(
+      hasNItemsOrMore(V3.begin(), V3.end(), 3, [](int x) { return x < 10; }));
+  EXPECT_FALSE(
+      hasNItemsOrMore(V3.begin(), V3.end(), 3, [](int x) { return x > 10; }));
+  EXPECT_TRUE(
+      hasNItemsOrMore(V3.begin(), V3.end(), 2, [](int x) { return x < 5; }));
+}
+
+TEST(STLExtras, hasNItemsOrLess) {
+  const std::list<int> V0 = {}, V1 = {1}, V2 = {1, 2};
+  const std::list<int> V3 = {1, 3, 5};
+
+  EXPECT_TRUE(hasNItemsOrLess(V0, 0));
+  EXPECT_TRUE(hasNItemsOrLess(V0, 1));
+  EXPECT_TRUE(hasNItemsOrLess(V0, 2));
+
+  EXPECT_FALSE(hasNItemsOrLess(V1, 0));
+  EXPECT_TRUE(hasNItemsOrLess(V1, 1));
+  EXPECT_TRUE(hasNItemsOrLess(V1, 2));
+
+  EXPECT_FALSE(hasNItemsOrLess(V2, 0));
+  EXPECT_FALSE(hasNItemsOrLess(V2, 1));
+  EXPECT_TRUE(hasNItemsOrLess(V2, 2));
+  EXPECT_TRUE(hasNItemsOrLess(V2, 3));
+
+  EXPECT_FALSE(hasNItemsOrLess(V3, 0));
+  EXPECT_FALSE(hasNItemsOrLess(V3, 1));
+  EXPECT_FALSE(hasNItemsOrLess(V3, 2));
+  EXPECT_TRUE(hasNItemsOrLess(V3, 3));
+  EXPECT_TRUE(hasNItemsOrLess(V3, 4));
+
+  EXPECT_TRUE(
+      hasNItemsOrLess(V3.begin(), V3.end(), 1, [](int x) { return x == 1; }));
+  EXPECT_TRUE(
+      hasNItemsOrLess(V3.begin(), V3.end(), 2, [](int x) { return x < 5; }));
+  EXPECT_TRUE(
+      hasNItemsOrLess(V3.begin(), V3.end(), 5, [](int x) { return x < 5; }));
+  EXPECT_FALSE(
+      hasNItemsOrLess(V3.begin(), V3.end(), 2, [](int x) { return x < 10; }));
+}
 } // namespace