rustc_data_structures/
sorted_map.rs

Help
1use std::borrow::Borrow;
2use std::cmp::Ordering;
3use std::fmt::Debug;
4use std::mem;
5use std::ops::{Bound, Index, IndexMut, RangeBounds};
6
7use rustc_macros::{Decodable_NoContext, Encodable_NoContext};
8
9use crate::stable_hasher::{HashStable, StableHasher, StableOrd};
10
11mod index_map;
12
13pub use index_map::SortedIndexMultiMap;
14
15/// `SortedMap` is a data structure with similar characteristics as BTreeMap but
16/// slightly different trade-offs: lookup is *O*(log(*n*)), insertion and removal
17/// are *O*(*n*) but elements can be iterated in order cheaply.
18///
19/// `SortedMap` can be faster than a `BTreeMap` for small sizes (<50) since it
20/// stores data in a more compact way. It also supports accessing contiguous
21/// ranges of elements as a slice, and slices of already sorted elements can be
22/// inserted efficiently.
23#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Encodable_NoContext, Decodable_NoContext)]
24pub struct SortedMap<K, V> {
25    data: Vec<(K, V)>,
26}
27
28impl<K, V> Default for SortedMap<K, V> {
29    #[inline]
30    fn default() -> SortedMap<K, V> {
31        SortedMap { data: Vec::new() }
32    }
33}
34
35impl<K, V> SortedMap<K, V> {
36    #[inline]
37    pub const fn new() -> SortedMap<K, V> {
38        SortedMap { data: Vec::new() }
39    }
40}
41
42impl<K: Ord, V> SortedMap<K, V> {
43    /// Construct a `SortedMap` from a presorted set of elements. This is faster
44    /// than creating an empty map and then inserting the elements individually.
45    ///
46    /// It is up to the caller to make sure that the elements are sorted by key
47    /// and that there are no duplicates.
48    #[inline]
49    pub fn from_presorted_elements(elements: Vec<(K, V)>) -> SortedMap<K, V> {
50        debug_assert!(elements.array_windows().all(|[fst, snd]| fst.0 < snd.0));
51
52        SortedMap { data: elements }
53    }
54
55    #[inline]
56    pub fn insert(&mut self, key: K, value: V) -> Option<V> {
57        match self.lookup_index_for(&key) {
58            Ok(index) => {
59                let slot = unsafe { self.data.get_unchecked_mut(index) };
60                Some(mem::replace(&mut slot.1, value))
61            }
62            Err(index) => {
63                self.data.insert(index, (key, value));
64                None
65            }
66        }
67    }
68
69    #[inline]
70    pub fn remove(&mut self, key: &K) -> Option<V> {
71        match self.lookup_index_for(key) {
72            Ok(index) => Some(self.data.remove(index).1),
73            Err(_) => None,
74        }
75    }
76
77    #[inline]
78    pub fn get<Q>(&self, key: &Q) -> Option<&V>
79    where
80        K: Borrow<Q>,
81        Q: Ord + ?Sized,
82    {
83        match self.lookup_index_for(key) {
84            Ok(index) => unsafe { Some(&self.data.get_unchecked(index).1) },
85            Err(_) => None,
86        }
87    }
88
89    #[inline]
90    pub fn get_mut<Q>(&mut self, key: &Q) -> Option<&mut V>
91    where
92        K: Borrow<Q>,
93        Q: Ord + ?Sized,
94    {
95        match self.lookup_index_for(key) {
96            Ok(index) => unsafe { Some(&mut self.data.get_unchecked_mut(index).1) },
97            Err(_) => None,
98        }
99    }
100
101    /// Gets a mutable reference to the value in the entry, or insert a new one.
102    #[inline]
103    pub fn get_mut_or_insert_default(&mut self, key: K) -> &mut V
104    where
105        K: Eq,
106        V: Default,
107    {
108        let index = match self.lookup_index_for(&key) {
109            Ok(index) => index,
110            Err(index) => {
111                self.data.insert(index, (key, V::default()));
112                index
113            }
114        };
115        unsafe { &mut self.data.get_unchecked_mut(index).1 }
116    }
117
118    #[inline]
119    pub fn clear(&mut self) {
120        self.data.clear();
121    }
122
123    /// Iterate over elements, sorted by key
124    #[inline]
125    pub fn iter(&self) -> std::slice::Iter<'_, (K, V)> {
126        self.data.iter()
127    }
128
129    /// Iterate over the keys, sorted
130    #[inline]
131    pub fn keys(&self) -> impl ExactSizeIterator<Item = &K> + DoubleEndedIterator {
132        self.data.iter().map(|(k, _)| k)
133    }
134
135    /// Iterate over values, sorted by key
136    #[inline]
137    pub fn values(&self) -> impl ExactSizeIterator<Item = &V> + DoubleEndedIterator {
138        self.data.iter().map(|(_, v)| v)
139    }
140
141    #[inline]
142    pub fn len(&self) -> usize {
143        self.data.len()
144    }
145
146    #[inline]
147    pub fn is_empty(&self) -> bool {
148        self.len() == 0
149    }
150
151    #[inline]
152    pub fn range<R>(&self, range: R) -> &[(K, V)]
153    where
154        R: RangeBounds<K>,
155    {
156        let (start, end) = self.range_slice_indices(range);
157        &self.data[start..end]
158    }
159
160    /// `sm.range_is_empty(r)` == `sm.range(r).is_empty()`, but is faster.
161    #[inline]
162    pub fn range_is_empty<R>(&self, range: R) -> bool
163    where
164        R: RangeBounds<K>,
165    {
166        // `range` must (via `range_slice_indices`) search for the start and
167        // end separately. But here we can do a single binary search for the
168        // entire range. If a single `x` matching `range` is found then the
169        // range is *not* empty.
170        self.data
171            .binary_search_by(|(x, _)| {
172                // Is `x` below `range`?
173                match range.start_bound() {
174                    Bound::Included(start) if x < start => return Ordering::Less,
175                    Bound::Excluded(start) if x <= start => return Ordering::Less,
176                    _ => {}
177                };
178
179                // Is `x` above `range`?
180                match range.end_bound() {
181                    Bound::Included(end) if x > end => return Ordering::Greater,
182                    Bound::Excluded(end) if x >= end => return Ordering::Greater,
183                    _ => {}
184                };
185
186                // `x` must be within `range`.
187                Ordering::Equal
188            })
189            .is_err()
190    }
191
192    #[inline]
193    pub fn remove_range<R>(&mut self, range: R)
194    where
195        R: RangeBounds<K>,
196    {
197        let (start, end) = self.range_slice_indices(range);
198        self.data.splice(start..end, std::iter::empty());
199    }
200
201    /// Mutate all keys with the given function `f`. This mutation must not
202    /// change the sort-order of keys.
203    #[inline]
204    pub fn offset_keys<F>(&mut self, f: F)
205    where
206        F: Fn(&mut K),
207    {
208        self.data.iter_mut().map(|(k, _)| k).for_each(f);
209    }
210
211    /// Inserts a presorted range of elements into the map. If the range can be
212    /// inserted as a whole in between to existing elements of the map, this
213    /// will be faster than inserting the elements individually.
214    ///
215    /// It is up to the caller to make sure that the elements are sorted by key
216    /// and that there are no duplicates.
217    #[inline]
218    pub fn insert_presorted(&mut self, elements: Vec<(K, V)>) {
219        if elements.is_empty() {
220            return;
221        }
222
223        debug_assert!(elements.array_windows().all(|[fst, snd]| fst.0 < snd.0));
224
225        let start_index = self.lookup_index_for(&elements[0].0);
226
227        let elements = match start_index {
228            Ok(index) => {
229                let mut elements = elements.into_iter();
230                self.data[index] = elements.next().unwrap();
231                elements
232            }
233            Err(index) => {
234                if index == self.data.len() || elements.last().unwrap().0 < self.data[index].0 {
235                    // We can copy the whole range without having to mix with
236                    // existing elements.
237                    self.data.splice(index..index, elements);
238                    return;
239                }
240
241                let mut elements = elements.into_iter();
242                self.data.insert(index, elements.next().unwrap());
243                elements
244            }
245        };
246
247        // Insert the rest
248        for (k, v) in elements {
249            self.insert(k, v);
250        }
251    }
252
253    /// Looks up the key in `self.data` via `slice::binary_search()`.
254    #[inline(always)]
255    fn lookup_index_for<Q>(&self, key: &Q) -> Result<usize, usize>
256    where
257        K: Borrow<Q>,
258        Q: Ord + ?Sized,
259    {
260        self.data.binary_search_by(|(x, _)| x.borrow().cmp(key))
261    }
262
263    #[inline]
264    fn range_slice_indices<R>(&self, range: R) -> (usize, usize)
265    where
266        R: RangeBounds<K>,
267    {
268        let start = match range.start_bound() {
269            Bound::Included(k) => match self.lookup_index_for(k) {
270                Ok(index) | Err(index) => index,
271            },
272            Bound::Excluded(k) => match self.lookup_index_for(k) {
273                Ok(index) => index + 1,
274                Err(index) => index,
275            },
276            Bound::Unbounded => 0,
277        };
278
279        let end = match range.end_bound() {
280            Bound::Included(k) => match self.lookup_index_for(k) {
281                Ok(index) => index + 1,
282                Err(index) => index,
283            },
284            Bound::Excluded(k) => match self.lookup_index_for(k) {
285                Ok(index) | Err(index) => index,
286            },
287            Bound::Unbounded => self.data.len(),
288        };
289
290        (start, end)
291    }
292
293    #[inline]
294    pub fn contains_key<Q>(&self, key: &Q) -> bool
295    where
296        K: Borrow<Q>,
297        Q: Ord + ?Sized,
298    {
299        self.get(key).is_some()
300    }
301}
302
303impl<K: Ord, V> IntoIterator for SortedMap<K, V> {
304    type Item = (K, V);
305    type IntoIter = std::vec::IntoIter<(K, V)>;
306
307    fn into_iter(self) -> Self::IntoIter {
308        self.data.into_iter()
309    }
310}
311
312impl<'a, K, Q, V> Index<&'a Q> for SortedMap<K, V>
313where
314    K: Ord + Borrow<Q>,
315    Q: Ord + ?Sized,
316{
317    type Output = V;
318
319    fn index(&self, key: &Q) -> &Self::Output {
320        self.get(key).expect("no entry found for key")
321    }
322}
323
324impl<'a, K, Q, V> IndexMut<&'a Q> for SortedMap<K, V>
325where
326    K: Ord + Borrow<Q>,
327    Q: Ord + ?Sized,
328{
329    fn index_mut(&mut self, key: &Q) -> &mut Self::Output {
330        self.get_mut(key).expect("no entry found for key")
331    }
332}
333
334impl<K: Ord, V> FromIterator<(K, V)> for SortedMap<K, V> {
335    fn from_iter<T: IntoIterator<Item = (K, V)>>(iter: T) -> Self {
336        let mut data: Vec<(K, V)> = iter.into_iter().collect();
337
338        data.sort_unstable_by(|(k1, _), (k2, _)| k1.cmp(k2));
339        data.dedup_by(|(k1, _), (k2, _)| k1 == k2);
340
341        SortedMap { data }
342    }
343}
344
345impl<K: HashStable<CTX> + StableOrd, V: HashStable<CTX>, CTX> HashStable<CTX> for SortedMap<K, V> {
346    #[inline]
347    fn hash_stable(&self, ctx: &mut CTX, hasher: &mut StableHasher) {
348        self.data.hash_stable(ctx, hasher);
349    }
350}
351
352impl<K: Debug, V: Debug> Debug for SortedMap<K, V> {
353    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
354        f.debug_map().entries(self.data.iter().map(|(a, b)| (a, b))).finish()
355    }
356}
357
358#[cfg(test)]
359mod tests;
rustc_data_structures/sorted_map.rs

rustc_data_structures/
sorted_map.rs
