[llvm] 61211fe - [NFC] Add filters to hasNItems and hasNItemsOrMore

[via llvm-commits]

Author: Tyker
Date: 2020-03-12T10:10:21+01:00
New Revision: 61211fec864917cfc24988c21853d19747e6a1e4

URL: https://github.com/llvm/llvm-project/commit/61211fec864917cfc24988c21853d19747e6a1e4
DIFF: https://github.com/llvm/llvm-project/commit/61211fec864917cfc24988c21853d19747e6a1e4.diff

LOG: [NFC] Add filters to hasNItems and hasNItemsOrMore

Reviewers: lebedev.ri, jdoerfert

Reviewed By: jdoerfert

Subscribers: jdoerfert, dexonsmith, llvm-commits

Tags: #llvm

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

Added: 
    

Modified: 
    llvm/include/llvm/ADT/STLExtras.h

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/ADT/STLExtras.h b/llvm/include/llvm/ADT/STLExtras.h
index cdde0cb67a43..638895ee4dc6 100644
--- a/llvm/include/llvm/ADT/STLExtras.h
+++ b/llvm/include/llvm/ADT/STLExtras.h
@@ -1520,33 +1520,45 @@ decltype(auto) apply_tuple(F &&f, Tuple &&t) {
 
 /// Return true if the sequence [Begin, End) has exactly N items. Runs in O(N)
 /// time. Not meant for use with random-access iterators.
-template <typename IterTy>
+/// Can optionally take a predicate to filter lazily some items.
+template<typename IterTy,
+         typename Pred = bool (*)(const decltype(*std::declval<IterTy>()) &)>
 bool hasNItems(
     IterTy &&Begin, IterTy &&End, unsigned N,
+    Pred &&ShouldBeCounted =
+        [](const decltype(*std::declval<IterTy>()) &) { return true; },
     std::enable_if_t<
         !std::is_same<typename std::iterator_traits<std::remove_reference_t<
                           decltype(Begin)>>::iterator_category,
                       std::random_access_iterator_tag>::value,
         void> * = nullptr) {
-  for (; N; --N, ++Begin)
+  for (; N; ++Begin) {
     if (Begin == End)
       return false; // Too few.
+    N -= ShouldBeCounted(*Begin);
+  }
   return Begin == End;
 }
 
 /// 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.
-template <typename IterTy>
+/// Can optionally take a predicate to filter lazily some items.
+template<typename IterTy,
+         typename Pred = bool (*)(const decltype(*std::declval<IterTy>()) &)>
 bool hasNItemsOrMore(
     IterTy &&Begin, IterTy &&End, unsigned N,
+    Pred &&ShouldBeCounted =
+        [](const decltype(*std::declval<IterTy>()) &) { return true; },
     std::enable_if_t<
         !std::is_same<typename std::iterator_traits<std::remove_reference_t<
                           decltype(Begin)>>::iterator_category,
                       std::random_access_iterator_tag>::value,
         void> * = nullptr) {
-  for (; N; --N, ++Begin)
+  for (; N; ++Begin) {
     if (Begin == End)
       return false; // Too few.
+    N -= ShouldBeCounted(*Begin);
+  }
   return true;
 }