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.
17template <u64 kSize, typename AddressSpaceViewTy = LocalAddressSpaceView>
18class FlatByteMap {
19 public:
20  using AddressSpaceView = AddressSpaceViewTy;
21  void Init() {
22    internal_memset(map_, 0, sizeof(map_));
23  }
24
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.
44template <u64 kSize1, u64 kSize2,
45          typename AddressSpaceViewTy = LocalAddressSpaceView,
46          class MapUnmapCallback = NoOpMapUnmapCallback>
47class TwoLevelByteMap {
48 public:
49  using AddressSpaceView = AddressSpaceViewTy;
50  void Init() {
51    internal_memset(map1_, 0, sizeof(map1_));
52    mu_.Init();
53  }
54
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
64  uptr size() const { return kSize1 * kSize2; }
65  uptr size1() const { return kSize1; }
66  uptr size2() const { return kSize2; }
67
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:
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
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