1 files changed, 9 insertions, 6 deletions
diff --git a/lib/sanitizer_common/sanitizer_allocator_primary32.h b/lib/sanitizer_common/sanitizer_allocator_primary32.h
index e85821543..8f8a71d67 100644
--- a/lib/sanitizer_common/sanitizer_allocator_primary32.h
+++ b/lib/sanitizer_common/sanitizer_allocator_primary32.h
@@ -41,6 +41,7 @@ template<class SizeClassAllocator> struct SizeClassAllocator32LocalCache;
 struct SizeClassAllocator32FlagMasks {  //  Bit masks.
   enum {
     kRandomShuffleChunks = 1,
+    kUseSeparateSizeClassForBatch = 2,
   };
 };
 
@@ -55,8 +56,10 @@ class SizeClassAllocator32 {
   typedef typename Params::ByteMap ByteMap;
   typedef typename Params::MapUnmapCallback MapUnmapCallback;
 
-  static const bool kRandomShuffleChunks =
-      Params::kFlags & SizeClassAllocator32FlagMasks::kRandomShuffleChunks;
+  static const bool kRandomShuffleChunks = Params::kFlags &
+      SizeClassAllocator32FlagMasks::kRandomShuffleChunks;
+  static const bool kUseSeparateSizeClassForBatch = Params::kFlags &
+      SizeClassAllocator32FlagMasks::kUseSeparateSizeClassForBatch;
 
   struct TransferBatch {
     static const uptr kMaxNumCached = SizeClassMap::kMaxNumCachedHint - 2;
@@ -94,11 +97,11 @@ class SizeClassAllocator32 {
 
   static const uptr kBatchSize = sizeof(TransferBatch);
   COMPILER_CHECK((kBatchSize & (kBatchSize - 1)) == 0);
-  COMPILER_CHECK(sizeof(TransferBatch) ==
-                 SizeClassMap::kMaxNumCachedHint * sizeof(uptr));
+  COMPILER_CHECK(kBatchSize == SizeClassMap::kMaxNumCachedHint * sizeof(uptr));
 
   static uptr ClassIdToSize(uptr class_id) {
-    return SizeClassMap::Size(class_id);
+    return (class_id == SizeClassMap::kBatchClassID) ?
+        kBatchSize : SizeClassMap::Size(class_id);
   }
 
   typedef SizeClassAllocator32<Params> ThisT;
@@ -161,9 +164,9 @@ class SizeClassAllocator32 {
   NOINLINE void DeallocateBatch(AllocatorStats *stat, uptr class_id,
                                 TransferBatch *b) {
     CHECK_LT(class_id, kNumClasses);
+    CHECK_GT(b->Count(), 0);
     SizeClassInfo *sci = GetSizeClassInfo(class_id);
     SpinMutexLock l(&sci->mutex);
-    CHECK_GT(b->Count(), 0);
     sci->free_list.push_front(b);
   }