5 files changed, 62 insertions, 15 deletions

diff --git a/compiler-rt/lib/asan/asan_allocator.cc b/compiler-rt/lib/asan/asan_allocator.cc
index 60d8671c4af..f9da120cc10 100644
--- a/compiler-rt/lib/asan/asan_allocator.cc
+++ b/compiler-rt/lib/asan/asan_allocator.cc
@@ -362,9 +362,6 @@ class MallocInfo {
     return FindChunkByAddr(addr);
   }
 
-  // TODO(glider): AllocationSize() may become very slow if the size of
-  // page_groups_ grows. This can be fixed by increasing kMinMmapSize,
-  // but a better solution is to speed up the search somehow.
   size_t AllocationSize(uintptr_t ptr) {
     if (!ptr) return 0;
     ScopedLock lock(&mu_);
@@ -413,12 +410,32 @@ class MallocInfo {
 
  private:
   PageGroup *FindPageGroupUnlocked(uintptr_t addr) {
-    for (int i = 0; i < n_page_groups_; i++) {
-      PageGroup *g = page_groups_[i];
-      if (g->InRange(addr)) {
-        return g;
+    int n = n_page_groups_;
+    // If the page groups are not sorted yet, sort them.
+    if (n_sorterd_page_groups < n) {
+      SortArray((uintptr_t*)page_groups_, n);
+      n_sorterd_page_groups = n;
+    }
+    // Binray search over the page groups.
+    int beg = 0, end = n;
+    while (beg < end) {
+      int med = (beg + end) / 2;
+      uintptr_t g = (uintptr_t)page_groups_[med];
+      if (addr > g) {
+        // 'g' points to the end of the group, so 'addr'
+        // may not belong to page_groups_[med] or any previous group.
+        beg = med + 1;
+      } else {
+        // 'addr' may belong to page_groups_[med] or a previous group.
+        end = med;
       }
     }
+    if (beg >= n)
+      return NULL;
+    PageGroup *g = page_groups_[beg];
+    CHECK(g);
+    if (g->InRange(addr))
+      return g;
     return NULL;
   }
 
@@ -546,6 +563,7 @@ class MallocInfo {
 
   PageGroup *page_groups_[kMaxAvailableRam / kMinMmapSize];
   int n_page_groups_;  // atomic
+  int n_sorterd_page_groups;
 };
 
 static MallocInfo malloc_info(LINKER_INITIALIZED);
diff --git a/compiler-rt/lib/asan/asan_internal.h b/compiler-rt/lib/asan/asan_internal.h
index 8acd36274e1..42e8829f3bb 100644
--- a/compiler-rt/lib/asan/asan_internal.h
+++ b/compiler-rt/lib/asan/asan_internal.h
@@ -205,6 +205,8 @@ void Report(const char *format, ...);
 template<class T> T Min(T a, T b) { return a < b ? a : b; }
 template<class T> T Max(T a, T b) { return a > b ? a : b; }
 
+void SortArray(uintptr_t *array, size_t size);
+
 // asan_poisoning.cc
 // Poisons the shadow memory for "size" bytes starting from "addr".
 void PoisonShadow(uintptr_t addr, size_t size, uint8_t value);
diff --git a/compiler-rt/lib/asan/asan_posix.cc b/compiler-rt/lib/asan/asan_posix.cc
index 43a5ffee058..cc7a128254d 100644
--- a/compiler-rt/lib/asan/asan_posix.cc
+++ b/compiler-rt/lib/asan/asan_posix.cc
@@ -30,6 +30,10 @@
 #include <sys/atomics.h>
 #endif
 
+// Should not add dependency on libstdc++,
+// since most of the stuff here is inlinable.
+#include <algorithm>
+
 namespace __asan {
 
 static void MaybeInstallSigaction(int signum,
@@ -112,6 +116,10 @@ int AtomicInc(int *a) {
 #endif
 }
 
+void SortArray(uintptr_t *array, size_t size) {
+  std::sort(array, array + size);
+}
+
 // ---------------------- TSD ---------------- {{{1
 
 static pthread_key_t tsd_key;
diff --git a/compiler-rt/lib/asan/asan_win.cc b/compiler-rt/lib/asan/asan_win.cc
index d47ff08e9e0..5df110330b6 100644
--- a/compiler-rt/lib/asan/asan_win.cc
+++ b/compiler-rt/lib/asan/asan_win.cc
@@ -26,6 +26,10 @@
 #include "asan_procmaps.h"
 #include "asan_thread.h"
 
+// Should not add dependency on libstdc++,
+// since most of the stuff here is inlinable.
+#include <algorithm>
+
 namespace __asan {
 
 // ---------------------- Memory management ---------------- {{{1
@@ -267,6 +271,10 @@ int Atexit(void (*function)(void)) {
   return atexit(function);
 }
 
+void SortArray(uintptr_t *array, size_t size) {
+  std::sort(array, array + size);
+}
+
 }  // namespace __asan
 
 #endif  // _WIN32
diff --git a/compiler-rt/lib/asan/tests/asan_interface_test.cc b/compiler-rt/lib/asan/tests/asan_interface_test.cc
index 6f6a9636e58..43889058b65 100644
--- a/compiler-rt/lib/asan/tests/asan_interface_test.cc
+++ b/compiler-rt/lib/asan/tests/asan_interface_test.cc
@@ -363,12 +363,23 @@ TEST(AddressSanitizerInterface, SetErrorReportCallbackTest) {
   __asan_set_error_report_callback(NULL);
 }
 
-TEST(AddressSanitizerInterface, DISABLED_GetOwnershipStressTest) {
-  std::vector<void *> v;
-  for (size_t i = 0; i < 3000; i++)
-    v.push_back(malloc(i * 1000));
-  for (size_t i = 0; i < 1000000; i++)
-    __asan_get_ownership(&v);
-  for (size_t i = 0, n = v.size(); i < n; i++)
-    free(v[i]);
+TEST(AddressSanitizerInterface, GetOwnershipStressTest) {
+  std::vector<char *> pointers;
+  std::vector<size_t> sizes;
+  const size_t kNumMallocs =
+      (__WORDSIZE <= 32 || ASAN_LOW_MEMORY) ? 1 << 10 : 1 << 14;
+  for (size_t i = 0; i < kNumMallocs; i++) {
+    size_t size = i * 100 + 1;
+    pointers.push_back((char*)malloc(size));
+    sizes.push_back(size);
+  }
+  for (size_t i = 0; i < 4000000; i++) {
+    EXPECT_FALSE(__asan_get_ownership(&pointers));
+    EXPECT_FALSE(__asan_get_ownership((void*)0x1234));
+    size_t idx = i % kNumMallocs;
+    EXPECT_TRUE(__asan_get_ownership(pointers[idx]));
+    EXPECT_EQ(sizes[idx], __asan_get_allocated_size(pointers[idx]));
+  }
+  for (size_t i = 0, n = pointers.size(); i < n; i++)
+    free(pointers[i]);
 }