compiler/rustc_data_structures/src/sorted_map.rs

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