1 //===-- sanitizer_flat_map.h ------------------------------------*- C++ -*-===// 2 // 3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. 4 // See https://llvm.org/LICENSE.txt for license information. 5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception 6 // 7 //===----------------------------------------------------------------------===// 8 // 9 // Part of the Sanitizer Allocator. 10 // 11 //===----------------------------------------------------------------------===// 12 13 #ifndef SANITIZER_FLAT_MAP_H 14 #define SANITIZER_FLAT_MAP_H 15 16 #include "sanitizer_atomic.h" 17 #include "sanitizer_common.h" 18 #include "sanitizer_internal_defs.h" 19 #include "sanitizer_local_address_space_view.h" 20 #include "sanitizer_mutex.h" 21 22 namespace __sanitizer { 23 24 // Maps integers in rage [0, kSize) to values. 25 template <typename T, u64 kSize, 26 typename AddressSpaceViewTy = LocalAddressSpaceView> 27 class FlatMap { 28 public: 29 using AddressSpaceView = AddressSpaceViewTy; 30 void Init() { internal_memset(map_, 0, sizeof(map_)); } 31 32 constexpr uptr size() const { return kSize; } 33 34 bool contains(uptr idx) const { 35 CHECK_LT(idx, kSize); 36 return true; 37 } 38 39 T &operator[](uptr idx) { 40 DCHECK_LT(idx, kSize); 41 return map_[idx]; 42 } 43 44 const T &operator[](uptr idx) const { 45 DCHECK_LT(idx, kSize); 46 return map_[idx]; 47 } 48 49 private: 50 T map_[kSize]; 51 }; 52 53 // TwoLevelMap maps integers in range [0, kSize1*kSize2) to values. 54 // It is implemented as a two-dimensional array: array of kSize1 pointers 55 // to kSize2-byte arrays. The secondary arrays are mmaped on demand. 56 // Each value is initially zero and can be set to something else only once. 57 // Setting and getting values from multiple threads is safe w/o extra locking. 58 template <typename T, u64 kSize1, u64 kSize2, 59 typename AddressSpaceViewTy = LocalAddressSpaceView> 60 class TwoLevelMap { 61 static_assert(IsPowerOfTwo(kSize2), "Use a power of two for performance."); 62 63 public: 64 using AddressSpaceView = AddressSpaceViewTy; 65 void Init() { 66 mu_.Init(); 67 internal_memset(map1_, 0, sizeof(map1_)); 68 } 69 70 void TestOnlyUnmap() { 71 for (uptr i = 0; i < kSize1; i++) { 72 T *p = Get(i); 73 if (!p) 74 continue; 75 UnmapOrDie(p, kSize2); 76 } 77 Init(); 78 } 79 80 uptr MemoryUsage() const { 81 uptr res = 0; 82 for (uptr i = 0; i < kSize1; i++) { 83 T *p = Get(i); 84 if (!p) 85 continue; 86 res += MmapSize(); 87 } 88 return res; 89 } 90 91 constexpr uptr size() const { return kSize1 * kSize2; } 92 constexpr uptr size1() const { return kSize1; } 93 constexpr uptr size2() const { return kSize2; } 94 95 bool contains(uptr idx) const { 96 CHECK_LT(idx, kSize1 * kSize2); 97 return Get(idx / kSize2); 98 } 99 100 const T &operator[](uptr idx) const { 101 DCHECK_LT(idx, kSize1 * kSize2); 102 T *map2 = GetOrCreate(idx / kSize2); 103 return *AddressSpaceView::Load(&map2[idx % kSize2]); 104 } 105 106 T &operator[](uptr idx) { 107 DCHECK_LT(idx, kSize1 * kSize2); 108 T *map2 = GetOrCreate(idx / kSize2); 109 return *AddressSpaceView::LoadWritable(&map2[idx % kSize2]); 110 } 111 112 private: 113 constexpr uptr MmapSize() const { 114 return RoundUpTo(kSize2 * sizeof(T), GetPageSizeCached()); 115 } 116 117 T *Get(uptr idx) const { 118 DCHECK_LT(idx, kSize1); 119 return reinterpret_cast<T *>( 120 atomic_load(&map1_[idx], memory_order_acquire)); 121 } 122 123 T *GetOrCreate(uptr idx) const { 124 DCHECK_LT(idx, kSize1); 125 // This code needs to use memory_order_acquire/consume, but we use 126 // memory_order_relaxed for performance reasons (matters for arm64). We 127 // expect memory_order_relaxed to be effectively equivalent to 128 // memory_order_consume in this case for all relevant architectures: all 129 // dependent data is reachable only by dereferencing the resulting pointer. 130 // If relaxed load fails to see stored ptr, the code will fall back to 131 // Create() and reload the value again with locked mutex as a memory 132 // barrier. 133 T *res = reinterpret_cast<T *>(atomic_load_relaxed(&map1_[idx])); 134 if (LIKELY(res)) 135 return res; 136 return Create(idx); 137 } 138 139 NOINLINE T *Create(uptr idx) const { 140 SpinMutexLock l(&mu_); 141 T *res = Get(idx); 142 if (!res) { 143 res = reinterpret_cast<T *>(MmapOrDie(MmapSize(), "TwoLevelMap")); 144 atomic_store(&map1_[idx], reinterpret_cast<uptr>(res), 145 memory_order_release); 146 } 147 return res; 148 } 149 150 mutable StaticSpinMutex mu_; 151 mutable atomic_uintptr_t map1_[kSize1]; 152 }; 153 154 template <u64 kSize, typename AddressSpaceViewTy = LocalAddressSpaceView> 155 using FlatByteMap = FlatMap<u8, kSize, AddressSpaceViewTy>; 156 157 template <u64 kSize1, u64 kSize2, 158 typename AddressSpaceViewTy = LocalAddressSpaceView> 159 using TwoLevelByteMap = TwoLevelMap<u8, kSize1, kSize2, AddressSpaceViewTy>; 160 } // namespace __sanitizer 161 162 #endif 163