[ADT] llvm::bsearch, binary search for mere mortals
Summary:
Add to STLExtras a binary search function with a simple mental model:
You provide a range and a predicate which is true above a certain point.
bsearch() tells you that point.
Overloads are provided for integers, iterators, and containers.
This is more suitable than std:: alternatives in many cases:
- std::binary_search only indicates presence/absence
- upper_bound/lower_bound give you the opportunity to pick the wrong one
- all of the options have confusing names and definitions when your predicate
doesn't have simple "less than" semantics
- all of the options require iterators
- we plumb around a useless `value` parameter that should be a lambda capture
The API is inspired by Go's standard library, but we add an extra parameter as
well as some overloads and templates to show how clever C++ is.
Reviewers: ilya-biryukov, gribozavr
Subscribers: dexonsmith, kristina, llvm-commits
Tags: #llvm
Differential Revision: https://reviews.llvm.org/D60779
git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@358540 91177308-0d34-0410-b5e6-96231b3b80d8
diff --git a/include/llvm/ADT/STLExtras.h b/include/llvm/ADT/STLExtras.h
index 8a5f6a8..d63b819 100644
--- a/include/llvm/ADT/STLExtras.h
+++ b/include/llvm/ADT/STLExtras.h
@@ -1304,6 +1304,47 @@
return std::upper_bound(adl_begin(Range), adl_end(Range),
std::forward<T>(Value), C);
}
+
+/// Binary search for the first index where a predicate is true.
+/// Returns the first I in [Lo, Hi) where C(I) is true, or Hi if it never is.
+/// Requires that C is always false below some limit, and always true above it.
+///
+/// Example:
+/// size_t DawnModernEra = bsearch(1776, 2050, [](size_t Year){
+/// return Presidents.for(Year).twitterHandle() != None;
+/// });
+///
+/// Note the return value differs from std::binary_search!
+template <typename Predicate>
+size_t bsearch(size_t Lo, size_t Hi, Predicate P) {
+ while (Lo != Hi) {
+ assert(Hi > Lo);
+ size_t Mid = Lo + (Hi - Lo) / 2;
+ if (P(Mid))
+ Hi = Mid;
+ else
+ Lo = Mid + 1;
+ }
+ return Hi;
+}
+
+/// Binary search for the first iterator where a predicate is true.
+/// Returns the first I in [Lo, Hi) where C(*I) is true, or Hi if it never is.
+/// Requires that C is always false below some limit, and always true above it.
+template <typename It, typename Predicate,
+ typename Val = decltype(*std::declval<It>())>
+It bsearch(It Lo, It Hi, Predicate P) {
+ return std::lower_bound(Lo, Hi, 0u,
+ [&](const Val &V, unsigned) { return !P(V); });
+}
+
+/// Binary search for the first iterator in a range where a predicate is true.
+/// Requires that C is always false below some limit, and always true above it.
+template <typename R, typename Predicate>
+auto bsearch(R &&Range, Predicate P) -> decltype(adl_begin(Range)) {
+ return bsearch(adl_begin(Range), adl_end(Range), P);
+}
+
/// Wrapper function around std::equal to detect if all elements
/// in a container are same.
template <typename R>
diff --git a/unittests/ADT/STLExtrasTest.cpp b/unittests/ADT/STLExtrasTest.cpp
index 26196fb..24844f4 100644
--- a/unittests/ADT/STLExtrasTest.cpp
+++ b/unittests/ADT/STLExtrasTest.cpp
@@ -469,4 +469,25 @@
EXPECT_EQ(V1, to_address(V3));
}
+TEST(STLExtrasTest, bsearch) {
+ // Integer version.
+ EXPECT_EQ(7u, bsearch(5, 10, [](unsigned X) { return X >= 7; }));
+ EXPECT_EQ(5u, bsearch(5, 10, [](unsigned X) { return X >= 0; }));
+ EXPECT_EQ(10u, bsearch(5, 10, [](unsigned X) { return X >= 50; }));
+
+ // Iterator version.
+ std::vector<int> V = {1, 3, 5, 7, 9};
+ EXPECT_EQ(V.begin() + 3,
+ bsearch(V.begin(), V.end(), [](unsigned X) { return X >= 7; }));
+ EXPECT_EQ(V.begin(),
+ bsearch(V.begin(), V.end(), [](unsigned X) { return X >= 0; }));
+ EXPECT_EQ(V.end(),
+ bsearch(V.begin(), V.end(), [](unsigned X) { return X >= 50; }));
+
+ // Range version.
+ EXPECT_EQ(V.begin() + 3, bsearch(V, [](unsigned X) { return X >= 7; }));
+ EXPECT_EQ(V.begin(), bsearch(V, [](unsigned X) { return X >= 0; }));
+ EXPECT_EQ(V.end(), bsearch(V, [](unsigned X) { return X >= 50; }));
+}
+
} // namespace
