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