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