fixing binary search for cases when element is not in array

Subscribers: kubabrecka

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

llvm-svn: 287078

3 files changed, 52 insertions, 12 deletions

diff --git a/compiler-rt/lib/sanitizer_common/sanitizer_common.h b/compiler-rt/lib/sanitizer_common/sanitizer_common.h
index 779e96fe638..63956aed94d 100644
--- a/compiler-rt/lib/sanitizer_common/sanitizer_common.h
+++ b/compiler-rt/lib/sanitizer_common/sanitizer_common.h
@@ -635,20 +635,19 @@ void InternalSort(Container *v, uptr size, Compare comp) {
   }
 }
 
+// Works like std::lower_bound: finds the first element that is not less
+// than the val.
 template<class Container, class Value, class Compare>
 uptr InternalBinarySearch(const Container &v, uptr first, uptr last,
                           const Value &val, Compare comp) {
-  uptr not_found = last + 1;
-  while (last >= first) {
+  while (last > first) {
     uptr mid = (first + last) / 2;
     if (comp(v[mid], val))
       first = mid + 1;
-    else if (comp(val, v[mid]))
-      last = mid - 1;
     else
-      return mid;
+      last = mid;
   }
-  return not_found;
+  return first;
 }
 
 // Represents a binary loaded into virtual memory (e.g. this can be an
diff --git a/compiler-rt/lib/sanitizer_common/sanitizer_stackdepot.cc b/compiler-rt/lib/sanitizer_common/sanitizer_stackdepot.cc
index 985193d1ed5..2ec7a82d162 100644
--- a/compiler-rt/lib/sanitizer_common/sanitizer_stackdepot.cc
+++ b/compiler-rt/lib/sanitizer_common/sanitizer_stackdepot.cc
@@ -155,7 +155,7 @@ StackTrace StackDepotReverseMap::Get(u32 id) {
   IdDescPair pair = {id, nullptr};
   uptr idx = InternalBinarySearch(map_, 0, map_.size(), pair,
                                   IdDescPair::IdComparator);
-  if (idx > map_.size())
+  if (idx > map_.size() || map_[idx].id != id)
     return StackTrace();
   return map_[idx].desc->load();
 }
diff --git a/compiler-rt/lib/sanitizer_common/tests/sanitizer_common_test.cc b/compiler-rt/lib/sanitizer_common/tests/sanitizer_common_test.cc
index 6fc308ad14d..0980cf4e47c 100644
--- a/compiler-rt/lib/sanitizer_common/tests/sanitizer_common_test.cc
+++ b/compiler-rt/lib/sanitizer_common/tests/sanitizer_common_test.cc
@@ -10,6 +10,8 @@
 // This file is a part of ThreadSanitizer/AddressSanitizer runtime.
 //
 //===----------------------------------------------------------------------===//
+#include <algorithm>
+
 #include "sanitizer_common/sanitizer_allocator_internal.h"
 #include "sanitizer_common/sanitizer_common.h"
 #include "sanitizer_common/sanitizer_flags.h"
@@ -172,13 +174,52 @@ bool UptrLess(uptr a, uptr b) {
 
 TEST(SanitizerCommon, InternalBinarySearch) {
   static const uptr kSize = 5;
-  uptr arr[kSize];
-  for (uptr i = 0; i < kSize; i++) arr[i] = i * i;
+  int arr[kSize];
+  arr[0] = 1;
+  arr[1] = 3;
+  arr[2] = 5;
+  arr[3] = 7;
+  arr[4] = 11;
+
+  EXPECT_EQ(0u, InternalBinarySearch(arr, 0, kSize, 0, UptrLess));
+  EXPECT_EQ(0u, InternalBinarySearch(arr, 0, kSize, 1, UptrLess));
+  EXPECT_EQ(1u, InternalBinarySearch(arr, 0, kSize, 2, UptrLess));
+  EXPECT_EQ(1u, InternalBinarySearch(arr, 0, kSize, 3, UptrLess));
+  EXPECT_EQ(2u, InternalBinarySearch(arr, 0, kSize, 4, UptrLess));
+  EXPECT_EQ(2u, InternalBinarySearch(arr, 0, kSize, 5, UptrLess));
+  EXPECT_EQ(3u, InternalBinarySearch(arr, 0, kSize, 6, UptrLess));
+  EXPECT_EQ(3u, InternalBinarySearch(arr, 0, kSize, 7, UptrLess));
+  EXPECT_EQ(4u, InternalBinarySearch(arr, 0, kSize, 8, UptrLess));
+  EXPECT_EQ(4u, InternalBinarySearch(arr, 0, kSize, 9, UptrLess));
+  EXPECT_EQ(4u, InternalBinarySearch(arr, 0, kSize, 10, UptrLess));
+  EXPECT_EQ(4u, InternalBinarySearch(arr, 0, kSize, 11, UptrLess));
+  EXPECT_EQ(5u, InternalBinarySearch(arr, 0, kSize, 12, UptrLess));
+}
 
-  for (uptr i = 0; i < kSize; i++)
-    ASSERT_EQ(InternalBinarySearch(arr, 0, kSize, i * i, UptrLess), i);
+TEST(SanitizerCommon, InternalBinarySearchVsLowerBound) {
+  std::vector<int> data;
+  auto create_item = [] (size_t i, size_t j) {
+    auto v = i * 10000 + j;
+    return ((v << 6) + (v >> 6) + 0x9e3779b9) % 100;
+  };
+  for (size_t i = 0; i < 1000; ++i) {
+    data.resize(i);
+    for (size_t j = 0; j < i; ++j) {
+      data[j] = create_item(i, j);
+    }
 
-  ASSERT_EQ(InternalBinarySearch(arr, 0, kSize, 7, UptrLess), kSize + 1);
+    std::sort(data.begin(), data.end());
+
+    for (size_t j = 0; j < i; ++j) {
+      int val = create_item(i, j);
+      for (auto to_find : {val - 1, val, val + 1}) {
+        uptr expected =
+            std::lower_bound(data.begin(), data.end(), to_find) - data.begin();
+        EXPECT_EQ(expected, InternalBinarySearch(data.data(), 0, data.size(),
+                                                 to_find, std::less<int>()));
+      }
+    }
+  }
 }
 
 #if SANITIZER_LINUX && !SANITIZER_ANDROID