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