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