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