[sanitizer] Minor 32-bit primary improvements

[cryptoad]

Summary

For the 32-bit TransferBatch:

• SetFromArray callers have bounds count, so relax the CHECK to DCHECK;
• same for Add;
• mark CopyToArray as const;

For the 32-bit Primary:

• {Dea,A}llocateBatch are only called from places that check class_id, relax the CHECK to DCHECK;
• same for AllocateRegion;
• remove GetRegionBeginBySizeClass that is not used;
• use a local variable for the random shuffle state, so that the compiler can use a register instead of reading and writing to the SizeClassInfo at every iteration;

For the 32-bit local cache:

• pass the count to drain instead of doing a Min everytime which is at times superfluous.

[cryptoad]

And I forgot:

• the SizeClassInfo mutex should be a StaticSpinMutex since it's zero initialized and doesn't need to be "constructed";

[alekseyshl]

lib/sanitizer_common/sanitizer_allocator_local_cache.h

183 ↗

(On Diff #145981)

Drain(c, allocator, class_id, c->count / 2); is the same, but reads a bit easier, more logical ("drain half of the elements").

192 ↗

(On Diff #145981)

Why the loop, why don't we drain the class in one call?

Drain(c, allocator, i, c->count);

We used to have to have a loop because of the Drain API, but now we can specify the number of elements to drain.

lib/sanitizer_common/sanitizer_allocator_primary32.h

306

Add a comment why you need a local variable here, otherwise someone will "optimize" it back later or, better yet, add a local variable in RandomShuffle and let the compiler do its job generating the best code possible.

[cryptoad]

Addressing a couple of comments from Aleksey.

[cryptoad]

lib/sanitizer_common/sanitizer_allocator_local_cache.h

192 ↗

(On Diff #145981)

Since the TransferBatch that we are using can only hold max_count / 2, we can't directly drain count.
I could change it to something like if count > max_count / 2 then we can drain max_count / 2, then drain count elements.

[alekseyshl]

lib/sanitizer_common/sanitizer_allocator_local_cache.h

192 ↗

(On Diff #145981)

Ah, I see.

This restriction was not obvious before and it is even more obfuscated now. Was Min() call in Drain so taxing that it is worth complicating the already not so easy to follow logic of the code? If it is not, I'd rather not change the code. If it is though, well, do it, but add a DCHECK to Drain to document the restriction of this API: DCHECK_LE(count, c->max_count / 2); and ignore my previous comment, Drain(c, allocator, class_id, c->max_count / 2);

[cryptoad]

lib/sanitizer_common/sanitizer_allocator_local_cache.h

192 ↗

(On Diff #145981)

Yeah you are right it's not worth it. I'll revert this part to its original state.

[cryptoad]

The changes to the local cache were not worth doing. Reverting those.

lib/sanitizer_common/sanitizer_allocator.h

INLINE u32 Rand(u32 *state) {  // ANSI C linear congruential PRNG.
   return (*state = *state * 1103515245 + 12345) >> 16;
 }
 
 INLINE u32 RandN(u32 *state, u32 n) { return Rand(state) % n; }  // [0, n)
 
 template<typename T>
 INLINE void RandomShuffle(T *a, u32 n, u32 *rand_state) {
   if (n <= 1) return;
+  u32 state = *rand_state;
   for (u32 i = n - 1; i > 0; i--)
-    Swap(a[i], a[RandN(rand_state, i + 1)]);
+    Swap(a[i], a[RandN(&state, i + 1)]);
+  *rand_state = state;
 }
 
 #include "sanitizer_allocator_size_class_map.h"
 #include "sanitizer_allocator_stats.h"
 #include "sanitizer_allocator_primary64.h"
 #include "sanitizer_allocator_bytemap.h"
 #include "sanitizer_allocator_primary32.h"
 #include "sanitizer_allocator_local_cache.h"
 #include "sanitizer_allocator_secondary.h"
 #include "sanitizer_allocator_combined.h"
 
 } // namespace __sanitizer
 
 #endif // SANITIZER_ALLOCATOR_H

lib/sanitizer_common/sanitizer_allocator_primary32.h

public:
   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;
     void SetFromArray(void *batch[], uptr count) {
+      DCHECK_LE(count, kMaxNumCached);
       count_ = count;
-      CHECK_LE(count_, kMaxNumCached);
       for (uptr i = 0; i < count; i++)
         batch_[i] = batch[i];
     }
     uptr Count() const { return count_; }
     void Clear() { count_ = 0; }
     void Add(void *ptr) {
       batch_[count_++] = ptr;
-      CHECK_LE(count_, kMaxNumCached);
+      DCHECK_LE(count_, kMaxNumCached);
     }
-    void CopyToArray(void *to_batch[]) {
+    void CopyToArray(void *to_batch[]) const {
       for (uptr i = 0, n = Count(); i < n; i++)
         to_batch[i] = batch_[i];
     }
 
     // How much memory do we need for a batch containing n elements.
     static uptr AllocationSizeRequiredForNElements(uptr n) {
       return sizeof(uptr) * 2 + sizeof(void *) * n;
     }
void *GetMetaData(const void *p) {
     u32 offset = mem - beg;
     uptr n = offset / (u32)size;  // 32-bit division
     uptr meta = (beg + kRegionSize) - (n + 1) * kMetadataSize;
     return reinterpret_cast<void*>(meta);
   }
 
   NOINLINE TransferBatch *AllocateBatch(AllocatorStats *stat, AllocatorCache *c,
                                         uptr class_id) {
-    CHECK_LT(class_id, kNumClasses);
+    DCHECK_LT(class_id, kNumClasses);
     SizeClassInfo *sci = GetSizeClassInfo(class_id);
     SpinMutexLock l(&sci->mutex);
     if (sci->free_list.empty()) {
       if (UNLIKELY(!PopulateFreeList(stat, c, sci, class_id)))
         return nullptr;
       DCHECK(!sci->free_list.empty());
     }
     TransferBatch *b = sci->free_list.front();
     sci->free_list.pop_front();
     return b;
   }
 
   NOINLINE void DeallocateBatch(AllocatorStats *stat, uptr class_id,
                                 TransferBatch *b) {
-    CHECK_LT(class_id, kNumClasses);
+    DCHECK_LT(class_id, kNumClasses);
     CHECK_GT(b->Count(), 0);
     SizeClassInfo *sci = GetSizeClassInfo(class_id);
     SpinMutexLock l(&sci->mutex);
     sci->free_list.push_front(b);
   }
 
-  uptr GetRegionBeginBySizeClass(uptr class_id) { return 0; }
-
   bool PointerIsMine(const void *p) {
     uptr mem = reinterpret_cast<uptr>(p);
     if (mem < kSpaceBeg || mem >= kSpaceBeg + kSpaceSize)
       return false;
     return GetSizeClass(p) != 0;
   }
 
   uptr GetSizeClass(const void *p) {
for (uptr region = 0; region < kNumPossibleRegions; region++)
              chunk < region_beg + max_chunks_in_region * chunk_size;
              chunk += chunk_size) {
           // Too slow: CHECK_EQ((void *)chunk, GetBlockBegin((void *)chunk));
           callback(chunk, arg);
         }
       }
   }
 
-  void PrintStats() {
-  }
+  void PrintStats() {}
 
-  static uptr AdditionalSize() {
-    return 0;
-  }
+  static uptr AdditionalSize() { return 0; }
 
   typedef SizeClassMap SizeClassMapT;
   static const uptr kNumClasses = SizeClassMap::kNumClasses;
 
  private:
   static const uptr kRegionSize = 1 << kRegionSizeLog;
   static const uptr kNumPossibleRegions = kSpaceSize / kRegionSize;
 
   struct ALIGNED(SANITIZER_CACHE_LINE_SIZE) SizeClassInfo {
-    SpinMutex mutex;
+    StaticSpinMutex mutex;
     IntrusiveList<TransferBatch> free_list;
     u32 rand_state;
   };
   COMPILER_CHECK(sizeof(SizeClassInfo) % kCacheLineSize == 0);
 
   uptr ComputeRegionId(uptr mem) {
     const uptr res = mem >> kRegionSizeLog;
     CHECK_LT(res, kNumPossibleRegions);
     return res;
   }
 
   uptr ComputeRegionBeg(uptr mem) {
     return mem & ~(kRegionSize - 1);
   }
 
   uptr AllocateRegion(AllocatorStats *stat, uptr class_id) {
-    CHECK_LT(class_id, kNumClasses);
-    uptr res = reinterpret_cast<uptr>(MmapAlignedOrDieOnFatalError(
+    DCHECK_LT(class_id, kNumClasses);
+    const uptr res = reinterpret_cast<uptr>(MmapAlignedOrDieOnFatalError(
         kRegionSize, kRegionSize, PrimaryAllocatorName));
     if (UNLIKELY(!res))
       return 0;
     MapUnmapCallback().OnMap(res, kRegionSize);
     stat->Add(AllocatorStatMapped, kRegionSize);
     CHECK(IsAligned(res, kRegionSize));
     possible_regions.set(ComputeRegionId(res), static_cast<u8>(class_id));
     return res;
   }
 
   SizeClassInfo *GetSizeClassInfo(uptr class_id) {
     DCHECK_LT(class_id, kNumClasses);
     return &size_class_info_array[class_id];
   }
 
   bool PopulateBatches(AllocatorCache *c, SizeClassInfo *sci, uptr class_id,
                        TransferBatch **current_batch, uptr max_count,
                        uptr *pointers_array, uptr count) {
     // If using a separate class for batches, we do not need to shuffle it.
     if (kRandomShuffleChunks && (!kUseSeparateSizeClassForBatch ||
         class_id != SizeClassMap::kBatchClassID))
       RandomShuffle(pointers_array, count, &sci->rand_state);
     TransferBatch *b = *current_batch;
     for (uptr i = 0; i < count; i++) {
       if (!b) {
         b = c->CreateBatch(class_id, this, (TransferBatch*)pointers_array[i]);
         if (UNLIKELY(!b))
           return false;
         b->Clear();
       }
       b->Add((void*)pointers_array[i]);