compiler/rustc_data_structures/src/sharded.rs

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