1 //===-- sanitizer_allocator_bytemap.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 #ifndef SANITIZER_ALLOCATOR_H
13 #error This file must be included inside sanitizer_allocator.h
14 #endif
15 
16 // Maps integers in rage [0, kSize) to u8 values.
17 template <u64 kSize, typename AddressSpaceViewTy = LocalAddressSpaceView>
18 class FlatByteMap {
19  public:
20   using AddressSpaceView = AddressSpaceViewTy;
Init()21   void Init() {
22     internal_memset(map_, 0, sizeof(map_));
23   }
24 
set(uptr idx,u8 val)25   void set(uptr idx, u8 val) {
26     CHECK_LT(idx, kSize);
27     CHECK_EQ(0U, map_[idx]);
28     map_[idx] = val;
29   }
30   u8 operator[] (uptr idx) {
31     CHECK_LT(idx, kSize);
32     // FIXME: CHECK may be too expensive here.
33     return map_[idx];
34   }
35  private:
36   u8 map_[kSize];
37 };
38 
39 // TwoLevelByteMap maps integers in range [0, kSize1*kSize2) to u8 values.
40 // It is implemented as a two-dimensional array: array of kSize1 pointers
41 // to kSize2-byte arrays. The secondary arrays are mmaped on demand.
42 // Each value is initially zero and can be set to something else only once.
43 // Setting and getting values from multiple threads is safe w/o extra locking.
44 template <u64 kSize1, u64 kSize2,
45           typename AddressSpaceViewTy = LocalAddressSpaceView,
46           class MapUnmapCallback = NoOpMapUnmapCallback>
47 class TwoLevelByteMap {
48  public:
49   using AddressSpaceView = AddressSpaceViewTy;
Init()50   void Init() {
51     internal_memset(map1_, 0, sizeof(map1_));
52     mu_.Init();
53   }
54 
TestOnlyUnmap()55   void TestOnlyUnmap() {
56     for (uptr i = 0; i < kSize1; i++) {
57       u8 *p = Get(i);
58       if (!p) continue;
59       MapUnmapCallback().OnUnmap(reinterpret_cast<uptr>(p), kSize2);
60       UnmapOrDie(p, kSize2);
61     }
62   }
63 
size()64   uptr size() const { return kSize1 * kSize2; }
size1()65   uptr size1() const { return kSize1; }
size2()66   uptr size2() const { return kSize2; }
67 
set(uptr idx,u8 val)68   void set(uptr idx, u8 val) {
69     CHECK_LT(idx, kSize1 * kSize2);
70     u8 *map2 = GetOrCreate(idx / kSize2);
71     CHECK_EQ(0U, map2[idx % kSize2]);
72     map2[idx % kSize2] = val;
73   }
74 
75   u8 operator[] (uptr idx) const {
76     CHECK_LT(idx, kSize1 * kSize2);
77     u8 *map2 = Get(idx / kSize2);
78     if (!map2) return 0;
79     auto value_ptr = AddressSpaceView::Load(&map2[idx % kSize2]);
80     return *value_ptr;
81   }
82 
83  private:
Get(uptr idx)84   u8 *Get(uptr idx) const {
85     CHECK_LT(idx, kSize1);
86     return reinterpret_cast<u8 *>(
87         atomic_load(&map1_[idx], memory_order_acquire));
88   }
89 
GetOrCreate(uptr idx)90   u8 *GetOrCreate(uptr idx) {
91     u8 *res = Get(idx);
92     if (!res) {
93       SpinMutexLock l(&mu_);
94       if (!(res = Get(idx))) {
95         res = (u8*)MmapOrDie(kSize2, "TwoLevelByteMap");
96         MapUnmapCallback().OnMap(reinterpret_cast<uptr>(res), kSize2);
97         atomic_store(&map1_[idx], reinterpret_cast<uptr>(res),
98                      memory_order_release);
99       }
100     }
101     return res;
102   }
103 
104   atomic_uintptr_t map1_[kSize1];
105   StaticSpinMutex mu_;
106 };
107 
108