142 lines
3.7 KiB
C++
142 lines
3.7 KiB
C++
//
|
|
// Copyright Aliaksei Levin (levlam@telegram.org), Arseny Smirnov (arseny30@gmail.com) 2014-2022
|
|
//
|
|
// Distributed under the Boost Software License, Version 1.0. (See accompanying
|
|
// file LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)
|
|
//
|
|
#pragma once
|
|
|
|
#include "td/utils/common.h"
|
|
#include "td/utils/FlatHashSet.h"
|
|
#include "td/utils/HashTableUtils.h"
|
|
|
|
#include <functional>
|
|
|
|
namespace td {
|
|
|
|
template <class KeyT, class HashT = Hash<KeyT>, class EqT = std::equal_to<KeyT>>
|
|
class WaitFreeHashSet {
|
|
static constexpr size_t MAX_STORAGE_COUNT = 1 << 8;
|
|
static_assert((MAX_STORAGE_COUNT & (MAX_STORAGE_COUNT - 1)) == 0, "");
|
|
static constexpr uint32 DEFAULT_STORAGE_SIZE = 1 << 12;
|
|
|
|
FlatHashSet<KeyT, HashT, EqT> default_set_;
|
|
struct WaitFreeStorage {
|
|
WaitFreeHashSet sets_[MAX_STORAGE_COUNT];
|
|
};
|
|
unique_ptr<WaitFreeStorage> wait_free_storage_;
|
|
uint32 hash_mult_ = 1;
|
|
uint32 max_storage_size_ = DEFAULT_STORAGE_SIZE;
|
|
|
|
uint32 get_wait_free_index(const KeyT &key) const {
|
|
return randomize_hash(HashT()(key) * hash_mult_) & (MAX_STORAGE_COUNT - 1);
|
|
}
|
|
|
|
WaitFreeHashSet &get_wait_free_storage(const KeyT &key) {
|
|
return wait_free_storage_->sets_[get_wait_free_index(key)];
|
|
}
|
|
|
|
const WaitFreeHashSet &get_wait_free_storage(const KeyT &key) const {
|
|
return wait_free_storage_->sets_[get_wait_free_index(key)];
|
|
}
|
|
|
|
void split_storage() {
|
|
CHECK(wait_free_storage_ == nullptr);
|
|
wait_free_storage_ = make_unique<WaitFreeStorage>();
|
|
uint32 next_hash_mult = hash_mult_ * 1000000007;
|
|
for (uint32 i = 0; i < MAX_STORAGE_COUNT; i++) {
|
|
auto &set = wait_free_storage_->sets_[i];
|
|
set.hash_mult_ = next_hash_mult;
|
|
set.max_storage_size_ = DEFAULT_STORAGE_SIZE + i * next_hash_mult % DEFAULT_STORAGE_SIZE;
|
|
}
|
|
for (auto &it : default_set_) {
|
|
get_wait_free_storage(it).insert(it);
|
|
}
|
|
default_set_.clear();
|
|
}
|
|
|
|
public:
|
|
void insert(const KeyT &key) {
|
|
if (wait_free_storage_ != nullptr) {
|
|
return get_wait_free_storage(key).insert(key);
|
|
}
|
|
|
|
default_set_.insert(key);
|
|
if (default_set_.size() == max_storage_size_) {
|
|
split_storage();
|
|
}
|
|
}
|
|
|
|
size_t count(const KeyT &key) const {
|
|
if (wait_free_storage_ != nullptr) {
|
|
return get_wait_free_storage(key).count(key);
|
|
}
|
|
|
|
return default_set_.count(key);
|
|
}
|
|
|
|
size_t erase(const KeyT &key) {
|
|
if (wait_free_storage_ != nullptr) {
|
|
return get_wait_free_storage(key).erase(key);
|
|
}
|
|
|
|
return default_set_.erase(key);
|
|
}
|
|
|
|
void foreach(const std::function<void(const KeyT &key)> &callback) const {
|
|
if (wait_free_storage_ == nullptr) {
|
|
for (auto &it : default_set_) {
|
|
callback(it);
|
|
}
|
|
return;
|
|
}
|
|
|
|
for (auto &it : wait_free_storage_->sets_) {
|
|
it.foreach(callback);
|
|
}
|
|
}
|
|
|
|
KeyT get_random() const {
|
|
if (wait_free_storage_ != nullptr) {
|
|
for (size_t i = 0; i < MAX_STORAGE_COUNT; i++) {
|
|
if (!wait_free_storage_->sets_[i].empty()) {
|
|
return wait_free_storage_->sets_[i].get_random();
|
|
}
|
|
}
|
|
// no need to explicitly return KeyT()
|
|
}
|
|
|
|
if (default_set_.empty()) {
|
|
return KeyT();
|
|
}
|
|
return *default_set_.begin();
|
|
}
|
|
|
|
size_t calc_size() const {
|
|
if (wait_free_storage_ == nullptr) {
|
|
return default_set_.size();
|
|
}
|
|
|
|
size_t result = 0;
|
|
for (size_t i = 0; i < MAX_STORAGE_COUNT; i++) {
|
|
result += wait_free_storage_->sets_[i].calc_size();
|
|
}
|
|
return result;
|
|
}
|
|
|
|
bool empty() const {
|
|
if (wait_free_storage_ == nullptr) {
|
|
return default_set_.empty();
|
|
}
|
|
|
|
for (size_t i = 0; i < MAX_STORAGE_COUNT; i++) {
|
|
if (!wait_free_storage_->sets_[i].empty()) {
|
|
return false;
|
|
}
|
|
}
|
|
return true;
|
|
}
|
|
};
|
|
|
|
} // namespace td
|