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