gcc/libsanitizer/sanitizer_common/sanitizer_flat_map.h

174 lines
5.1 KiB
C++

//===-- sanitizer_flat_map.h ------------------------------------*- C++ -*-===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
//
// Part of the Sanitizer Allocator.
//
//===----------------------------------------------------------------------===//
#ifndef SANITIZER_FLAT_MAP_H
#define SANITIZER_FLAT_MAP_H
#include "sanitizer_atomic.h"
#include "sanitizer_common.h"
#include "sanitizer_internal_defs.h"
#include "sanitizer_local_address_space_view.h"
#include "sanitizer_mutex.h"
namespace __sanitizer {
// Call these callbacks on mmap/munmap.
struct NoOpMapUnmapCallback {
void OnMap(uptr p, uptr size) const {}
void OnUnmap(uptr p, uptr size) const {}
};
// Maps integers in rage [0, kSize) to values.
template <typename T, u64 kSize,
typename AddressSpaceViewTy = LocalAddressSpaceView>
class FlatMap {
public:
using AddressSpaceView = AddressSpaceViewTy;
void Init() { internal_memset(map_, 0, sizeof(map_)); }
constexpr uptr size() const { return kSize; }
bool contains(uptr idx) const {
CHECK_LT(idx, kSize);
return true;
}
T &operator[](uptr idx) {
DCHECK_LT(idx, kSize);
return map_[idx];
}
const T &operator[](uptr idx) const {
DCHECK_LT(idx, kSize);
return map_[idx];
}
private:
T map_[kSize];
};
// TwoLevelMap maps integers in range [0, kSize1*kSize2) to values.
// It is implemented as a two-dimensional array: array of kSize1 pointers
// to kSize2-byte arrays. The secondary arrays are mmaped on demand.
// Each value is initially zero and can be set to something else only once.
// Setting and getting values from multiple threads is safe w/o extra locking.
template <typename T, u64 kSize1, u64 kSize2,
typename AddressSpaceViewTy = LocalAddressSpaceView,
class MapUnmapCallback = NoOpMapUnmapCallback>
class TwoLevelMap {
static_assert(IsPowerOfTwo(kSize2), "Use a power of two for performance.");
public:
using AddressSpaceView = AddressSpaceViewTy;
void Init() {
mu_.Init();
internal_memset(map1_, 0, sizeof(map1_));
}
void TestOnlyUnmap() {
for (uptr i = 0; i < kSize1; i++) {
T *p = Get(i);
if (!p)
continue;
MapUnmapCallback().OnUnmap(reinterpret_cast<uptr>(p), MmapSize());
UnmapOrDie(p, kSize2);
}
Init();
}
uptr MemoryUsage() const {
uptr res = 0;
for (uptr i = 0; i < kSize1; i++) {
T *p = Get(i);
if (!p)
continue;
res += MmapSize();
}
return res;
}
constexpr uptr size() const { return kSize1 * kSize2; }
constexpr uptr size1() const { return kSize1; }
constexpr uptr size2() const { return kSize2; }
bool contains(uptr idx) const {
CHECK_LT(idx, kSize1 * kSize2);
return Get(idx / kSize2);
}
const T &operator[](uptr idx) const {
DCHECK_LT(idx, kSize1 * kSize2);
T *map2 = GetOrCreate(idx / kSize2);
return *AddressSpaceView::Load(&map2[idx % kSize2]);
}
T &operator[](uptr idx) {
DCHECK_LT(idx, kSize1 * kSize2);
T *map2 = GetOrCreate(idx / kSize2);
return *AddressSpaceView::LoadWritable(&map2[idx % kSize2]);
}
private:
constexpr uptr MmapSize() const {
return RoundUpTo(kSize2 * sizeof(T), GetPageSizeCached());
}
T *Get(uptr idx) const {
DCHECK_LT(idx, kSize1);
return reinterpret_cast<T *>(
atomic_load(&map1_[idx], memory_order_acquire));
}
T *GetOrCreate(uptr idx) const {
DCHECK_LT(idx, kSize1);
// This code needs to use memory_order_acquire/consume, but we use
// memory_order_relaxed for performance reasons (matters for arm64). We
// expect memory_order_relaxed to be effectively equivalent to
// memory_order_consume in this case for all relevant architectures: all
// dependent data is reachable only by dereferencing the resulting pointer.
// If relaxed load fails to see stored ptr, the code will fall back to
// Create() and reload the value again with locked mutex as a memory
// barrier.
T *res = reinterpret_cast<T *>(atomic_load_relaxed(&map1_[idx]));
if (LIKELY(res))
return res;
return Create(idx);
}
NOINLINE T *Create(uptr idx) const {
SpinMutexLock l(&mu_);
T *res = Get(idx);
if (!res) {
res = reinterpret_cast<T *>(MmapOrDie(MmapSize(), "TwoLevelMap"));
MapUnmapCallback().OnMap(reinterpret_cast<uptr>(res), kSize2);
atomic_store(&map1_[idx], reinterpret_cast<uptr>(res),
memory_order_release);
}
return res;
}
mutable StaticSpinMutex mu_;
mutable atomic_uintptr_t map1_[kSize1];
};
template <u64 kSize, typename AddressSpaceViewTy = LocalAddressSpaceView>
using FlatByteMap = FlatMap<u8, kSize, AddressSpaceViewTy>;
template <u64 kSize1, u64 kSize2,
typename AddressSpaceViewTy = LocalAddressSpaceView,
class MapUnmapCallback = NoOpMapUnmapCallback>
using TwoLevelByteMap =
TwoLevelMap<u8, kSize1, kSize2, AddressSpaceViewTy, MapUnmapCallback>;
} // namespace __sanitizer
#endif