[NFC][sanitizer] Refactor Maps into templates

Depends on D111599.

Reviewed By: dvyukov

Differential Revision: https://reviews.llvm.org/D111605

compiler-rt/lib/sanitizer_common/sanitizer_allocator_primary32.h

@@ -199,7 +199,8 @@ class SizeClassAllocator32 {
   }
 
   uptr GetSizeClass(const void *p) const {
-    return possible_regions[ComputeRegionId(reinterpret_cast<uptr>(p))];
+    uptr id = ComputeRegionId(reinterpret_cast<uptr>(p));
+    return possible_regions.contains(id) ? possible_regions[id] : 0;
   }
 
   void *GetBlockBegin(const void *p) {
@@ -253,7 +254,7 @@ class SizeClassAllocator32 {
   // The allocator must be locked when calling this function.
   void ForEachChunk(ForEachChunkCallback callback, void *arg) const {
     for (uptr region = 0; region < kNumPossibleRegions; region++)
-      if (possible_regions[region]) {
+      if (possible_regions.contains(region) && possible_regions[region]) {
         uptr chunk_size = ClassIdToSize(possible_regions[region]);
         uptr max_chunks_in_region = kRegionSize / (chunk_size + kMetadataSize);
         uptr region_beg = region * kRegionSize;
@@ -305,7 +306,7 @@ class SizeClassAllocator32 {
     MapUnmapCallback().OnMap(res, kRegionSize);
     stat->Add(AllocatorStatMapped, kRegionSize);
     CHECK(IsAligned(res, kRegionSize));
-    possible_regions.set(ComputeRegionId(res), static_cast<u8>(class_id));
+    possible_regions[ComputeRegionId(res)] = class_id;
     return res;
   }
 

compiler-rt/lib/sanitizer_common/sanitizer_flat_map.h

@@ -27,8 +27,9 @@ struct NoOpMapUnmapCallback {
   void OnUnmap(uptr p, uptr size) const {}
 };
 
-// Maps integers in rage [0, kSize) to u8 values.
+// Maps integers in rage [0, kSize) to values.
-template <u64 kSize, typename AddressSpaceViewTy = LocalAddressSpaceView>
+template <typename T, u64 kSize,
+          typename AddressSpaceViewTy = LocalAddressSpaceView>
 class FlatMap {
  public:
   using AddressSpaceView = AddressSpaceViewTy;
@@ -36,43 +37,49 @@ class FlatMap {
 
   constexpr uptr size() const { return kSize; }
 
-  void set(uptr idx, u8 val) {
+  bool contains(uptr idx) const {
     CHECK_LT(idx, kSize);
-    CHECK_EQ(0U, map_[idx]);
+    return true;
-    map_[idx] = val;
   }
-  u8 operator[](uptr idx) const {
+
+  T &operator[](uptr idx) {
+    CHECK_LT(idx, kSize);
+    // FIXME: CHECK may be too expensive here.
+    return map_[idx];
+  }
+
+  const T &operator[](uptr idx) const {
     CHECK_LT(idx, kSize);
     // FIXME: CHECK may be too expensive here.
     return map_[idx];
   }
 
  private:
-  u8 map_[kSize];
+  T map_[kSize];
 };
 
-// TwoLevelByteMap maps integers in range [0, kSize1*kSize2) to u8 values.
+// TwoLevelMap maps integers in range [0, kSize1*kSize2) to values.
 // It is implemented as a two-dimensional array: array of kSize1 pointers
 // to kSize2-byte arrays. The secondary arrays are mmaped on demand.
 // Each value is initially zero and can be set to something else only once.
 // Setting and getting values from multiple threads is safe w/o extra locking.
-template <u64 kSize1, u64 kSize2,
+template <typename T, u64 kSize1, u64 kSize2,
           typename AddressSpaceViewTy = LocalAddressSpaceView,
           class MapUnmapCallback = NoOpMapUnmapCallback>
 class TwoLevelMap {
  public:
   using AddressSpaceView = AddressSpaceViewTy;
   void Init() {
-    internal_memset(map1_, 0, sizeof(map1_));
     mu_.Init();
+    internal_memset(map1_, 0, sizeof(map1_));
   }
 
   void TestOnlyUnmap() {
     for (uptr i = 0; i < kSize1; i++) {
-      u8 *p = Get(i);
+      T *p = Get(i);
       if (!p)
         continue;
-      MapUnmapCallback().OnUnmap(reinterpret_cast<uptr>(p), kSize2);
+      MapUnmapCallback().OnUnmap(reinterpret_cast<uptr>(p), kSize2 * sizeof(T));
       UnmapOrDie(p, kSize2);
     }
   }
@@ -81,56 +88,61 @@ class TwoLevelMap {
   constexpr uptr size1() const { return kSize1; }
   constexpr uptr size2() const { return kSize2; }
 
-  void set(uptr idx, u8 val) {
+  bool contains(uptr idx) const {
     CHECK_LT(idx, kSize1 * kSize2);
-    u8 *map2 = GetOrCreate(idx / kSize2);
+    return Get(idx / kSize2);
-    CHECK_EQ(0U, map2[idx % kSize2]);
-    map2[idx % kSize2] = val;
   }
 
-  u8 operator[](uptr idx) const {
+  const T &operator[](uptr idx) const {
     CHECK_LT(idx, kSize1 * kSize2);
-    u8 *map2 = Get(idx / kSize2);
+    T *map2 = GetOrCreate(idx / kSize2);
-    if (!map2)
+    return *AddressSpaceView::Load(&map2[idx % kSize2]);
-      return 0;
+  }
-    auto value_ptr = AddressSpaceView::Load(&map2[idx % kSize2]);
+
-    return *value_ptr;
+  T &operator[](uptr idx) {
+    CHECK_LT(idx, kSize1 * kSize2);
+    T *map2 = GetOrCreate(idx / kSize2);
+    return *AddressSpaceView::LoadWritable(&map2[idx % kSize2]);
   }
 
  private:
-  u8 *Get(uptr idx) const {
+  T *Get(uptr idx) const {
     CHECK_LT(idx, kSize1);
-    return reinterpret_cast<u8 *>(
+    return reinterpret_cast<T *>(
         atomic_load(&map1_[idx], memory_order_acquire));
   }
 
-  u8 *GetOrCreate(uptr idx) {
+  T *GetOrCreate(uptr idx) const {
-    u8 *res = Get(idx);
+    T *res = Get(idx);
+    if (LIKELY(res))
+      return res;
+    return Create(idx);
+  }
+
+  NOINLINE T *Create(uptr idx) const {
+    SpinMutexLock l(&mu_);
+    T *res = Get(idx);
     if (!res) {
-      SpinMutexLock l(&mu_);
+      res = reinterpret_cast<T *>(MmapOrDie(kSize2 * sizeof(T), "TwoLevelMap"));
-      if (!(res = Get(idx))) {
+      MapUnmapCallback().OnMap(reinterpret_cast<uptr>(res), kSize2);
-        res = (u8 *)MmapOrDie(kSize2, "TwoLevelMap");
+      atomic_store(&map1_[idx], reinterpret_cast<uptr>(res),
-        MapUnmapCallback().OnMap(reinterpret_cast<uptr>(res), kSize2);
+                   memory_order_release);
-        atomic_store(&map1_[idx], reinterpret_cast<uptr>(res),
-                     memory_order_release);
-      }
     }
     return res;
   }
 
-  atomic_uintptr_t map1_[kSize1];
+  mutable StaticSpinMutex mu_;
-  StaticSpinMutex mu_;
+  mutable atomic_uintptr_t map1_[kSize1];
 };
 
 template <u64 kSize, typename AddressSpaceViewTy = LocalAddressSpaceView>
-using FlatByteMap = FlatMap<kSize, AddressSpaceViewTy>;
+using FlatByteMap = FlatMap<u8, kSize, AddressSpaceViewTy>;
 
 template <u64 kSize1, u64 kSize2,
           typename AddressSpaceViewTy = LocalAddressSpaceView,
           class MapUnmapCallback = NoOpMapUnmapCallback>
 using TwoLevelByteMap =
-    TwoLevelMap<kSize1, kSize2, AddressSpaceViewTy, MapUnmapCallback>;
+    TwoLevelMap<u8, kSize1, kSize2, AddressSpaceViewTy, MapUnmapCallback>;
-
 }  // namespace __sanitizer
 
 #endif

compiler-rt/lib/sanitizer_common/tests/sanitizer_flat_map_test.cpp

@@ -22,58 +22,83 @@ struct TestMapUnmapCallback1 {
 int TestMapUnmapCallback1::map_count;
 int TestMapUnmapCallback1::unmap_count;
 
-TEST(FlatMapTest, TwoLevelByteMap) {
+struct TestStruct {
+  int data[125] = {};
+  TestStruct(uptr v = 0) { data[11] = v; }
+  bool operator==(const TestStruct &other) const {
+    return 0 == memcmp(data, other.data, sizeof(data));
+  }
+};
+
+template <typename T>
+class FlatMapTest : public ::testing::Test {};
+
+using FlatMapTestTypes = ::testing::Types<u8, u64, TestStruct>;
+TYPED_TEST_SUITE(FlatMapTest, FlatMapTestTypes, );
+
+TYPED_TEST(FlatMapTest, TwoLevelByteMap) {
   const u64 kSize1 = 1 << 6, kSize2 = 1 << 12;
   const u64 n = kSize1 * kSize2;
-  TwoLevelByteMap<kSize1, kSize2> m;
+  TwoLevelMap<TypeParam, kSize1, kSize2> m;
   m.Init();
+
+  m[7] = {10};
+  for (u64 i = 0; i < kSize2; ++i) {
+    EXPECT_TRUE(m.contains(i));
+  }
+  EXPECT_FALSE(m.contains(kSize2));
+
   for (u64 i = 0; i < n; i += 7) {
-    m.set(i, (i % 100) + 1);
+    m[i] = TypeParam((i % 100) + 1);
   }
   for (u64 j = 0; j < n; j++) {
+    EXPECT_TRUE(m.contains(j));
     if (j % 7)
-      EXPECT_EQ(m[j], 0);
+      EXPECT_EQ(m[j], TypeParam());
     else
-      EXPECT_EQ(m[j], (j % 100) + 1);
+      EXPECT_EQ(m[j], TypeParam((j % 100) + 1));
   }
 
   m.TestOnlyUnmap();
 }
 
-template <typename AddressSpaceView>
+template <typename TypeParam, typename AddressSpaceView>
-using TestByteMapASVT =
+using TestMapASVT = TwoLevelMap<TypeParam, 1 << 8, 1 << 7, AddressSpaceView,
-    TwoLevelByteMap<1 << 12, 1 << 13, AddressSpaceView, TestMapUnmapCallback1>;
+                                TestMapUnmapCallback1>;
-using TestByteMap = TestByteMapASVT<LocalAddressSpaceView>;
+template <typename TypeParam>
+using TestMap = TestMapASVT<TypeParam, LocalAddressSpaceView>;
 
-struct TestByteMapParam {
+template <typename TypeParam>
-  TestByteMap *m;
+struct TestMapParam {
+  TestMap<TypeParam> *m;
   size_t shard;
   size_t num_shards;
 };
 
-static void *TwoLevelByteMapUserThread(void *param) {
+template <typename TypeParam>
-  TestByteMapParam *p = (TestByteMapParam *)param;
+static void *TwoLevelMapUserThread(void *param) {
+  TestMapParam<TypeParam> *p = (TestMapParam<TypeParam> *)param;
   for (size_t i = p->shard; i < p->m->size(); i += p->num_shards) {
-    size_t val = (i % 100) + 1;
+    TypeParam val = (i % 100) + 1;
-    p->m->set(i, val);
+    (*p->m)[i] = val;
     EXPECT_EQ((*p->m)[i], val);
   }
   return 0;
 }
 
-TEST(FlatMapTest, ThreadedTwoLevelByteMap) {
+TYPED_TEST(FlatMapTest, ThreadedTwoLevelByteMap) {
-  TestByteMap m;
+  TestMap<TypeParam> m;
   m.Init();
   TestMapUnmapCallback1::map_count = 0;
   TestMapUnmapCallback1::unmap_count = 0;
   static const int kNumThreads = 4;
   pthread_t t[kNumThreads];
-  TestByteMapParam p[kNumThreads];
+  TestMapParam<TypeParam> p[kNumThreads];
   for (int i = 0; i < kNumThreads; i++) {
     p[i].m = &m;
     p[i].shard = i;
     p[i].num_shards = kNumThreads;
-    PTHREAD_CREATE(&t[i], 0, TwoLevelByteMapUserThread, &p[i]);
+    PTHREAD_CREATE(&t[i], 0, TwoLevelMapUserThread<TypeParam>, &p[i]);
   }
   for (int i = 0; i < kNumThreads; i++) {
     PTHREAD_JOIN(t[i], 0);