2 files changed, 112 insertions, 10 deletions

diff --git a/compiler-rt/lib/sanitizer_common/sanitizer_allocator64.h b/compiler-rt/lib/sanitizer_common/sanitizer_allocator64.h
index 4b25158a41d..0bac2374202 100644
--- a/compiler-rt/lib/sanitizer_common/sanitizer_allocator64.h
+++ b/compiler-rt/lib/sanitizer_common/sanitizer_allocator64.h
@@ -97,12 +97,17 @@ class SizeClassAllocator64 {
     CHECK_EQ(AllocBeg(), reinterpret_cast<uptr>(MmapFixedNoReserve(
              AllocBeg(), AllocSize())));
   }
-  NOINLINE
-  void *Allocate(uptr size) {
-    CHECK_LE(size, SizeClassMap::kMaxSize);
+
+  bool CanAllocate(uptr size, uptr alignment) {
+    return size <= SizeClassMap::kMaxSize &&
+      alignment <= SizeClassMap::kMaxSize;
+  }
+
+  void *Allocate(uptr size, uptr alignment) {
+    CHECK(CanAllocate(size, alignment));
     return AllocateBySizeClass(SizeClassMap::ClassID(size));
   }
-  NOINLINE
+
   void Deallocate(void *p) {
     CHECK(PointerIsMine(p));
     DeallocateBySizeClass(p, GetSizeClass(p));
@@ -236,7 +241,8 @@ class LargeMmapAllocator {
   void Init() {
     internal_memset(this, 0, sizeof(*this));
   }
-  void *Allocate(uptr size) {
+  void *Allocate(uptr size, uptr alignment) {
+    CHECK_LE(alignment, kPageSize);  // Not implemented. Do we need it?
     uptr map_size = RoundUpMapSize(size);
     void *map = MmapOrDie(map_size, "LargeMmapAllocator");
     void *res = reinterpret_cast<void*>(reinterpret_cast<uptr>(map)
@@ -318,6 +324,63 @@ class LargeMmapAllocator {
   uptr lock_;  // FIXME
 };
 
+// This class implements a complete memory allocator by using two
+// internal allocators:
+// PrimaryAllocator is efficient, but may not allocate some sizes (alignments).
+//  When allocating 2^x bytes it should return 2^x aligned chunk.
+// SecondaryAllocator can allocate anything, but is not efficient.
+template <class PrimaryAllocator, class SecondaryAllocator>
+class CombinedAllocator {
+ public:
+  void Init() {
+    primary_.Init();
+    secondary_.Init();
+  }
+
+  void *Allocate(uptr size, uptr alignment) {
+    CHECK_GT(size, 0);
+    if (alignment > 8)
+      size = RoundUpTo(size, alignment);
+    void *res;
+    if (primary_.CanAllocate(size, alignment))
+      res = primary_.Allocate(size, alignment);
+    else
+      res = secondary_.Allocate(size, alignment);
+    if (alignment > 8)
+      CHECK_EQ(reinterpret_cast<uptr>(res) & (alignment - 1), 0);
+    return res;
+  }
+
+  void Deallocate(void *p) {
+    if (primary_.PointerIsMine(p))
+      primary_.Deallocate(p);
+    else
+      secondary_.Deallocate(p);
+  }
+
+  bool PointerIsMine(void *p) {
+    if (primary_.PointerIsMine(p))
+      return true;
+    return secondary_.PointerIsMine(p);
+  }
+
+  void *GetMetaData(void *p) {
+    if (primary_.PointerIsMine(p))
+      return primary_.GetMetaData(p);
+    return secondary_.GetMetaData(p);
+  }
+
+  uptr TotalMemoryUsed() {
+    return primary_.TotalMemoryUsed() + secondary_.TotalMemoryUsed();
+  }
+
+  void TestOnlyUnmap() { primary_.TestOnlyUnmap(); }
+
+ private:
+  PrimaryAllocator primary_;
+  SecondaryAllocator secondary_;
+};
+
 }  // namespace __sanitizer
 
 #endif  // SANITIZER_ALLOCATOR_H
diff --git a/compiler-rt/lib/sanitizer_common/tests/sanitizer_allocator64_test.cc b/compiler-rt/lib/sanitizer_common/tests/sanitizer_allocator64_test.cc
index c0b7b7bd9f5..ba636439f90 100644
--- a/compiler-rt/lib/sanitizer_common/tests/sanitizer_allocator64_test.cc
+++ b/compiler-rt/lib/sanitizer_common/tests/sanitizer_allocator64_test.cc
@@ -71,7 +71,7 @@ TEST(SanitizerCommon, SizeClassAllocator64) {
       // printf("s = %ld\n", size);
       uptr n_iter = std::max((uptr)2, 1000000 / size);
       for (uptr i = 0; i < n_iter; i++) {
-        void *x = a.Allocate(size);
+        void *x = a.Allocate(size, 1);
         allocated.push_back(x);
         CHECK(a.PointerIsMine(x));
         uptr class_id = a.GetSizeClass(x);
@@ -112,7 +112,7 @@ TEST(SanitizerCommon, SizeClassAllocator64MetadataStress) {
   void *allocated[kNumAllocs];
   for (uptr i = 0; i < kNumAllocs; i++) {
     uptr size = (i % 4096) + 1;
-    void *x = a.Allocate(size);
+    void *x = a.Allocate(size, 1);
     allocated[i] = x;
   }
   // Get Metadata kNumAllocs^2 times.
@@ -134,7 +134,7 @@ void FailInAssertionOnOOM() {
   a.Init();
   const uptr size = 1 << 20;
   for (int i = 0; i < 1000000; i++) {
-    a.Allocate(size);
+    a.Allocate(size, 1);
   }
 
   a.TestOnlyUnmap();
@@ -154,7 +154,7 @@ TEST(SanitizerCommon, LargeMmapAllocator) {
   static const uptr size = 1000;
   // Allocate some.
   for (int i = 0; i < kNumAllocs; i++) {
-    allocated[i] = a.Allocate(size);
+    allocated[i] = a.Allocate(size, 1);
   }
   // Deallocate all.
   CHECK_GT(a.TotalMemoryUsed(), size * kNumAllocs);
@@ -168,7 +168,7 @@ TEST(SanitizerCommon, LargeMmapAllocator) {
 
   // Allocate some more, also add metadata.
   for (int i = 0; i < kNumAllocs; i++) {
-    void *x = a.Allocate(size);
+    void *x = a.Allocate(size, 1);
     uptr *meta = reinterpret_cast<uptr*>(a.GetMetaData(x));
     *meta = i;
     allocated[i] = x;
@@ -185,3 +185,42 @@ TEST(SanitizerCommon, LargeMmapAllocator) {
   }
   CHECK_EQ(a.TotalMemoryUsed(), 0);
 }
+
+TEST(SanitizerCommon, CombinedAllocator) {
+  typedef DefaultSizeClassMap SCMap;
+  typedef SizeClassAllocator64<kAllocatorSpace, kAllocatorSize,
+          16, SCMap> PrimaryAllocator;
+  typedef LargeMmapAllocator SecondaryAllocator;
+  typedef CombinedAllocator<PrimaryAllocator, SecondaryAllocator> Allocator;
+
+  Allocator a;
+  a.Init();
+  const uptr kNumAllocs = 100000;
+  const uptr kNumIter = 10;
+  for (uptr iter = 0; iter < kNumIter; iter++) {
+    std::vector<void*> allocated;
+    for (uptr i = 0; i < kNumAllocs; i++) {
+      uptr size = (i % (1 << 14)) + 1;
+      if ((i % 1024) == 0)
+        size = 1 << (10 + (i % 14));
+      void *x = a.Allocate(size, 1);
+      uptr *meta = reinterpret_cast<uptr*>(a.GetMetaData(x));
+      CHECK_EQ(*meta, 0);
+      *meta = size;
+      allocated.push_back(x);
+    }
+
+    random_shuffle(allocated.begin(), allocated.end());
+
+    for (uptr i = 0; i < kNumAllocs; i++) {
+      void *x = allocated[i];
+      uptr *meta = reinterpret_cast<uptr*>(a.GetMetaData(x));
+      CHECK_NE(*meta, 0);
+      CHECK(a.PointerIsMine(x));
+      *meta = 0;
+      a.Deallocate(x);
+    }
+    allocated.clear();
+  }
+  a.TestOnlyUnmap();
+}