compiler/rustc_data_structures/src/sharded.rs

   1 use crate::fx::{FxHashMap, FxHasher};
   2 use crate::sync::{Lock, LockGuard};
   3 use std::borrow::Borrow;
   4 use std::collections::hash_map::RawEntryMut;
   5 use std::hash::{Hash, Hasher};
   6 use std::mem;
   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 number of shards.
  16 const SHARD_BITS: usize = 5;
  17
  18 #[cfg(not(parallel_compiler))]
  19 const SHARD_BITS: usize = 0;
  20
  21 pub 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         Self::new(T::default)
  33     }
  34 }
  35
  36 impl<T> Sharded<T> {
  37     #[inline]
  38     pub fn new(mut value: impl FnMut() -> T) -> Self {
  39         Sharded { shards: [(); SHARDS].map(|()| CacheAligned(Lock::new(value()))) }
  40     }
  41
  42     /// The shard is selected by hashing `val` with `FxHasher`.
  43     #[inline]
  44     pub fn get_shard_by_value<K: Hash + ?Sized>(&self, val: &K) -> &Lock<T> {
  45         if SHARDS == 1 { &self.shards[0].0 } else { self.get_shard_by_hash(make_hash(val)) }
  46     }
  47
  48     #[inline]
  49     pub fn get_shard_by_hash(&self, hash: u64) -> &Lock<T> {
  50         &self.shards[get_shard_index_by_hash(hash)].0
  51     }
  52
  53     #[inline]
  54     pub fn get_shard_by_index(&self, i: usize) -> &Lock<T> {
  55         &self.shards[i].0
  56     }
  57
  58     pub fn lock_shards(&self) -> Vec<LockGuard<'_, T>> {
  59         (0..SHARDS).map(|i| self.shards[i].0.lock()).collect()
  60     }
  61
  62     pub fn try_lock_shards(&self) -> Option<Vec<LockGuard<'_, T>>> {
  63         (0..SHARDS).map(|i| self.shards[i].0.try_lock()).collect()
  64     }
  65 }
  66
  67 pub type ShardedHashMap<K, V> = Sharded<FxHashMap<K, V>>;
  68
  69 impl<K: Eq, V> ShardedHashMap<K, V> {
  70     pub fn len(&self) -> usize {
  71         self.lock_shards().iter().map(|shard| shard.len()).sum()
  72     }
  73 }
  74
  75 impl<K: Eq + Hash + Copy> ShardedHashMap<K, ()> {
  76     #[inline]
  77     pub fn intern_ref<Q: ?Sized>(&self, value: &Q, make: impl FnOnce() -> K) -> K
  78     where
  79         K: Borrow<Q>,
  80         Q: Hash + Eq,
  81     {
  82         let hash = make_hash(value);
  83         let mut shard = self.get_shard_by_hash(hash).lock();
  84         let entry = shard.raw_entry_mut().from_key_hashed_nocheck(hash, value);
  85
  86         match entry {
  87             RawEntryMut::Occupied(e) => *e.key(),
  88             RawEntryMut::Vacant(e) => {
  89                 let v = make();
  90                 e.insert_hashed_nocheck(hash, v, ());
  91                 v
  92             }
  93         }
  94     }
  95
  96     #[inline]
  97     pub fn intern<Q>(&self, value: Q, make: impl FnOnce(Q) -> K) -> K
  98     where
  99         K: Borrow<Q>,
 100         Q: Hash + Eq,
 101     {
 102         let hash = make_hash(&value);
 103         let mut shard = self.get_shard_by_hash(hash).lock();
 104         let entry = shard.raw_entry_mut().from_key_hashed_nocheck(hash, &value);
 105
 106         match entry {
 107             RawEntryMut::Occupied(e) => *e.key(),
 108             RawEntryMut::Vacant(e) => {
 109                 let v = make(value);
 110                 e.insert_hashed_nocheck(hash, v, ());
 111                 v
 112             }
 113         }
 114     }
 115 }
 116
 117 pub trait IntoPointer {
 118     /// Returns a pointer which outlives `self`.
 119     fn into_pointer(&self) -> *const ();
 120 }
 121
 122 impl<K: Eq + Hash + Copy + IntoPointer> ShardedHashMap<K, ()> {
 123     pub fn contains_pointer_to<T: Hash + IntoPointer>(&self, value: &T) -> bool {
 124         let hash = make_hash(&value);
 125         let shard = self.get_shard_by_hash(hash).lock();
 126         let value = value.into_pointer();
 127         shard.raw_entry().from_hash(hash, |entry| entry.into_pointer() == value).is_some()
 128     }
 129 }
 130
 131 #[inline]
 132 fn make_hash<K: Hash + ?Sized>(val: &K) -> u64 {
 133     let mut state = FxHasher::default();
 134     val.hash(&mut state);
 135     state.finish()
 136 }
 137
 138 /// Get a shard with a pre-computed hash value. If `get_shard_by_value` is
 139 /// ever used in combination with `get_shard_by_hash` on a single `Sharded`
 140 /// instance, then `hash` must be computed with `FxHasher`. Otherwise,
 141 /// `hash` can be computed with any hasher, so long as that hasher is used
 142 /// consistently for each `Sharded` instance.
 143 #[inline]
 144 pub fn get_shard_index_by_hash(hash: u64) -> usize {
 145     let hash_len = mem::size_of::<usize>();
 146     // Ignore the top 7 bits as hashbrown uses these and get the next SHARD_BITS highest bits.
 147     // hashbrown also uses the lowest bits, so we can't use those
 148     let bits = (hash >> (hash_len * 8 - 7 - SHARD_BITS)) as usize;
 149     bits % SHARDS
 150 }