rustc_data_structures/
sharded.rs

1use std::borrow::Borrow;
2use std::collections::hash_map::RawEntryMut;
3use std::hash::{Hash, Hasher};
4use std::{iter, mem};
5
6use either::Either;
7
8use crate::fx::{FxHashMap, FxHasher};
9use crate::sync::{CacheAligned, Lock, LockGuard, Mode, is_dyn_thread_safe};
10
11// 32 shards is sufficient to reduce contention on an 8-core Ryzen 7 1700,
12// but this should be tested on higher core count CPUs. How the `Sharded` type gets used
13// may also affect the ideal number of shards.
14const SHARD_BITS: usize = 5;
15
16const SHARDS: usize = 1 << SHARD_BITS;
17
18/// An array of cache-line aligned inner locked structures with convenience methods.
19/// A single field is used when the compiler uses only one thread.
20pub enum Sharded<T> {
21    Single(Lock<T>),
22    Shards(Box<[CacheAligned<Lock<T>>; SHARDS]>),
23}
24
25impl<T: Default> Default for Sharded<T> {
26    #[inline]
27    fn default() -> Self {
28        Self::new(T::default)
29    }
30}
31
32impl<T> Sharded<T> {
33    #[inline]
34    pub fn new(mut value: impl FnMut() -> T) -> Self {
35        if is_dyn_thread_safe() {
36            return Sharded::Shards(Box::new(
37                [(); SHARDS].map(|()| CacheAligned(Lock::new(value()))),
38            ));
39        }
40
41        Sharded::Single(Lock::new(value()))
42    }
43
44    /// The shard is selected by hashing `val` with `FxHasher`.
45    #[inline]
46    pub fn get_shard_by_value<K: Hash + ?Sized>(&self, _val: &K) -> &Lock<T> {
47        match self {
48            Self::Single(single) => single,
49            Self::Shards(..) => self.get_shard_by_hash(make_hash(_val)),
50        }
51    }
52
53    #[inline]
54    pub fn get_shard_by_hash(&self, hash: u64) -> &Lock<T> {
55        self.get_shard_by_index(get_shard_hash(hash))
56    }
57
58    #[inline]
59    pub fn get_shard_by_index(&self, _i: usize) -> &Lock<T> {
60        match self {
61            Self::Single(single) => single,
62            Self::Shards(shards) => {
63                // SAFETY: The index gets ANDed with the shard mask, ensuring it is always inbounds.
64                unsafe { &shards.get_unchecked(_i & (SHARDS - 1)).0 }
65            }
66        }
67    }
68
69    /// The shard is selected by hashing `val` with `FxHasher`.
70    #[inline]
71    #[track_caller]
72    pub fn lock_shard_by_value<K: Hash + ?Sized>(&self, _val: &K) -> LockGuard<'_, T> {
73        match self {
74            Self::Single(single) => {
75                // Synchronization is disabled so use the `lock_assume_no_sync` method optimized
76                // for that case.
77
78                // SAFETY: We know `is_dyn_thread_safe` was false when creating the lock thus
79                // `might_be_dyn_thread_safe` was also false.
80                unsafe { single.lock_assume(Mode::NoSync) }
81            }
82            Self::Shards(..) => self.lock_shard_by_hash(make_hash(_val)),
83        }
84    }
85
86    #[inline]
87    #[track_caller]
88    pub fn lock_shard_by_hash(&self, hash: u64) -> LockGuard<'_, T> {
89        self.lock_shard_by_index(get_shard_hash(hash))
90    }
91
92    #[inline]
93    #[track_caller]
94    pub fn lock_shard_by_index(&self, _i: usize) -> LockGuard<'_, T> {
95        match self {
96            Self::Single(single) => {
97                // Synchronization is disabled so use the `lock_assume_no_sync` method optimized
98                // for that case.
99
100                // SAFETY: We know `is_dyn_thread_safe` was false when creating the lock thus
101                // `might_be_dyn_thread_safe` was also false.
102                unsafe { single.lock_assume(Mode::NoSync) }
103            }
104            Self::Shards(shards) => {
105                // Synchronization is enabled so use the `lock_assume_sync` method optimized
106                // for that case.
107
108                // SAFETY (get_unchecked): The index gets ANDed with the shard mask, ensuring it is
109                // always inbounds.
110                // SAFETY (lock_assume_sync): We know `is_dyn_thread_safe` was true when creating
111                // the lock thus `might_be_dyn_thread_safe` was also true.
112                unsafe { shards.get_unchecked(_i & (SHARDS - 1)).0.lock_assume(Mode::Sync) }
113            }
114        }
115    }
116
117    #[inline]
118    pub fn lock_shards(&self) -> impl Iterator<Item = LockGuard<'_, T>> {
119        match self {
120            Self::Single(single) => Either::Left(iter::once(single.lock())),
121            Self::Shards(shards) => Either::Right(shards.iter().map(|shard| shard.0.lock())),
122        }
123    }
124
125    #[inline]
126    pub fn try_lock_shards(&self) -> impl Iterator<Item = Option<LockGuard<'_, T>>> {
127        match self {
128            Self::Single(single) => Either::Left(iter::once(single.try_lock())),
129            Self::Shards(shards) => Either::Right(shards.iter().map(|shard| shard.0.try_lock())),
130        }
131    }
132}
133
134#[inline]
135pub fn shards() -> usize {
136    if is_dyn_thread_safe() {
137        return SHARDS;
138    }
139
140    1
141}
142
143pub type ShardedHashMap<K, V> = Sharded<FxHashMap<K, V>>;
144
145impl<K: Eq, V> ShardedHashMap<K, V> {
146    pub fn len(&self) -> usize {
147        self.lock_shards().map(|shard| shard.len()).sum()
148    }
149}
150
151impl<K: Eq + Hash + Copy> ShardedHashMap<K, ()> {
152    #[inline]
153    pub fn intern_ref<Q: ?Sized>(&self, value: &Q, make: impl FnOnce() -> K) -> K
154    where
155        K: Borrow<Q>,
156        Q: Hash + Eq,
157    {
158        let hash = make_hash(value);
159        let mut shard = self.lock_shard_by_hash(hash);
160        let entry = shard.raw_entry_mut().from_key_hashed_nocheck(hash, value);
161
162        match entry {
163            RawEntryMut::Occupied(e) => *e.key(),
164            RawEntryMut::Vacant(e) => {
165                let v = make();
166                e.insert_hashed_nocheck(hash, v, ());
167                v
168            }
169        }
170    }
171
172    #[inline]
173    pub fn intern<Q>(&self, value: Q, make: impl FnOnce(Q) -> K) -> K
174    where
175        K: Borrow<Q>,
176        Q: Hash + Eq,
177    {
178        let hash = make_hash(&value);
179        let mut shard = self.lock_shard_by_hash(hash);
180        let entry = shard.raw_entry_mut().from_key_hashed_nocheck(hash, &value);
181
182        match entry {
183            RawEntryMut::Occupied(e) => *e.key(),
184            RawEntryMut::Vacant(e) => {
185                let v = make(value);
186                e.insert_hashed_nocheck(hash, v, ());
187                v
188            }
189        }
190    }
191}
192
193pub trait IntoPointer {
194    /// Returns a pointer which outlives `self`.
195    fn into_pointer(&self) -> *const ();
196}
197
198impl<K: Eq + Hash + Copy + IntoPointer> ShardedHashMap<K, ()> {
199    pub fn contains_pointer_to<T: Hash + IntoPointer>(&self, value: &T) -> bool {
200        let hash = make_hash(&value);
201        let shard = self.lock_shard_by_hash(hash);
202        let value = value.into_pointer();
203        shard.raw_entry().from_hash(hash, |entry| entry.into_pointer() == value).is_some()
204    }
205}
206
207#[inline]
208pub fn make_hash<K: Hash + ?Sized>(val: &K) -> u64 {
209    let mut state = FxHasher::default();
210    val.hash(&mut state);
211    state.finish()
212}
213
214/// Get a shard with a pre-computed hash value. If `get_shard_by_value` is
215/// ever used in combination with `get_shard_by_hash` on a single `Sharded`
216/// instance, then `hash` must be computed with `FxHasher`. Otherwise,
217/// `hash` can be computed with any hasher, so long as that hasher is used
218/// consistently for each `Sharded` instance.
219#[inline]
220fn get_shard_hash(hash: u64) -> usize {
221    let hash_len = mem::size_of::<usize>();
222    // Ignore the top 7 bits as hashbrown uses these and get the next SHARD_BITS highest bits.
223    // hashbrown also uses the lowest bits, so we can't use those
224    (hash >> (hash_len * 8 - 7 - SHARD_BITS)) as usize
225}
rustc_data_structures/sharded.rs

rustc_data_structures/
sharded.rs
