src/librustc_data_structures/sorted_map.rs

   1 // Copyright 2018 The Rust Project Developers. See the COPYRIGHT
   2 // file at the top-level directory of this distribution and at
   3 // http://rust-lang.org/COPYRIGHT.
   4 //
   5 // Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
   6 // http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
   7 // <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
   8 // option. This file may not be copied, modified, or distributed
   9 // except according to those terms.
  10
  11 use std::borrow::Borrow;
  12 use std::cmp::Ordering;
  13 use std::convert::From;
  14 use std::mem;
  15 use std::ops::{RangeBounds, Bound, Index, IndexMut};
  16
  17 /// `SortedMap` is a data structure with similar characteristics as BTreeMap but
  18 /// slightly different trade-offs: lookup, inseration, and removal are O(log(N))
  19 /// and elements can be iterated in order cheaply.
  20 ///
  21 /// `SortedMap` can be faster than a `BTreeMap` for small sizes (<50) since it
  22 /// stores data in a more compact way. It also supports accessing contiguous
  23 /// ranges of elements as a slice, and slices of already sorted elements can be
  24 /// inserted efficiently.
  25 #[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Default, Debug, RustcEncodable,
  26          RustcDecodable)]
  27 pub struct SortedMap<K: Ord, V> {
  28     data: Vec<(K,V)>
  29 }
  30
  31 impl<K: Ord, V> SortedMap<K, V> {
  32
  33     #[inline]
  34     pub fn new() -> SortedMap<K, V> {
  35         SortedMap {
  36             data: vec![]
  37         }
  38     }
  39
  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     {
  48         debug_assert!(elements.windows(2).all(|w| w[0].0 < w[1].0));
  49
  50         SortedMap {
  51             data: elements
  52         }
  53     }
  54
  55     #[inline]
  56     pub fn insert(&mut self, key: K, mut value: V) -> Option<V> {
  57         match self.lookup_index_for(&key) {
  58             Ok(index) => {
  59                 let slot = unsafe {
  60                     self.data.get_unchecked_mut(index)
  61                 };
  62                 mem::swap(&mut slot.1, &mut value);
  63                 Some(value)
  64             }
  65             Err(index) => {
  66                 self.data.insert(index, (key, value));
  67                 None
  68             }
  69         }
  70     }
  71
  72     #[inline]
  73     pub fn remove(&mut self, key: &K) -> Option<V> {
  74         match self.lookup_index_for(key) {
  75             Ok(index) => {
  76                 Some(self.data.remove(index).1)
  77             }
  78             Err(_) => {
  79                 None
  80             }
  81         }
  82     }
  83
  84     #[inline]
  85     pub fn get(&self, key: &K) -> Option<&V> {
  86         match self.lookup_index_for(key) {
  87             Ok(index) => {
  88                 unsafe {
  89                     Some(&self.data.get_unchecked(index).1)
  90                 }
  91             }
  92             Err(_) => {
  93                 None
  94             }
  95         }
  96     }
  97
  98     #[inline]
  99     pub fn get_mut(&mut self, key: &K) -> Option<&mut V> {
 100         match self.lookup_index_for(key) {
 101             Ok(index) => {
 102                 unsafe {
 103                     Some(&mut self.data.get_unchecked_mut(index).1)
 104                 }
 105             }
 106             Err(_) => {
 107                 None
 108             }
 109         }
 110     }
 111
 112     #[inline]
 113     pub fn clear(&mut self) {
 114         self.data.clear();
 115     }
 116
 117     /// Iterate over elements, sorted by key
 118     #[inline]
 119     pub fn iter(&self) -> ::std::slice::Iter<(K, V)> {
 120         self.data.iter()
 121     }
 122
 123     /// Iterate over the keys, sorted
 124     #[inline]
 125     pub fn keys(&self) -> impl Iterator<Item=&K> + ExactSizeIterator {
 126         self.data.iter().map(|&(ref k, _)| k)
 127     }
 128
 129     /// Iterate over values, sorted by key
 130     #[inline]
 131     pub fn values(&self) -> impl Iterator<Item=&V> + ExactSizeIterator {
 132         self.data.iter().map(|&(_, ref v)| v)
 133     }
 134
 135     #[inline]
 136     pub fn len(&self) -> usize {
 137         self.data.len()
 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(&self, key: &K) -> Result<usize, usize> {
 211         self.data.binary_search_by(|&(ref x, _)| x.cmp(key))
 212     }
 213
 214     #[inline]
 215     fn range_slice_indices<R>(&self, range: R) -> (usize, usize)
 216         where R: RangeBounds<K>
 217     {
 218         let start = match range.start_bound() {
 219             Bound::Included(ref k) => {
 220                 match self.lookup_index_for(k) {
 221                     Ok(index) | Err(index) => index
 222                 }
 223             }
 224             Bound::Excluded(ref k) => {
 225                 match self.lookup_index_for(k) {
 226                     Ok(index) => index + 1,
 227                     Err(index) => index,
 228                 }
 229             }
 230             Bound::Unbounded => 0,
 231         };
 232
 233         let end = match range.end_bound() {
 234             Bound::Included(ref k) => {
 235                 match self.lookup_index_for(k) {
 236                     Ok(index) => index + 1,
 237                     Err(index) => index,
 238                 }
 239             }
 240             Bound::Excluded(ref k) => {
 241                 match self.lookup_index_for(k) {
 242                     Ok(index) | Err(index) => index,
 243                 }
 244             }
 245             Bound::Unbounded => self.data.len(),
 246         };
 247
 248         (start, end)
 249     }
 250 }
 251
 252 impl<K: Ord, V> IntoIterator for SortedMap<K, V> {
 253     type Item = (K, V);
 254     type IntoIter = ::std::vec::IntoIter<(K, V)>;
 255     fn into_iter(self) -> Self::IntoIter {
 256         self.data.into_iter()
 257     }
 258 }
 259
 260 impl<K: Ord, V, Q: Borrow<K>> Index<Q> for SortedMap<K, V> {
 261     type Output = V;
 262     fn index(&self, index: Q) -> &Self::Output {
 263         let k: &K = index.borrow();
 264         self.get(k).unwrap()
 265     }
 266 }
 267
 268 impl<K: Ord, V, Q: Borrow<K>> IndexMut<Q> for SortedMap<K, V> {
 269     fn index_mut(&mut self, index: Q) -> &mut Self::Output {
 270         let k: &K = index.borrow();
 271         self.get_mut(k).unwrap()
 272     }
 273 }
 274
 275 impl<K: Ord, V, I: Iterator<Item=(K, V)>> From<I> for SortedMap<K, V> {
 276     fn from(data: I) -> Self {
 277         let mut data: Vec<(K, V)> = data.collect();
 278         data.sort_unstable_by(|&(ref k1, _), &(ref k2, _)| k1.cmp(k2));
 279         data.dedup_by(|&mut (ref k1, _), &mut (ref k2, _)| {
 280             k1.cmp(k2) == Ordering::Equal
 281         });
 282         SortedMap {
 283             data
 284         }
 285     }
 286 }
 287
 288 #[cfg(test)]
 289 mod tests {
 290     use super::SortedMap;
 291
 292     #[test]
 293     fn test_insert_and_iter() {
 294         let mut map = SortedMap::new();
 295         let mut expected = Vec::new();
 296
 297         for x in 0 .. 100 {
 298             assert_eq!(map.iter().cloned().collect::<Vec<_>>(), expected);
 299
 300             let x = 1000 - x * 2;
 301             map.insert(x, x);
 302             expected.insert(0, (x, x));
 303         }
 304     }
 305
 306     #[test]
 307     fn test_get_and_index() {
 308         let mut map = SortedMap::new();
 309         let mut expected = Vec::new();
 310
 311         for x in 0 .. 100 {
 312             let x = 1000 - x;
 313             if x & 1 == 0 {
 314                 map.insert(x, x);
 315             }
 316             expected.push(x);
 317         }
 318
 319         for mut x in expected {
 320             if x & 1 == 0 {
 321                 assert_eq!(map.get(&x), Some(&x));
 322                 assert_eq!(map.get_mut(&x), Some(&mut x));
 323                 assert_eq!(map[&x], x);
 324                 assert_eq!(&mut map[&x], &mut x);
 325             } else {
 326                 assert_eq!(map.get(&x), None);
 327                 assert_eq!(map.get_mut(&x), None);
 328             }
 329         }
 330     }
 331
 332     #[test]
 333     fn test_range() {
 334         let mut map = SortedMap::new();
 335         map.insert(1, 1);
 336         map.insert(3, 3);
 337         map.insert(6, 6);
 338         map.insert(9, 9);
 339
 340         let keys = |s: &[(_, _)]| {
 341             s.into_iter().map(|e| e.0).collect::<Vec<u32>>()
 342         };
 343
 344         for start in 0 .. 11 {
 345             for end in 0 .. 11 {
 346                 if end < start {
 347                     continue
 348                 }
 349
 350                 let mut expected = vec![1, 3, 6, 9];
 351                 expected.retain(|&x| x >= start && x < end);
 352
 353                 assert_eq!(keys(map.range(start..end)), expected, "range = {}..{}", start, end);
 354             }
 355         }
 356     }
 357
 358
 359     #[test]
 360     fn test_offset_keys() {
 361         let mut map = SortedMap::new();
 362         map.insert(1, 1);
 363         map.insert(3, 3);
 364         map.insert(6, 6);
 365
 366         map.offset_keys(|k| *k += 1);
 367
 368         let mut expected = SortedMap::new();
 369         expected.insert(2, 1);
 370         expected.insert(4, 3);
 371         expected.insert(7, 6);
 372
 373         assert_eq!(map, expected);
 374     }
 375
 376     fn keys(s: SortedMap<u32, u32>) -> Vec<u32> {
 377         s.into_iter().map(|(k, _)| k).collect::<Vec<u32>>()
 378     }
 379
 380     fn elements(s: SortedMap<u32, u32>) -> Vec<(u32, u32)> {
 381         s.into_iter().collect::<Vec<(u32, u32)>>()
 382     }
 383
 384     #[test]
 385     fn test_remove_range() {
 386         let mut map = SortedMap::new();
 387         map.insert(1, 1);
 388         map.insert(3, 3);
 389         map.insert(6, 6);
 390         map.insert(9, 9);
 391
 392         for start in 0 .. 11 {
 393             for end in 0 .. 11 {
 394                 if end < start {
 395                     continue
 396                 }
 397
 398                 let mut expected = vec![1, 3, 6, 9];
 399                 expected.retain(|&x| x < start || x >= end);
 400
 401                 let mut map = map.clone();
 402                 map.remove_range(start .. end);
 403
 404                 assert_eq!(keys(map), expected, "range = {}..{}", start, end);
 405             }
 406         }
 407     }
 408
 409     #[test]
 410     fn test_remove() {
 411         let mut map = SortedMap::new();
 412         let mut expected = Vec::new();
 413
 414         for x in 0..10 {
 415             map.insert(x, x);
 416             expected.push((x, x));
 417         }
 418
 419         for x in 0 .. 10 {
 420             let mut map = map.clone();
 421             let mut expected = expected.clone();
 422
 423             assert_eq!(map.remove(&x), Some(x));
 424             expected.remove(x as usize);
 425
 426             assert_eq!(map.iter().cloned().collect::<Vec<_>>(), expected);
 427         }
 428     }
 429
 430     #[test]
 431     fn test_insert_presorted_non_overlapping() {
 432         let mut map = SortedMap::new();
 433         map.insert(2, 0);
 434         map.insert(8, 0);
 435
 436         map.insert_presorted(vec![(3, 0), (7, 0)]);
 437
 438         let expected = vec![2, 3, 7, 8];
 439         assert_eq!(keys(map), expected);
 440     }
 441
 442     #[test]
 443     fn test_insert_presorted_first_elem_equal() {
 444         let mut map = SortedMap::new();
 445         map.insert(2, 2);
 446         map.insert(8, 8);
 447
 448         map.insert_presorted(vec![(2, 0), (7, 7)]);
 449
 450         let expected = vec![(2, 0), (7, 7), (8, 8)];
 451         assert_eq!(elements(map), expected);
 452     }
 453
 454     #[test]
 455     fn test_insert_presorted_last_elem_equal() {
 456         let mut map = SortedMap::new();
 457         map.insert(2, 2);
 458         map.insert(8, 8);
 459
 460         map.insert_presorted(vec![(3, 3), (8, 0)]);
 461
 462         let expected = vec![(2, 2), (3, 3), (8, 0)];
 463         assert_eq!(elements(map), expected);
 464     }
 465
 466     #[test]
 467     fn test_insert_presorted_shuffle() {
 468         let mut map = SortedMap::new();
 469         map.insert(2, 2);
 470         map.insert(7, 7);
 471
 472         map.insert_presorted(vec![(1, 1), (3, 3), (8, 8)]);
 473
 474         let expected = vec![(1, 1), (2, 2), (3, 3), (7, 7), (8, 8)];
 475         assert_eq!(elements(map), expected);
 476     }
 477
 478     #[test]
 479     fn test_insert_presorted_at_end() {
 480         let mut map = SortedMap::new();
 481         map.insert(1, 1);
 482         map.insert(2, 2);
 483
 484         map.insert_presorted(vec![(3, 3), (8, 8)]);
 485
 486         let expected = vec![(1, 1), (2, 2), (3, 3), (8, 8)];
 487         assert_eq!(elements(map), expected);
 488     }
 489 }