src/librustc_data_structures/sorted_map.rs

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