use crate::map::{
bucket::{self, BucketArray},
bucket_array_ref::BucketArrayRef,
DefaultHashBuilder,
};
use std::{
borrow::Borrow,
hash::{BuildHasher, Hash},
ptr,
sync::atomic::{self, AtomicUsize, Ordering},
};
use crossbeam_epoch::Atomic;
pub struct HashMap<K, V, S = DefaultHashBuilder> {
segments: Box<[Segment<K, V>]>,
build_hasher: S,
len: AtomicUsize,
segment_shift: u32,
}
#[cfg(feature = "num-cpus")]
impl<K, V> HashMap<K, V, DefaultHashBuilder> {
pub fn new() -> Self {
Self::with_num_segments_capacity_and_hasher(
default_num_segments(),
0,
DefaultHashBuilder::default(),
)
}
pub fn with_capacity(capacity: usize) -> Self {
Self::with_num_segments_capacity_and_hasher(
default_num_segments(),
capacity,
DefaultHashBuilder::default(),
)
}
}
#[cfg(feature = "num-cpus")]
impl<K, V, S: BuildHasher> HashMap<K, V, S> {
pub fn with_hasher(build_hasher: S) -> Self {
Self::with_num_segments_capacity_and_hasher(default_num_segments(), 0, build_hasher)
}
pub fn with_capacity_and_hasher(capacity: usize, build_hasher: S) -> Self {
Self::with_num_segments_capacity_and_hasher(default_num_segments(), capacity, build_hasher)
}
}
impl<K, V> HashMap<K, V, DefaultHashBuilder> {
pub fn with_num_segments(num_segments: usize) -> Self {
Self::with_num_segments_capacity_and_hasher(num_segments, 0, DefaultHashBuilder::default())
}
pub fn with_num_segments_and_capacity(num_segments: usize, capacity: usize) -> Self {
Self::with_num_segments_capacity_and_hasher(
num_segments,
capacity,
DefaultHashBuilder::default(),
)
}
}
impl<K, V, S> HashMap<K, V, S> {
pub fn with_num_segments_and_hasher(num_segments: usize, build_hasher: S) -> Self {
Self::with_num_segments_capacity_and_hasher(num_segments, 0, build_hasher)
}
pub fn with_num_segments_capacity_and_hasher(
num_segments: usize,
capacity: usize,
build_hasher: S,
) -> Self {
assert!(num_segments > 0);
let actual_num_segments = num_segments.next_power_of_two();
let segment_shift = 64 - actual_num_segments.trailing_zeros();
let mut segments = Vec::with_capacity(actual_num_segments);
if capacity == 0 {
unsafe {
ptr::write_bytes(segments.as_mut_ptr(), 0, actual_num_segments);
segments.set_len(actual_num_segments);
}
} else {
let actual_capacity = (capacity * 2).next_power_of_two();
for _ in 0..actual_num_segments {
segments.push(Segment {
bucket_array: Atomic::new(BucketArray::with_length(0, actual_capacity)),
len: AtomicUsize::new(0),
});
}
}
let segments = segments.into_boxed_slice();
Self {
segments,
build_hasher,
len: AtomicUsize::new(0),
segment_shift,
}
}
pub fn len(&self) -> usize {
self.len.load(Ordering::Relaxed)
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
pub fn capacity(&self) -> usize {
let guard = &crossbeam_epoch::pin();
self.segments
.iter()
.map(|s| s.bucket_array.load_consume(guard))
.map(|p| unsafe { p.as_ref() })
.map(|a| a.map(BucketArray::capacity).unwrap_or(0))
.min()
.unwrap()
}
pub fn segment_capacity(&self, index: usize) -> usize {
assert!(index < self.segments.len());
let guard = &crossbeam_epoch::pin();
unsafe {
self.segments[index]
.bucket_array
.load_consume(guard)
.as_ref()
}
.map(BucketArray::capacity)
.unwrap_or(0)
}
pub fn num_segments(&self) -> usize {
self.segments.len()
}
}
impl<K, V, S: BuildHasher> HashMap<K, V, S> {
pub fn segment_index<Q: Hash>(&self, key: &Q) -> usize
where
K: Borrow<Q>,
{
let hash = bucket::hash(&self.build_hasher, key);
self.segment_index_from_hash(hash)
}
}
impl<K: Hash + Eq, V, S: BuildHasher> HashMap<K, V, S> {
#[inline]
pub fn get<Q: Hash + Eq + ?Sized>(&self, key: &Q) -> Option<V>
where
K: Borrow<Q>,
V: Clone,
{
self.get_key_value_and(key, |_, v| v.clone())
}
#[inline]
pub fn get_key_value<Q: Hash + Eq + ?Sized>(&self, key: &Q) -> Option<(K, V)>
where
K: Borrow<Q> + Clone,
V: Clone,
{
self.get_key_value_and(key, |k, v| (k.clone(), v.clone()))
}
#[inline]
pub fn get_and<Q: Hash + Eq + ?Sized, F: FnOnce(&V) -> T, T>(
&self,
key: &Q,
with_value: F,
) -> Option<T>
where
K: Borrow<Q>,
{
self.get_key_value_and(key, move |_, v| with_value(v))
}
#[inline]
pub fn get_key_value_and<Q: Hash + Eq + ?Sized, F: FnOnce(&K, &V) -> T, T>(
&self,
key: &Q,
with_entry: F,
) -> Option<T>
where
K: Borrow<Q>,
{
let hash = bucket::hash(&self.build_hasher, &key);
self.bucket_array_ref(hash)
.get_key_value_and(key, hash, with_entry)
}
#[inline]
pub fn insert(&self, key: K, value: V) -> Option<V>
where
V: Clone,
{
self.insert_entry_and(key, value, |_, v| v.clone())
}
#[inline]
pub fn insert_entry(&self, key: K, value: V) -> Option<(K, V)>
where
K: Clone,
V: Clone,
{
self.insert_entry_and(key, value, |k, v| (k.clone(), v.clone()))
}
#[inline]
pub fn insert_and<F: FnOnce(&V) -> T, T>(
&self,
key: K,
value: V,
with_previous_value: F,
) -> Option<T> {
self.insert_entry_and(key, value, move |_, v| with_previous_value(v))
}
#[inline]
pub fn insert_entry_and<F: FnOnce(&K, &V) -> T, T>(
&self,
key: K,
value: V,
with_previous_entry: F,
) -> Option<T> {
let hash = bucket::hash(&self.build_hasher, &key);
let result =
self.bucket_array_ref(hash)
.insert_entry_and(key, hash, value, with_previous_entry);
if result.is_none() {
self.len.fetch_add(1, Ordering::Relaxed);
}
result
}
#[inline]
pub fn remove<Q: Hash + Eq + ?Sized>(&self, key: &Q) -> Option<V>
where
K: Borrow<Q>,
V: Clone,
{
self.remove_entry_if_and(key, |_, _| true, |_, v| v.clone())
}
#[inline]
pub fn remove_entry<Q: Hash + Eq + ?Sized>(&self, key: &Q) -> Option<(K, V)>
where
K: Borrow<Q> + Clone,
V: Clone,
{
self.remove_entry_if_and(key, |_, _| true, |k, v| (k.clone(), v.clone()))
}
#[inline]
pub fn remove_and<Q: Hash + Eq + ?Sized, F: FnOnce(&V) -> T, T>(
&self,
key: &Q,
with_previous_value: F,
) -> Option<T>
where
K: Borrow<Q>,
{
self.remove_entry_if_and(key, |_, _| true, move |_, v| with_previous_value(v))
}
#[inline]
pub fn remove_entry_and<Q: Hash + Eq + ?Sized, F: FnOnce(&K, &V) -> T, T>(
&self,
key: &Q,
with_previous_entry: F,
) -> Option<T>
where
K: Borrow<Q>,
{
self.remove_entry_if_and(key, |_, _| true, with_previous_entry)
}
pub fn remove_if<Q: Hash + Eq + ?Sized, F: FnMut(&K, &V) -> bool>(
&self,
key: &Q,
condition: F,
) -> Option<V>
where
K: Borrow<Q>,
V: Clone,
{
self.remove_entry_if_and(key, condition, move |_, v| v.clone())
}
#[inline]
pub fn remove_entry_if<Q: Hash + Eq + ?Sized, F: FnMut(&K, &V) -> bool>(
&self,
key: &Q,
condition: F,
) -> Option<(K, V)>
where
K: Clone + Borrow<Q>,
V: Clone,
{
self.remove_entry_if_and(key, condition, move |k, v| (k.clone(), v.clone()))
}
#[inline]
pub fn remove_if_and<Q: Hash + Eq + ?Sized, F: FnMut(&K, &V) -> bool, G: FnOnce(&V) -> T, T>(
&self,
key: &Q,
condition: F,
with_previous_value: G,
) -> Option<T>
where
K: Borrow<Q>,
{
self.remove_entry_if_and(key, condition, move |_, v| with_previous_value(v))
}
#[inline]
pub fn remove_entry_if_and<
Q: Hash + Eq + ?Sized,
F: FnMut(&K, &V) -> bool,
G: FnOnce(&K, &V) -> T,
T,
>(
&self,
key: &Q,
condition: F,
with_previous_entry: G,
) -> Option<T>
where
K: Borrow<Q>,
{
let hash = bucket::hash(&self.build_hasher, &key);
self.bucket_array_ref(hash)
.remove_entry_if_and(key, hash, condition, move |k, v| {
self.len.fetch_sub(1, Ordering::Relaxed);
with_previous_entry(k, v)
})
}
#[inline]
pub fn insert_or_modify<F: FnMut(&K, &V) -> V>(
&self,
key: K,
value: V,
on_modify: F,
) -> Option<V>
where
V: Clone,
{
self.insert_with_or_modify_entry_and(key, move || value, on_modify, |_, v| v.clone())
}
#[inline]
pub fn insert_or_modify_entry<F: FnMut(&K, &V) -> V>(
&self,
key: K,
value: V,
on_modify: F,
) -> Option<(K, V)>
where
K: Clone,
V: Clone,
{
self.insert_with_or_modify_entry_and(
key,
move || value,
on_modify,
|k, v| (k.clone(), v.clone()),
)
}
#[inline]
pub fn insert_with_or_modify<F: FnOnce() -> V, G: FnMut(&K, &V) -> V>(
&self,
key: K,
on_insert: F,
on_modify: G,
) -> Option<V>
where
V: Clone,
{
self.insert_with_or_modify_entry_and(key, on_insert, on_modify, |_, v| v.clone())
}
#[inline]
pub fn insert_with_or_modify_entry<F: FnOnce() -> V, G: FnMut(&K, &V) -> V>(
&self,
key: K,
on_insert: F,
on_modify: G,
) -> Option<(K, V)>
where
K: Clone,
V: Clone,
{
self.insert_with_or_modify_entry_and(key, on_insert, on_modify, |k, v| {
(k.clone(), v.clone())
})
}
#[inline]
pub fn insert_or_modify_and<F: FnMut(&K, &V) -> V, G: FnOnce(&V) -> T, T>(
&self,
key: K,
value: V,
on_modify: F,
with_old_value: G,
) -> Option<T> {
self.insert_with_or_modify_entry_and(
key,
move || value,
on_modify,
move |_, v| with_old_value(v),
)
}
#[inline]
pub fn insert_or_modify_entry_and<F: FnMut(&K, &V) -> V, G: FnOnce(&K, &V) -> T, T>(
&self,
key: K,
value: V,
on_modify: F,
with_old_entry: G,
) -> Option<T> {
self.insert_with_or_modify_entry_and(key, move || value, on_modify, with_old_entry)
}
#[inline]
pub fn insert_with_or_modify_and<
F: FnOnce() -> V,
G: FnMut(&K, &V) -> V,
H: FnOnce(&V) -> T,
T,
>(
&self,
key: K,
on_insert: F,
on_modify: G,
with_old_value: H,
) -> Option<T> {
self.insert_with_or_modify_entry_and(key, on_insert, on_modify, move |_, v| {
with_old_value(v)
})
}
#[inline]
pub fn insert_with_or_modify_entry_and<
F: FnOnce() -> V,
G: FnMut(&K, &V) -> V,
H: FnOnce(&K, &V) -> T,
T,
>(
&self,
key: K,
on_insert: F,
on_modify: G,
with_old_entry: H,
) -> Option<T> {
let hash = bucket::hash(&self.build_hasher, &key);
let result = self.bucket_array_ref(hash).insert_with_or_modify_entry_and(
key,
hash,
on_insert,
on_modify,
with_old_entry,
);
if result.is_none() {
self.len.fetch_add(1, Ordering::Relaxed);
}
result
}
#[inline]
pub fn modify<F: FnMut(&K, &V) -> V>(&self, key: K, on_modify: F) -> Option<V>
where
V: Clone,
{
self.modify_entry_and(key, on_modify, |_, v| v.clone())
}
#[inline]
pub fn modify_entry<F: FnMut(&K, &V) -> V>(&self, key: K, on_modify: F) -> Option<(K, V)>
where
K: Clone,
V: Clone,
{
self.modify_entry_and(key, on_modify, |k, v| (k.clone(), v.clone()))
}
#[inline]
pub fn modify_and<F: FnMut(&K, &V) -> V, G: FnOnce(&V) -> T, T>(
&self,
key: K,
on_modify: F,
with_old_value: G,
) -> Option<T> {
self.modify_entry_and(key, on_modify, move |_, v| with_old_value(v))
}
#[inline]
pub fn modify_entry_and<F: FnMut(&K, &V) -> V, G: FnOnce(&K, &V) -> T, T>(
&self,
key: K,
on_modify: F,
with_old_entry: G,
) -> Option<T> {
let hash = bucket::hash(&self.build_hasher, &key);
self.bucket_array_ref(hash)
.modify_entry_and(key, hash, on_modify, with_old_entry)
}
}
#[cfg(feature = "num-cpus")]
impl<K, V, S: Default> Default for HashMap<K, V, S> {
fn default() -> Self {
HashMap::with_num_segments_capacity_and_hasher(default_num_segments(), 0, S::default())
}
}
impl<K, V, S> Drop for HashMap<K, V, S> {
fn drop(&mut self) {
let guard = unsafe { &crossbeam_epoch::unprotected() };
atomic::fence(Ordering::Acquire);
for Segment {
bucket_array: this_bucket_array,
..
} in self.segments.iter()
{
let mut current_ptr = this_bucket_array.load(Ordering::Relaxed, guard);
while let Some(current_ref) = unsafe { current_ptr.as_ref() } {
let next_ptr = current_ref.next.load(Ordering::Relaxed, guard);
for this_bucket_ptr in current_ref
.buckets
.iter()
.map(|b| b.load(Ordering::Relaxed, guard))
.filter(|p| !p.is_null())
.filter(|p| next_ptr.is_null() || p.tag() & bucket::TOMBSTONE_TAG == 0)
{
unsafe { bucket::defer_acquire_destroy(guard, this_bucket_ptr) };
}
unsafe { bucket::defer_acquire_destroy(guard, current_ptr) };
current_ptr = next_ptr;
}
}
}
}
impl<K, V, S> HashMap<K, V, S> {
#[inline]
fn bucket_array_ref(&'_ self, hash: u64) -> BucketArrayRef<'_, K, V, S> {
let index = self.segment_index_from_hash(hash);
let Segment {
ref bucket_array,
ref len,
} = self.segments[index];
BucketArrayRef {
bucket_array,
build_hasher: &self.build_hasher,
len,
}
}
#[inline]
fn segment_index_from_hash(&'_ self, hash: u64) -> usize {
if self.segment_shift == 64 {
0
} else {
(hash >> self.segment_shift) as usize
}
}
}
struct Segment<K, V> {
bucket_array: Atomic<BucketArray<K, V>>,
len: AtomicUsize,
}
#[cfg(feature = "num-cpus")]
fn default_num_segments() -> usize {
num_cpus::get() * 2
}
#[cfg(test)]
mod tests {
use crate::write_test_cases_for_me;
use super::*;
write_test_cases_for_me!(HashMap);
#[test]
fn single_segment() {
let map = HashMap::with_num_segments(1);
assert!(map.is_empty());
assert_eq!(map.len(), 0);
assert_eq!(map.insert("foo", 5), None);
assert_eq!(map.get("foo"), Some(5));
assert!(!map.is_empty());
assert_eq!(map.len(), 1);
assert_eq!(map.remove("foo"), Some(5));
assert!(map.is_empty());
assert_eq!(map.len(), 0);
}
}