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.
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.
11 use std::borrow::Borrow;
12 use std::cmp::Ordering;
13 use std::iter::FromIterator;
15 use std::ops::{RangeBounds, Bound, Index, IndexMut};
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.
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,
27 pub struct SortedMap<K: Ord, V> {
31 impl<K: Ord, V> SortedMap<K, V> {
33 pub fn new() -> SortedMap<K, V> {
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.
42 /// It is up to the caller to make sure that the elements are sorted by key
43 /// and that there are no duplicates.
45 pub fn from_presorted_elements(elements: Vec<(K, V)>) -> SortedMap<K, V>
47 debug_assert!(elements.windows(2).all(|w| w[0].0 < w[1].0));
55 pub fn insert(&mut self, key: K, mut value: V) -> Option<V> {
56 match self.lookup_index_for(&key) {
59 self.data.get_unchecked_mut(index)
61 mem::swap(&mut slot.1, &mut value);
65 self.data.insert(index, (key, value));
72 pub fn remove(&mut self, key: &K) -> Option<V> {
73 match self.lookup_index_for(key) {
75 Some(self.data.remove(index).1)
84 pub fn get<Q>(&self, key: &Q) -> Option<&V>
88 match self.lookup_index_for(key) {
91 Some(&self.data.get_unchecked(index).1)
101 pub fn get_mut<Q>(&mut self, key: &Q) -> Option<&mut V>
105 match self.lookup_index_for(key) {
108 Some(&mut self.data.get_unchecked_mut(index).1)
118 pub fn clear(&mut self) {
122 /// Iterate over elements, sorted by key
124 pub fn iter(&self) -> ::std::slice::Iter<(K, V)> {
128 /// Iterate over the keys, sorted
130 pub fn keys(&self) -> impl Iterator<Item = &K> + ExactSizeIterator {
131 self.data.iter().map(|&(ref k, _)| k)
134 /// Iterate over values, sorted by key
136 pub fn values(&self) -> impl Iterator<Item = &V> + ExactSizeIterator {
137 self.data.iter().map(|&(_, ref v)| v)
141 pub fn len(&self) -> usize {
146 pub fn is_empty(&self) -> bool {
151 pub fn range<R>(&self, range: R) -> &[(K, V)]
152 where R: RangeBounds<K>
154 let (start, end) = self.range_slice_indices(range);
155 (&self.data[start .. end])
159 pub fn remove_range<R>(&mut self, range: R)
160 where R: RangeBounds<K>
162 let (start, end) = self.range_slice_indices(range);
163 self.data.splice(start .. end, ::std::iter::empty());
166 /// Mutate all keys with the given function `f`. This mutation must not
167 /// change the sort-order of keys.
169 pub fn offset_keys<F>(&mut self, f: F)
172 self.data.iter_mut().map(|&mut (ref mut k, _)| k).for_each(f);
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.
179 /// It is up to the caller to make sure that the elements are sorted by key
180 /// and that there are no duplicates.
182 pub fn insert_presorted(&mut self, mut elements: Vec<(K, V)>) {
183 if elements.is_empty() {
187 debug_assert!(elements.windows(2).all(|w| w[0].0 < w[1].0));
189 let start_index = self.lookup_index_for(&elements[0].0);
191 let drain = match start_index {
193 let mut drain = elements.drain(..);
194 self.data[index] = drain.next().unwrap();
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(..));
206 let mut drain = elements.drain(..);
207 self.data.insert(index, drain.next().unwrap());
213 for (k, v) in drain {
218 /// Looks up the key in `self.data` via `slice::binary_search()`.
220 fn lookup_index_for<Q>(&self, key: &Q) -> Result<usize, usize>
224 self.data.binary_search_by(|&(ref x, _)| x.borrow().cmp(key))
228 fn range_slice_indices<R>(&self, range: R) -> (usize, usize)
229 where R: RangeBounds<K>
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
237 Bound::Excluded(ref k) => {
238 match self.lookup_index_for(k) {
239 Ok(index) => index + 1,
243 Bound::Unbounded => 0,
246 let end = match range.end_bound() {
247 Bound::Included(ref k) => {
248 match self.lookup_index_for(k) {
249 Ok(index) => index + 1,
253 Bound::Excluded(ref k) => {
254 match self.lookup_index_for(k) {
255 Ok(index) | Err(index) => index,
258 Bound::Unbounded => self.data.len(),
265 pub fn contains_key<Q>(&self, key: &Q) -> bool
269 self.get(key).is_some()
273 impl<K: Ord, V> IntoIterator for SortedMap<K, V> {
275 type IntoIter = ::std::vec::IntoIter<(K, V)>;
277 fn into_iter(self) -> Self::IntoIter {
278 self.data.into_iter()
282 impl<'a, K, Q, V> Index<&'a Q> for SortedMap<K, V>
283 where K: Ord + Borrow<Q>,
288 fn index(&self, key: &Q) -> &Self::Output {
289 self.get(key).expect("no entry found for key")
293 impl<'a, K, Q, V> IndexMut<&'a Q> for SortedMap<K, V>
294 where K: Ord + Borrow<Q>,
297 fn index_mut(&mut self, key: &Q) -> &mut Self::Output {
298 self.get_mut(key).expect("no entry found for key")
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();
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
319 use super::SortedMap;
322 fn test_insert_and_iter() {
323 let mut map = SortedMap::new();
324 let mut expected = Vec::new();
327 assert_eq!(map.iter().cloned().collect::<Vec<_>>(), expected);
329 let x = 1000 - x * 2;
331 expected.insert(0, (x, x));
336 fn test_get_and_index() {
337 let mut map = SortedMap::new();
338 let mut expected = Vec::new();
348 for mut x in expected {
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);
355 assert_eq!(map.get(&x), None);
356 assert_eq!(map.get_mut(&x), None);
363 let mut map = SortedMap::new();
369 let keys = |s: &[(_, _)]| {
370 s.into_iter().map(|e| e.0).collect::<Vec<u32>>()
373 for start in 0 .. 11 {
379 let mut expected = vec![1, 3, 6, 9];
380 expected.retain(|&x| x >= start && x < end);
382 assert_eq!(keys(map.range(start..end)), expected, "range = {}..{}", start, end);
389 fn test_offset_keys() {
390 let mut map = SortedMap::new();
395 map.offset_keys(|k| *k += 1);
397 let mut expected = SortedMap::new();
398 expected.insert(2, 1);
399 expected.insert(4, 3);
400 expected.insert(7, 6);
402 assert_eq!(map, expected);
405 fn keys(s: SortedMap<u32, u32>) -> Vec<u32> {
406 s.into_iter().map(|(k, _)| k).collect::<Vec<u32>>()
409 fn elements(s: SortedMap<u32, u32>) -> Vec<(u32, u32)> {
410 s.into_iter().collect::<Vec<(u32, u32)>>()
414 fn test_remove_range() {
415 let mut map = SortedMap::new();
421 for start in 0 .. 11 {
427 let mut expected = vec![1, 3, 6, 9];
428 expected.retain(|&x| x < start || x >= end);
430 let mut map = map.clone();
431 map.remove_range(start .. end);
433 assert_eq!(keys(map), expected, "range = {}..{}", start, end);
440 let mut map = SortedMap::new();
441 let mut expected = Vec::new();
445 expected.push((x, x));
449 let mut map = map.clone();
450 let mut expected = expected.clone();
452 assert_eq!(map.remove(&x), Some(x));
453 expected.remove(x as usize);
455 assert_eq!(map.iter().cloned().collect::<Vec<_>>(), expected);
460 fn test_insert_presorted_non_overlapping() {
461 let mut map = SortedMap::new();
465 map.insert_presorted(vec![(3, 0), (7, 0)]);
467 let expected = vec![2, 3, 7, 8];
468 assert_eq!(keys(map), expected);
472 fn test_insert_presorted_first_elem_equal() {
473 let mut map = SortedMap::new();
477 map.insert_presorted(vec![(2, 0), (7, 7)]);
479 let expected = vec![(2, 0), (7, 7), (8, 8)];
480 assert_eq!(elements(map), expected);
484 fn test_insert_presorted_last_elem_equal() {
485 let mut map = SortedMap::new();
489 map.insert_presorted(vec![(3, 3), (8, 0)]);
491 let expected = vec![(2, 2), (3, 3), (8, 0)];
492 assert_eq!(elements(map), expected);
496 fn test_insert_presorted_shuffle() {
497 let mut map = SortedMap::new();
501 map.insert_presorted(vec![(1, 1), (3, 3), (8, 8)]);
503 let expected = vec![(1, 1), (2, 2), (3, 3), (7, 7), (8, 8)];
504 assert_eq!(elements(map), expected);
508 fn test_insert_presorted_at_end() {
509 let mut map = SortedMap::new();
513 map.insert_presorted(vec![(3, 3), (8, 8)]);
515 let expected = vec![(1, 1), (2, 2), (3, 3), (8, 8)];
516 assert_eq!(elements(map), expected);