use crate::utils::r#ref::RefMut;
use core::hash::{BuildHasher, Hash, Hasher};
use std::sync::Mutex;
use ahash::{AHasher, RandomState};
use hashbrown::HashMap;
use lazy_static::*;
lazy_static! {
static ref B_HASHER: AHasher = RandomState::new().build_hasher();
}
fn shard_amount() -> usize {
(num_cpus::get() * 4).next_power_of_two()
}
fn ncb(shard_amount: usize) -> usize {
shard_amount.trailing_zeros() as usize
}
pub struct DBDashMap<K, V, S = RandomState> {
shift: usize,
pub shards: Box<[Mutex<HashMap<K, V, S>>]>,
pub r_shards: Box<[HashMap<K, V, S>]>,
hasher: S,
}
impl<K, V, S> Default for DBDashMap<K, V, S>
where
K: Eq + Hash,
S: Default + BuildHasher + Clone,
{
fn default() -> Self {
Self::with_hasher(Default::default())
}
}
impl<'a, K: 'a + Eq + Hash, V: 'a> DBDashMap<K, V, RandomState> {
pub fn new() -> Self {
DBDashMap::with_hasher(RandomState::default())
}
pub fn with_capacity(capacity: usize) -> Self {
DBDashMap::with_capacity_and_hasher(capacity, RandomState::default())
}
}
impl<'a, K: 'a + Eq + Hash, V: 'a, S: BuildHasher + Clone + Default> DBDashMap<K, V, S> {
pub fn with_hasher(hasher: S) -> Self {
Self::with_capacity_and_hasher(0, hasher)
}
pub fn with_capacity_and_hasher(mut capacity: usize, hasher: S) -> Self {
let shard_amount = shard_amount();
let shift = ptr_size_bits() - ncb(shard_amount);
if capacity != 0 {
capacity = (capacity + (shard_amount - 1)) & !(shard_amount - 1);
}
let cps = capacity / shard_amount;
let shards = (0..shard_amount)
.map(|_| Mutex::new(HashMap::with_capacity_and_hasher(cps, hasher.clone())))
.collect();
let r_shards = (0..shard_amount)
.map(|_| HashMap::with_capacity_and_hasher(cps, hasher.clone()))
.collect();
Self {
shift,
shards,
r_shards,
hasher,
}
}
pub fn hash_usize<T: Hash>(&self, item: &T) -> usize {
let mut hasher = B_HASHER.clone();
item.hash(&mut hasher);
hasher.finish() as usize
}
pub fn determine_shard(&self, hash: usize) -> usize {
(hash << 7) >> self.shift
}
pub fn insert(&self, key: K, value: V) -> Option<V> {
let hash = self.hash_usize(&key);
let idx = self.determine_shard(hash);
let mut shard = self.shards[idx].lock().expect("error on lock");
shard.remove(&key);
shard.insert(key, value)
}
pub fn remove(&self, key: &K) -> Option<(K, V)> {
let hash = self.hash_usize(&key);
let idx = self.determine_shard(hash);
let mut shard = self.shards[idx].lock().expect("error on lock");
shard.remove_entry(key)
}
pub fn get_read(&'a self, key: &K) -> Option<&'a V> {
let hash = self.hash_usize(&key);
let idx = self.determine_shard(hash);
let shard = &self.r_shards[idx];
shard.get(key)
}
pub fn get_key_value(&self, key: &K) -> Option<(&K, &V)> {
let hash = self.hash_usize(&key);
let idx = self.determine_shard(hash);
let shard = &self.r_shards[idx];
shard.get_key_value(key)
}
pub fn get_write<'b>(&'b self, key: &K) -> Option<RefMut<'b, K, V, S>> {
let hash = self.hash_usize(&key);
let idx = self.determine_shard(hash);
let mut shard = self.shards[idx].lock().expect("error on lock");
match shard.get_mut(key) {
Some(r) => unsafe {
let re = &mut *(r as *mut V);
Some(RefMut::new(shard, re))
},
None => None,
}
}
pub fn lazy_update(&mut self) {
let shard_amount = shard_amount();
for i in 0..shard_amount {
unsafe {
std::ptr::swap(
self.shards[i].get_mut().expect("error on get_mut") as *mut HashMap<K, V, S>,
&mut self.r_shards[i] as *mut HashMap<K, V, S>,
)
}
self.shards[i].get_mut().expect("error on get_mut").clear();
}
}
pub fn merge_r_shards(&mut self) -> HashMap<K, V, S> {
let mut ris = HashMap::with_hasher(self.hasher.clone());
for i in 0..shard_amount() {
ris.extend(self.shards[i].get_mut().expect("error on get_mut").drain());
}
ris
}
pub fn len(&self) -> usize {
self.shards
.iter()
.map(|shard| shard.lock().expect("error on lock").len())
.sum()
}
pub fn r_len(&self) -> usize {
self.r_shards.iter().map(|shard| shard.len()).sum()
}
pub fn keys(&self) -> Vec<&K> {
let mut ris = vec![];
for shard in self.r_shards.iter() {
for key in shard.keys() {
ris.push(key);
}
}
ris
}
pub fn values(&self) -> Vec<&V> {
let mut ris = vec![];
for shard in self.r_shards.iter() {
for value in shard.values() {
ris.push(value);
}
}
ris
}
pub fn clear(&self) {
for shard in self.shards.iter() {
shard.lock().expect("error on lock").clear();
}
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
pub fn is_empty_r(&self) -> bool {
self.r_len() == 0
}
}
impl<'a, K: 'a + Eq + Hash + Clone, V: Clone + 'a> DBDashMap<K, V, RandomState> {
pub fn update(&mut self) {
let n = shard_amount();
for i in 0..n {
for (key, value) in self.shards[i].lock().expect("error on lock").iter() {
self.r_shards[i].remove(&key.clone());
self.r_shards[i].insert(key.clone(), value.clone());
}
}
}
pub fn apply_to_all_values<F>(&self, closure: F)
where
F: Fn(&V) -> V,
{
for shard in self.r_shards.iter() {
for (key, value) in shard.iter() {
self.insert(key.clone(), closure(value));
}
}
}
pub fn apply_to_all_values_write<F>(&self, closure: F)
where
F: Fn(&V) -> V,
{
for shard in self.shards.iter() {
for (_key, value) in shard.lock().expect("error on lock").iter_mut() {
*value = closure(value);
}
}
}
pub fn apply_to_all_values_read_write<F>(&self, closure: F)
where
F: Fn(&V) -> V,
{
for shard_id in 0..self.r_shards.len() {
for (key, value) in self.r_shards[shard_id].iter() {
let mut locs = self.shards[shard_id].lock().expect("error on lock");
if let Some(write_value) = locs.get_mut(key) {
*write_value = closure(write_value);
} else {
locs.insert(key.clone(), closure(value));
}
}
}
}
pub fn apply_to_all_keys<F>(&self, closure: F)
where
F: Fn(&V, &K) -> Option<K>,
{
for shard in self.r_shards.iter() {
for (key, value) in shard.iter() {
let result = closure(value, key);
if let Some(x) = result {
self.remove(key);
self.insert(x, value.clone());
}
}
}
}
pub fn w_keys(&self) -> Vec<K> {
let mut ris: Vec<K> = Vec::with_capacity(self.len());
for shard in self.shards.iter() {
for key in shard.lock().expect("error on lock").keys() {
let k = key.to_owned();
ris.push(k);
}
}
ris
}
}
const fn ptr_size_bits() -> usize {
std::mem::size_of::<usize>() * 8
}