1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132

#![doc = include_str!("../README.md")]

mod encapsulate;

use parking_lot::RwLock;
use std::{
    collections::{hash_map::RandomState, HashMap as Map},
    hash::{BuildHasher, Hash, Hasher},
};

pub use encapsulate::*;

pub struct HashMap<K, V, S = RandomState> {
    hasher: S,
    shift: usize,
    shards: Box<[RwLock<Map<K, V, S>>]>,
}

impl<K, V> HashMap<K, V> {
    /// Creates an empty `HashMap`.
    ///
    /// The hash map is initially created with a capacity of 0, so it will not allocate until it
    /// is first inserted into.
    #[inline]
    pub fn new() -> Self {
        Self::with_hasher(RandomState::default())
    }

    /// Creates an empty `HashMap` with the specified capacity.
    ///
    /// The hash map will be able to hold at least `capacity` elements without
    /// reallocating. If `capacity` is 0, the hash map will not allocate.
    #[inline]
    pub fn with_capacity(capacity: usize) -> Self {
        Self::with_capacity_and_hasher(capacity, RandomState::default())
    }
}

impl<K: Hash, V, S: BuildHasher> HashMap<K, V, S> {
    #[inline]
    pub fn read<'a>(&'a self, key: &'a K) -> Readable<'a, K, V, S> {
        Readable {
            map: self.shard(key).read(),
            key,
        }
    }

    #[inline]
    pub fn write(&self, key: K) -> Writeable<K, V, S> {
        Writeable {
            map: self.shard(&key).write(),
            key,
        }
    }

    #[inline]
    pub fn shard(&self, key: &K) -> &RwLock<Map<K, V, S>> {
        let idx = self.shard_idx(key);
        unsafe { self.shards.get_unchecked(idx) }
    }

    pub fn shard_idx(&self, key: &K) -> usize {
        let mut hasher = self.hasher.build_hasher();
        key.hash(&mut hasher);
        let hash = hasher.finish() as usize;
        // Leave the high 7 bits for the HashBrown SIMD tag.
        (hash << 7) >> self.shift
    }
}

impl<K, V, S: Clone> HashMap<K, V, S> {
    #[inline]
    /// Get a reference to the map's shards.
    pub fn shards(&self) -> &[RwLock<Map<K, V, S>>] {
        self.shards.as_ref()
    }

    /// Creates an empty `HashMap` which will use the given hash builder to hash
    /// keys.
    ///
    /// The created map has the default initial capacity.
    ///
    /// Warning: `hash_builder` is normally randomly generated, and
    /// is designed to allow HashMaps to be resistant to attacks that
    /// cause many collisions and very poor performance. Setting it
    /// manually using this function can expose a DoS attack vector.
    ///
    /// The `hash_builder` passed should implement the [`BuildHasher`] trait for
    /// the HashMap to be useful.
    #[inline]
    pub fn with_hasher(hasher: S) -> Self {
        Self::with_capacity_and_hasher(0, hasher)
    }

    /// Creates an empty `HashMap` with the specified capacity, using `hash_builder`
    /// to hash the keys.
    ///
    /// The hash map will be able to hold at least `capacity` elements without
    /// reallocating. If `capacity` is 0, the hash map will not allocate.
    ///
    /// Warning: `hash_builder` is normally randomly generated, and
    /// is designed to allow HashMaps to be resistant to attacks that
    /// cause many collisions and very poor performance. Setting it
    /// manually using this function can expose a DoS attack vector.
    ///
    /// The `hash_builder` passed should implement the [`BuildHasher`] trait for
    /// the HashMap to be useful.
    pub fn with_capacity_and_hasher(mut capacity: usize, hasher: S) -> Self {
        let shard_amount = (num_cpus::get() * 4).next_power_of_two();
        if capacity != 0 {
            capacity = (capacity + (shard_amount - 1)) & !(shard_amount - 1);
        }
        Self {
            shift: std::mem::size_of::<usize>() * 8 - shard_amount.trailing_zeros() as usize,
            shards: (0..shard_amount)
                .map(|_| {
                    RwLock::new(Map::with_capacity_and_hasher(
                        capacity / shard_amount,
                        hasher.clone(),
                    ))
                })
                .collect(),
            hasher,
        }
    }
}

impl<K, V> Default for HashMap<K, V> {
    fn default() -> Self {
        Self::new()
    }
}