1 files changed, 23 insertions, 7 deletions

diff --git a/libsanitizer/sanitizer_common/sanitizer_allocator_size_class_map.h b/libsanitizer/sanitizer_common/sanitizer_allocator_size_class_map.h
index 4cda021a947..19af55138d1 100644
--- a/libsanitizer/sanitizer_common/sanitizer_allocator_size_class_map.h
+++ b/libsanitizer/sanitizer_common/sanitizer_allocator_size_class_map.h
@@ -132,8 +132,9 @@ class SizeClassMap {
 
   static const uptr kMaxSize = 1UL << kMaxSizeLog;
   static const uptr kNumClasses =
-      kMidClass + ((kMaxSizeLog - kMidSizeLog) << S) + 1;
+      kMidClass + ((kMaxSizeLog - kMidSizeLog) << S) + 1 + 1;
   static const uptr kLargestClassID = kNumClasses - 2;
+  static const uptr kBatchClassID = kNumClasses - 1;
   COMPILER_CHECK(kNumClasses >= 16 && kNumClasses <= 256);
   static const uptr kNumClassesRounded =
       kNumClasses <= 32  ? 32 :
@@ -141,6 +142,11 @@ class SizeClassMap {
       kNumClasses <= 128 ? 128 : 256;
 
   static uptr Size(uptr class_id) {
+    // Estimate the result for kBatchClassID because this class does not know
+    // the exact size of TransferBatch. It's OK since we are using the actual
+    // sizeof(TransferBatch) where it matters.
+    if (UNLIKELY(class_id == kBatchClassID))
+      return kMaxNumCachedHint * sizeof(uptr);
     if (class_id <= kMidClass)
       return kMinSize * class_id;
     class_id -= kMidClass;
@@ -149,9 +155,10 @@ class SizeClassMap {
   }
 
   static uptr ClassID(uptr size) {
+    if (UNLIKELY(size > kMaxSize))
+      return 0;
     if (size <= kMidSize)
       return (size + kMinSize - 1) >> kMinSizeLog;
-    if (size > kMaxSize) return 0;
     uptr l = MostSignificantSetBitIndex(size);
     uptr hbits = (size >> (l - S)) & M;
     uptr lbits = size & ((1 << (l - S)) - 1);
@@ -160,7 +167,13 @@ class SizeClassMap {
   }
 
   static uptr MaxCachedHint(uptr class_id) {
-    if (class_id == 0) return 0;
+    // Estimate the result for kBatchClassID because this class does not know
+    // the exact size of TransferBatch. We need to cache fewer batches than user
+    // chunks, so this number can be small.
+    if (UNLIKELY(class_id == kBatchClassID))
+      return 16;
+    if (UNLIKELY(class_id == 0))
+      return 0;
     uptr n = (1UL << kMaxBytesCachedLog) / Size(class_id);
     return Max<uptr>(1, Min(kMaxNumCachedHint, n));
   }
@@ -176,6 +189,8 @@ class SizeClassMap {
       uptr p = prev_s ? (d * 100 / prev_s) : 0;
       uptr l = s ? MostSignificantSetBitIndex(s) : 0;
       uptr cached = MaxCachedHint(i) * s;
+      if (i == kBatchClassID)
+        d = p = l = 0;
       Printf("c%02zd => s: %zd diff: +%zd %02zd%% l %zd "
              "cached: %zd %zd; id %zd\n",
              i, Size(i), d, p, l, MaxCachedHint(i), cached, ClassID(s));
@@ -190,12 +205,13 @@ class SizeClassMap {
       // Printf("Validate: c%zd\n", c);
       uptr s = Size(c);
       CHECK_NE(s, 0U);
+      if (c == kBatchClassID)
+        continue;
       CHECK_EQ(ClassID(s), c);
-      if (c != kNumClasses - 1)
+      if (c < kLargestClassID)
         CHECK_EQ(ClassID(s + 1), c + 1);
       CHECK_EQ(ClassID(s - 1), c);
-      if (c)
-        CHECK_GT(Size(c), Size(c-1));
+      CHECK_GT(Size(c), Size(c - 1));
     }
     CHECK_EQ(ClassID(kMaxSize + 1), 0);
 
@@ -205,7 +221,7 @@ class SizeClassMap {
       CHECK_LT(c, kNumClasses);
       CHECK_GE(Size(c), s);
       if (c > 0)
-        CHECK_LT(Size(c-1), s);
+        CHECK_LT(Size(c - 1), s);
     }
   }
 };