[ADT] llvm::bsearch, binary search for mere mortals

[sammccall]

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.

[sammccall]

Predicate, not compare.

[ilya-biryukov]

include/llvm/ADT/STLExtras.h

1317

Maybe use /2 instead of >>1?

[jkorous]

include/llvm/ADT/STLExtras.h

1315

Wouldn't size_t be better than unsigned for the use-case below?

1317

I'd consider adding assert(Hi > Lo).

[liutianle]

hi sammccall, when I run "check-all", some waring/error print out in STLExtrasTest.cpp as follow. My version is llvm:0ee120077 and clang:d87ee8e678. Could you please fix it or guide me how to fix it?

myLLVM/llvm/unittests/ADT/STLExtrasTest.cpp: In lambda function:
myLLVM/llvm/unittests/ADT/STLExtrasTest.cpp:475:58: warning: comparison of unsigned expression >= 0 is always true [-Wtype-limits]

EXPECT_EQ(5u, bsearch(5, 10, [](unsigned X) { return X >= 0; }));
                                                       ^

myLLVM/llvm/utils/unittest/googletest/include/gtest/gtest_pred_impl.h:77:52: note: in definition of macro 'GTEST_ASSERT_'

if (const ::testing::AssertionResult gtest_ar = (expression)) \
                                                 ^

myLLVM/llvm/utils/unittest/googletest/include/gtest/gtest_pred_impl.h:162:3: note: in expansion of macro 'GTEST_PRED_FORMAT2_'

GTEST_PRED_FORMAT2_(pred_format, v1, v2, GTEST_NONFATAL_FAILURE_)
^

myLLVM/llvm/utils/unittest/googletest/include/gtest/gtest.h:1923:3: note: in expansion of macro 'EXPECT_PRED_FORMAT2'

EXPECT_PRED_FORMAT2(::testing::internal:: \
^

myLLVM/llvm/unittests/ADT/STLExtrasTest.cpp:475:3: note: in expansion of macro 'EXPECT_EQ'

EXPECT_EQ(5u, bsearch(5, 10, [](unsigned X) { return X >= 0; }));
^

myLLVM/llvm/unittests/ADT/STLExtrasTest.cpp: In lambda function:
myLLVM/llvm/unittests/ADT/STLExtrasTest.cpp:483:67: warning: comparison of unsigned expression >= 0 is always true [-Wtype-limits]

bsearch(V.begin(), V.end(), [](unsigned X) { return X >= 0; }));
                                                      ^

myLLVM/llvm/utils/unittest/googletest/include/gtest/gtest_pred_impl.h:77:52: note: in definition of macro 'GTEST_ASSERT_'

if (const ::testing::AssertionResult gtest_ar = (expression)) \
                                                    ^

myLLVM/llvm/utils/unittest/googletest/include/gtest/gtest_pred_impl.h:162:3: note: in expansion of macro 'GTEST_PRED_FORMAT2_'

GTEST_PRED_FORMAT2_(pred_format, v1, v2, GTEST_NONFATAL_FAILURE_)
^

myLLVM/llvm/utils/unittest/googletest/include/gtest/gtest.h:1923:3: note: in expansion of macro 'EXPECT_PRED_FORMAT2'

EXPECT_PRED_FORMAT2(::testing::internal:: \
^

myLLVM/llvm/unittests/ADT/STLExtrasTest.cpp:482:3: note: in expansion of macro 'EXPECT_EQ'

EXPECT_EQ(V.begin(),
^

myLLVM/llvm/unittests/ADT/STLExtrasTest.cpp: In lambda function:
myLLVM/llvm/unittests/ADT/STLExtrasTest.cpp:489:61: warning: comparison of unsigned expression >= 0 is always true [-Wtype-limits]

EXPECT_EQ(V.begin(), bsearch(V, [](unsigned X) { return X >= 0; }));
                                                          ^

myLLVM/llvm/utils/unittest/googletest/include/gtest/gtest_pred_impl.h:77:52: note: in definition of macro 'GTEST_ASSERT_'

if (const ::testing::AssertionResult gtest_ar = (expression)) \
                                                 ^

myLLVM/llvm/utils/unittest/googletest/include/gtest/gtest_pred_impl.h:162:3: note: in expansion of macro 'GTEST_PRED_FORMAT2_'

GTEST_PRED_FORMAT2_(pred_format, v1, v2, GTEST_NONFATAL_FAILURE_)
^

myLLVM/llvm/utils/unittest/googletest/include/gtest/gtest.h:1923:3: note: in expansion of macro 'EXPECT_PRED_FORMAT2'

EXPECT_PRED_FORMAT2(::testing::internal:: \
^

myLLVM/llvm/unittests/ADT/STLExtrasTest.cpp:489:3: note: in expansion of macro 'EXPECT_EQ'

EXPECT_EQ(V.begin(), bsearch(V, [](unsigned X) { return X >= 0; }));
^

---

include/llvm/ADT/STLExtras.h

auto upper_bound(R &&Range, ForwardIt I) -> decltype(adl_begin(Range)) {
   return std::upper_bound(adl_begin(Range), adl_end(Range), I);
 }
 
 template <typename R, typename ForwardIt, typename Compare>
 auto upper_bound(R &&Range, ForwardIt I, Compare C)
     -> decltype(adl_begin(Range)) {
   return std::upper_bound(adl_begin(Range), adl_end(Range), I, 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:
+///   int DawnModernEra = bsearch(1776, 2050, [](int Year){
+///     return Presidents.for(Year).twitterHandle() != None;
+///   });
+///
+/// Note the return value differs from std::binary_search!
+template <typename Compare>
+unsigned bsearch(unsigned Lo, unsigned Hi, Compare C) {
+  while (Lo != Hi) {
+    unsigned Mid = Lo + ((Hi - Lo) >> 1);
+    if (C(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 Compare,
+          typename = decltype(*std::declval<It>()) /* It must be an iterator */>
+It bsearch(It Lo, It Hi, Compare C) {
+  return Lo + bsearch(0, Hi - Lo, [&](unsigned I) { return C(*(Lo + I)); });
+}
+
+/// 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 Compare>
+auto bsearch(R &&Range, Compare C) -> decltype(adl_begin(Range)) {
+  return bsearch(adl_begin(Range), adl_end(Range), C);
+}
+
 /// Wrapper function around std::equal to detect if all elements
 /// in a container are same.
 template <typename R>
 bool is_splat(R &&Range) {
   size_t range_size = size(Range);
   return range_size != 0 && (range_size == 1 ||
          std::equal(adl_begin(Range) + 1, adl_end(Range), adl_begin(Range)));
 }

unittests/ADT/STLExtrasTest.cpp

TEST(STLExtrasTest, to_address) {
   EXPECT_EQ(V3.get(), V4.get());
   EXPECT_EQ(V3.get(), to_address(V3));
   EXPECT_EQ(V4.get(), to_address(V4));
 
   V3.reset(V1);
   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