src/librustc_data_structures/sharded.rs

   1 use std::hash::{Hasher, Hash};
   2 use std::mem;
   3 use std::borrow::Borrow;
   4 use std::collections::hash_map::RawEntryMut;
   5 use smallvec::SmallVec;
   6 use crate::fx::{FxHasher, FxHashMap};
   7 use crate::sync::{Lock, LockGuard};
   8
   9 #[derive(Clone, Default)]
  10 #[cfg_attr(parallel_compiler, repr(align(64)))]
  11 struct CacheAligned<T>(T);
  12
  13 #[cfg(parallel_compiler)]
  14 // 32 shards is sufficient to reduce contention on an 8-core Ryzen 7 1700,
  15 // but this should be tested on higher core count CPUs. How the `Sharded` type gets used
  16 // may also affect the ideal nunber of shards.
  17 const SHARD_BITS: usize = 5;
  18
  19 #[cfg(not(parallel_compiler))]
  20 const SHARD_BITS: usize = 0;
  21
  22 pub const SHARDS: usize = 1 << SHARD_BITS;
  23
  24 /// An array of cache-line aligned inner locked structures with convenience methods.
  25 #[derive(Clone)]
  26 pub struct Sharded<T> {
  27     shards: [CacheAligned<Lock<T>>; SHARDS],
  28 }
  29
  30 impl<T: Default> Default for Sharded<T> {
  31     #[inline]
  32     fn default() -> Self {
  33         Self::new(|| T::default())
  34     }
  35 }
  36
  37 impl<T> Sharded<T> {
  38     #[inline]
  39     pub fn new(mut value: impl FnMut() -> T) -> Self {
  40         // Create a vector of the values we want
  41         let mut values: SmallVec<[_; SHARDS]> = (0..SHARDS).map(|_| {
  42             CacheAligned(Lock::new(value()))
  43         }).collect();
  44
  45         // Create an unintialized array
  46         let mut shards: mem::MaybeUninit<[CacheAligned<Lock<T>>; SHARDS]> =
  47             mem::MaybeUninit::uninit();
  48
  49         unsafe {
  50             // Copy the values into our array
  51             let first = shards.as_mut_ptr() as *mut CacheAligned<Lock<T>>;
  52             values.as_ptr().copy_to_nonoverlapping(first, SHARDS);
  53
  54             // Ignore the content of the vector
  55             values.set_len(0);
  56
  57             Sharded {
  58                 shards: shards.assume_init(),
  59             }
  60         }
  61     }
  62
  63     #[inline]
  64     pub fn get_shard_by_value<K: Hash + ?Sized>(&self, val: &K) -> &Lock<T> {
  65         if SHARDS == 1 {
  66             &self.shards[0].0
  67         } else {
  68             self.get_shard_by_hash(make_hash(val))
  69         }
  70     }
  71
  72     #[inline]
  73     pub fn get_shard_by_hash(&self, hash: u64) -> &Lock<T> {
  74         let hash_len = mem::size_of::<usize>();
  75         // Ignore the top 7 bits as hashbrown uses these and get the next SHARD_BITS highest bits.
  76         // hashbrown also uses the lowest bits, so we can't use those
  77         let bits = (hash >> (hash_len * 8 - 7 - SHARD_BITS)) as usize;
  78         let i = bits % SHARDS;
  79         &self.shards[i].0
  80     }
  81
  82     pub fn lock_shards(&self) -> Vec<LockGuard<'_, T>> {
  83         (0..SHARDS).map(|i| self.shards[i].0.lock()).collect()
  84     }
  85
  86     pub fn try_lock_shards(&self) -> Option<Vec<LockGuard<'_, T>>> {
  87         (0..SHARDS).map(|i| self.shards[i].0.try_lock()).collect()
  88     }
  89 }
  90
  91 pub type ShardedHashMap<K, V> = Sharded<FxHashMap<K, V>>;
  92
  93 impl<K: Eq, V> ShardedHashMap<K, V> {
  94     pub fn len(&self) -> usize {
  95         self.lock_shards().iter().map(|shard| shard.len()).sum()
  96     }
  97 }
  98
  99 impl<K: Eq + Hash + Copy> ShardedHashMap<K, ()> {
 100     #[inline]
 101     pub fn intern_ref<Q: ?Sized>(&self, value: &Q, make: impl FnOnce() -> K) -> K
 102         where K: Borrow<Q>,
 103               Q: Hash + Eq
 104     {
 105         let hash = make_hash(value);
 106         let mut shard = self.get_shard_by_hash(hash).lock();
 107         let entry = shard.raw_entry_mut().from_key_hashed_nocheck(hash, value);
 108
 109         match entry {
 110             RawEntryMut::Occupied(e) => *e.key(),
 111             RawEntryMut::Vacant(e) => {
 112                 let v = make();
 113                 e.insert_hashed_nocheck(hash, v, ());
 114                 v
 115             }
 116         }
 117     }
 118
 119     #[inline]
 120     pub fn intern<Q>(&self, value: Q, make: impl FnOnce(Q) -> K) -> K
 121         where K: Borrow<Q>,
 122               Q: Hash + Eq
 123     {
 124         let hash = make_hash(&value);
 125         let mut shard = self.get_shard_by_hash(hash).lock();
 126         let entry = shard.raw_entry_mut().from_key_hashed_nocheck(hash, &value);
 127
 128         match entry {
 129             RawEntryMut::Occupied(e) => *e.key(),
 130             RawEntryMut::Vacant(e) => {
 131                 let v = make(value);
 132                 e.insert_hashed_nocheck(hash, v, ());
 133                 v
 134             }
 135         }
 136     }
 137 }
 138
 139 #[inline]
 140 fn make_hash<K: Hash + ?Sized>(val: &K) -> u64 {
 141     let mut state = FxHasher::default();
 142     val.hash(&mut state);
 143     state.finish()
 144 }