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::{RangeBounds, Bound, Index, IndexMut};
   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(Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Default, Debug, RustcEncodable,
  16          RustcDecodable)]
  17 pub struct SortedMap<K: Ord, V> {
  18     data: Vec<(K, V)>
  19 }
  20
  21 impl<K: Ord, V> SortedMap<K, V> {
  22     #[inline]
  23     pub fn new() -> SortedMap<K, V> {
  24         SortedMap {
  25             data: vec![]
  26         }
  27     }
  28
  29     /// Construct a `SortedMap` from a presorted set of elements. This is faster
  30     /// than creating an empty map and then inserting the elements individually.
  31     ///
  32     /// It is up to the caller to make sure that the elements are sorted by key
  33     /// and that there are no duplicates.
  34     #[inline]
  35     pub fn from_presorted_elements(elements: Vec<(K, V)>) -> SortedMap<K, V>
  36     {
  37         debug_assert!(elements.windows(2).all(|w| w[0].0 < w[1].0));
  38
  39         SortedMap {
  40             data: elements
  41         }
  42     }
  43
  44     #[inline]
  45     pub fn insert(&mut self, key: K, mut value: V) -> Option<V> {
  46         match self.lookup_index_for(&key) {
  47             Ok(index) => {
  48                 let slot = unsafe {
  49                     self.data.get_unchecked_mut(index)
  50                 };
  51                 mem::swap(&mut slot.1, &mut value);
  52                 Some(value)
  53             }
  54             Err(index) => {
  55                 self.data.insert(index, (key, value));
  56                 None
  57             }
  58         }
  59     }
  60
  61     #[inline]
  62     pub fn remove(&mut self, key: &K) -> Option<V> {
  63         match self.lookup_index_for(key) {
  64             Ok(index) => {
  65                 Some(self.data.remove(index).1)
  66             }
  67             Err(_) => {
  68                 None
  69             }
  70         }
  71     }
  72
  73     #[inline]
  74     pub fn get<Q>(&self, key: &Q) -> Option<&V>
  75         where K: Borrow<Q>,
  76               Q: Ord + ?Sized
  77     {
  78         match self.lookup_index_for(key) {
  79             Ok(index) => {
  80                 unsafe {
  81                     Some(&self.data.get_unchecked(index).1)
  82                 }
  83             }
  84             Err(_) => {
  85                 None
  86             }
  87         }
  88     }
  89
  90     #[inline]
  91     pub fn get_mut<Q>(&mut self, key: &Q) -> Option<&mut V>
  92         where K: Borrow<Q>,
  93               Q: Ord + ?Sized
  94     {
  95         match self.lookup_index_for(key) {
  96             Ok(index) => {
  97                 unsafe {
  98                     Some(&mut self.data.get_unchecked_mut(index).1)
  99                 }
 100             }
 101             Err(_) => {
 102                 None
 103             }
 104         }
 105     }
 106
 107     #[inline]
 108     pub fn clear(&mut self) {
 109         self.data.clear();
 110     }
 111
 112     /// Iterate over elements, sorted by key
 113     #[inline]
 114     pub fn iter(&self) -> ::std::slice::Iter<'_, (K, V)> {
 115         self.data.iter()
 116     }
 117
 118     /// Iterate over the keys, sorted
 119     #[inline]
 120     pub fn keys(&self) -> impl Iterator<Item = &K> + ExactSizeIterator {
 121         self.data.iter().map(|&(ref k, _)| k)
 122     }
 123
 124     /// Iterate over values, sorted by key
 125     #[inline]
 126     pub fn values(&self) -> impl Iterator<Item = &V> + ExactSizeIterator {
 127         self.data.iter().map(|&(_, ref v)| v)
 128     }
 129
 130     #[inline]
 131     pub fn len(&self) -> usize {
 132         self.data.len()
 133     }
 134
 135     #[inline]
 136     pub fn is_empty(&self) -> bool {
 137         self.len() == 0
 138     }
 139
 140     #[inline]
 141     pub fn range<R>(&self, range: R) -> &[(K, V)]
 142         where R: RangeBounds<K>
 143     {
 144         let (start, end) = self.range_slice_indices(range);
 145         (&self.data[start .. end])
 146     }
 147
 148     #[inline]
 149     pub fn remove_range<R>(&mut self, range: R)
 150         where R: RangeBounds<K>
 151     {
 152         let (start, end) = self.range_slice_indices(range);
 153         self.data.splice(start .. end, ::std::iter::empty());
 154     }
 155
 156     /// Mutate all keys with the given function `f`. This mutation must not
 157     /// change the sort-order of keys.
 158     #[inline]
 159     pub fn offset_keys<F>(&mut self, f: F)
 160         where F: Fn(&mut K)
 161     {
 162         self.data.iter_mut().map(|&mut (ref mut k, _)| k).for_each(f);
 163     }
 164
 165     /// Inserts a presorted range of elements into the map. If the range can be
 166     /// inserted as a whole in between to existing elements of the map, this
 167     /// will be faster than inserting the elements individually.
 168     ///
 169     /// It is up to the caller to make sure that the elements are sorted by key
 170     /// and that there are no duplicates.
 171     #[inline]
 172     pub fn insert_presorted(&mut self, mut elements: Vec<(K, V)>) {
 173         if elements.is_empty() {
 174             return
 175         }
 176
 177         debug_assert!(elements.windows(2).all(|w| w[0].0 < w[1].0));
 178
 179         let start_index = self.lookup_index_for(&elements[0].0);
 180
 181         let drain = match start_index {
 182             Ok(index) => {
 183                 let mut drain = elements.drain(..);
 184                 self.data[index] = drain.next().unwrap();
 185                 drain
 186             }
 187             Err(index) => {
 188                 if index == self.data.len() ||
 189                    elements.last().unwrap().0 < self.data[index].0 {
 190                     // We can copy the whole range without having to mix with
 191                     // existing elements.
 192                     self.data.splice(index .. index, elements.drain(..));
 193                     return
 194                 }
 195
 196                 let mut drain = elements.drain(..);
 197                 self.data.insert(index, drain.next().unwrap());
 198                 drain
 199             }
 200         };
 201
 202         // Insert the rest
 203         for (k, v) in drain {
 204             self.insert(k, v);
 205         }
 206     }
 207
 208     /// Looks up the key in `self.data` via `slice::binary_search()`.
 209     #[inline(always)]
 210     fn lookup_index_for<Q>(&self, key: &Q) -> Result<usize, usize>
 211         where K: Borrow<Q>,
 212               Q: Ord + ?Sized
 213     {
 214         self.data.binary_search_by(|&(ref x, _)| x.borrow().cmp(key))
 215     }
 216
 217     #[inline]
 218     fn range_slice_indices<R>(&self, range: R) -> (usize, usize)
 219         where R: RangeBounds<K>
 220     {
 221         let start = match range.start_bound() {
 222             Bound::Included(ref k) => {
 223                 match self.lookup_index_for(k) {
 224                     Ok(index) | Err(index) => index
 225                 }
 226             }
 227             Bound::Excluded(ref k) => {
 228                 match self.lookup_index_for(k) {
 229                     Ok(index) => index + 1,
 230                     Err(index) => index,
 231                 }
 232             }
 233             Bound::Unbounded => 0,
 234         };
 235
 236         let end = match range.end_bound() {
 237             Bound::Included(ref k) => {
 238                 match self.lookup_index_for(k) {
 239                     Ok(index) => index + 1,
 240                     Err(index) => index,
 241                 }
 242             }
 243             Bound::Excluded(ref k) => {
 244                 match self.lookup_index_for(k) {
 245                     Ok(index) | Err(index) => index,
 246                 }
 247             }
 248             Bound::Unbounded => self.data.len(),
 249         };
 250
 251         (start, end)
 252     }
 253
 254     #[inline]
 255     pub fn contains_key<Q>(&self, key: &Q) -> bool
 256         where K: Borrow<Q>,
 257               Q: Ord + ?Sized
 258     {
 259         self.get(key).is_some()
 260     }
 261 }
 262
 263 impl<K: Ord, V> IntoIterator for SortedMap<K, V> {
 264     type Item = (K, V);
 265     type IntoIter = ::std::vec::IntoIter<(K, V)>;
 266
 267     fn into_iter(self) -> Self::IntoIter {
 268         self.data.into_iter()
 269     }
 270 }
 271
 272 impl<'a, K, Q, V> Index<&'a Q> for SortedMap<K, V>
 273     where K: Ord + Borrow<Q>,
 274           Q: Ord + ?Sized
 275 {
 276     type Output = V;
 277
 278     fn index(&self, key: &Q) -> &Self::Output {
 279         self.get(key).expect("no entry found for key")
 280     }
 281 }
 282
 283 impl<'a, K, Q, V> IndexMut<&'a Q> for SortedMap<K, V>
 284     where K: Ord + Borrow<Q>,
 285           Q: Ord + ?Sized
 286 {
 287     fn index_mut(&mut self, key: &Q) -> &mut Self::Output {
 288         self.get_mut(key).expect("no entry found for key")
 289     }
 290 }
 291
 292 impl<K: Ord, V> FromIterator<(K, V)> for SortedMap<K, V> {
 293     fn from_iter<T: IntoIterator<Item = (K, V)>>(iter: T) -> Self {
 294         let mut data: Vec<(K, V)> = iter.into_iter().collect();
 295
 296         data.sort_unstable_by(|&(ref k1, _), &(ref k2, _)| k1.cmp(k2));
 297         data.dedup_by(|&mut (ref k1, _), &mut (ref k2, _)| {
 298             k1.cmp(k2) == Ordering::Equal
 299         });
 300
 301         SortedMap {
 302             data
 303         }
 304     }
 305 }
 306
 307 #[cfg(test)]
 308 mod tests {
 309     use super::SortedMap;
 310
 311     #[test]
 312     fn test_insert_and_iter() {
 313         let mut map = SortedMap::new();
 314         let mut expected = Vec::new();
 315
 316         for x in 0 .. 100 {
 317             assert_eq!(map.iter().cloned().collect::<Vec<_>>(), expected);
 318
 319             let x = 1000 - x * 2;
 320             map.insert(x, x);
 321             expected.insert(0, (x, x));
 322         }
 323     }
 324
 325     #[test]
 326     fn test_get_and_index() {
 327         let mut map = SortedMap::new();
 328         let mut expected = Vec::new();
 329
 330         for x in 0 .. 100 {
 331             let x = 1000 - x;
 332             if x & 1 == 0 {
 333                 map.insert(x, x);
 334             }
 335             expected.push(x);
 336         }
 337
 338         for mut x in expected {
 339             if x & 1 == 0 {
 340                 assert_eq!(map.get(&x), Some(&x));
 341                 assert_eq!(map.get_mut(&x), Some(&mut x));
 342                 assert_eq!(map[&x], x);
 343                 assert_eq!(&mut map[&x], &mut x);
 344             } else {
 345                 assert_eq!(map.get(&x), None);
 346                 assert_eq!(map.get_mut(&x), None);
 347             }
 348         }
 349     }
 350
 351     #[test]
 352     fn test_range() {
 353         let mut map = SortedMap::new();
 354         map.insert(1, 1);
 355         map.insert(3, 3);
 356         map.insert(6, 6);
 357         map.insert(9, 9);
 358
 359         let keys = |s: &[(_, _)]| {
 360             s.into_iter().map(|e| e.0).collect::<Vec<u32>>()
 361         };
 362
 363         for start in 0 .. 11 {
 364             for end in 0 .. 11 {
 365                 if end < start {
 366                     continue
 367                 }
 368
 369                 let mut expected = vec![1, 3, 6, 9];
 370                 expected.retain(|&x| x >= start && x < end);
 371
 372                 assert_eq!(keys(map.range(start..end)), expected, "range = {}..{}", start, end);
 373             }
 374         }
 375     }
 376
 377
 378     #[test]
 379     fn test_offset_keys() {
 380         let mut map = SortedMap::new();
 381         map.insert(1, 1);
 382         map.insert(3, 3);
 383         map.insert(6, 6);
 384
 385         map.offset_keys(|k| *k += 1);
 386
 387         let mut expected = SortedMap::new();
 388         expected.insert(2, 1);
 389         expected.insert(4, 3);
 390         expected.insert(7, 6);
 391
 392         assert_eq!(map, expected);
 393     }
 394
 395     fn keys(s: SortedMap<u32, u32>) -> Vec<u32> {
 396         s.into_iter().map(|(k, _)| k).collect::<Vec<u32>>()
 397     }
 398
 399     fn elements(s: SortedMap<u32, u32>) -> Vec<(u32, u32)> {
 400         s.into_iter().collect::<Vec<(u32, u32)>>()
 401     }
 402
 403     #[test]
 404     fn test_remove_range() {
 405         let mut map = SortedMap::new();
 406         map.insert(1, 1);
 407         map.insert(3, 3);
 408         map.insert(6, 6);
 409         map.insert(9, 9);
 410
 411         for start in 0 .. 11 {
 412             for end in 0 .. 11 {
 413                 if end < start {
 414                     continue
 415                 }
 416
 417                 let mut expected = vec![1, 3, 6, 9];
 418                 expected.retain(|&x| x < start || x >= end);
 419
 420                 let mut map = map.clone();
 421                 map.remove_range(start .. end);
 422
 423                 assert_eq!(keys(map), expected, "range = {}..{}", start, end);
 424             }
 425         }
 426     }
 427
 428     #[test]
 429     fn test_remove() {
 430         let mut map = SortedMap::new();
 431         let mut expected = Vec::new();
 432
 433         for x in 0..10 {
 434             map.insert(x, x);
 435             expected.push((x, x));
 436         }
 437
 438         for x in 0 .. 10 {
 439             let mut map = map.clone();
 440             let mut expected = expected.clone();
 441
 442             assert_eq!(map.remove(&x), Some(x));
 443             expected.remove(x as usize);
 444
 445             assert_eq!(map.iter().cloned().collect::<Vec<_>>(), expected);
 446         }
 447     }
 448
 449     #[test]
 450     fn test_insert_presorted_non_overlapping() {
 451         let mut map = SortedMap::new();
 452         map.insert(2, 0);
 453         map.insert(8, 0);
 454
 455         map.insert_presorted(vec![(3, 0), (7, 0)]);
 456
 457         let expected = vec![2, 3, 7, 8];
 458         assert_eq!(keys(map), expected);
 459     }
 460
 461     #[test]
 462     fn test_insert_presorted_first_elem_equal() {
 463         let mut map = SortedMap::new();
 464         map.insert(2, 2);
 465         map.insert(8, 8);
 466
 467         map.insert_presorted(vec![(2, 0), (7, 7)]);
 468
 469         let expected = vec![(2, 0), (7, 7), (8, 8)];
 470         assert_eq!(elements(map), expected);
 471     }
 472
 473     #[test]
 474     fn test_insert_presorted_last_elem_equal() {
 475         let mut map = SortedMap::new();
 476         map.insert(2, 2);
 477         map.insert(8, 8);
 478
 479         map.insert_presorted(vec![(3, 3), (8, 0)]);
 480
 481         let expected = vec![(2, 2), (3, 3), (8, 0)];
 482         assert_eq!(elements(map), expected);
 483     }
 484
 485     #[test]
 486     fn test_insert_presorted_shuffle() {
 487         let mut map = SortedMap::new();
 488         map.insert(2, 2);
 489         map.insert(7, 7);
 490
 491         map.insert_presorted(vec![(1, 1), (3, 3), (8, 8)]);
 492
 493         let expected = vec![(1, 1), (2, 2), (3, 3), (7, 7), (8, 8)];
 494         assert_eq!(elements(map), expected);
 495     }
 496
 497     #[test]
 498     fn test_insert_presorted_at_end() {
 499         let mut map = SortedMap::new();
 500         map.insert(1, 1);
 501         map.insert(2, 2);
 502
 503         map.insert_presorted(vec![(3, 3), (8, 8)]);
 504
 505         let expected = vec![(1, 1), (2, 2), (3, 3), (8, 8)];
 506         assert_eq!(elements(map), expected);
 507     }
 508 }