1 use std::borrow::Borrow;
2 use std::cmp::Ordering;
3 use std::iter::FromIterator;
5 use std::ops::{RangeBounds, Bound, Index, IndexMut};
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.
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,
17 pub struct SortedMap<K: Ord, V> {
21 impl<K: Ord, V> SortedMap<K, V> {
23 pub fn new() -> SortedMap<K, V> {
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.
32 /// It is up to the caller to make sure that the elements are sorted by key
33 /// and that there are no duplicates.
35 pub fn from_presorted_elements(elements: Vec<(K, V)>) -> SortedMap<K, V>
37 debug_assert!(elements.windows(2).all(|w| w[0].0 < w[1].0));
45 pub fn insert(&mut self, key: K, mut value: V) -> Option<V> {
46 match self.lookup_index_for(&key) {
49 self.data.get_unchecked_mut(index)
51 mem::swap(&mut slot.1, &mut value);
55 self.data.insert(index, (key, value));
62 pub fn remove(&mut self, key: &K) -> Option<V> {
63 match self.lookup_index_for(key) {
65 Some(self.data.remove(index).1)
74 pub fn get<Q>(&self, key: &Q) -> Option<&V>
78 match self.lookup_index_for(key) {
81 Some(&self.data.get_unchecked(index).1)
91 pub fn get_mut<Q>(&mut self, key: &Q) -> Option<&mut V>
95 match self.lookup_index_for(key) {
98 Some(&mut self.data.get_unchecked_mut(index).1)
108 pub fn clear(&mut self) {
112 /// Iterate over elements, sorted by key
114 pub fn iter(&self) -> ::std::slice::Iter<'_, (K, V)> {
118 /// Iterate over the keys, sorted
120 pub fn keys(&self) -> impl Iterator<Item = &K> + ExactSizeIterator {
121 self.data.iter().map(|&(ref k, _)| k)
124 /// Iterate over values, sorted by key
126 pub fn values(&self) -> impl Iterator<Item = &V> + ExactSizeIterator {
127 self.data.iter().map(|&(_, ref v)| v)
131 pub fn len(&self) -> usize {
136 pub fn is_empty(&self) -> bool {
141 pub fn range<R>(&self, range: R) -> &[(K, V)]
142 where R: RangeBounds<K>
144 let (start, end) = self.range_slice_indices(range);
145 (&self.data[start .. end])
149 pub fn remove_range<R>(&mut self, range: R)
150 where R: RangeBounds<K>
152 let (start, end) = self.range_slice_indices(range);
153 self.data.splice(start .. end, ::std::iter::empty());
156 /// Mutate all keys with the given function `f`. This mutation must not
157 /// change the sort-order of keys.
159 pub fn offset_keys<F>(&mut self, f: F)
162 self.data.iter_mut().map(|&mut (ref mut k, _)| k).for_each(f);
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.
169 /// It is up to the caller to make sure that the elements are sorted by key
170 /// and that there are no duplicates.
172 pub fn insert_presorted(&mut self, mut elements: Vec<(K, V)>) {
173 if elements.is_empty() {
177 debug_assert!(elements.windows(2).all(|w| w[0].0 < w[1].0));
179 let start_index = self.lookup_index_for(&elements[0].0);
181 let drain = match start_index {
183 let mut drain = elements.drain(..);
184 self.data[index] = drain.next().unwrap();
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(..));
196 let mut drain = elements.drain(..);
197 self.data.insert(index, drain.next().unwrap());
203 for (k, v) in drain {
208 /// Looks up the key in `self.data` via `slice::binary_search()`.
210 fn lookup_index_for<Q>(&self, key: &Q) -> Result<usize, usize>
214 self.data.binary_search_by(|&(ref x, _)| x.borrow().cmp(key))
218 fn range_slice_indices<R>(&self, range: R) -> (usize, usize)
219 where R: RangeBounds<K>
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
227 Bound::Excluded(ref k) => {
228 match self.lookup_index_for(k) {
229 Ok(index) => index + 1,
233 Bound::Unbounded => 0,
236 let end = match range.end_bound() {
237 Bound::Included(ref k) => {
238 match self.lookup_index_for(k) {
239 Ok(index) => index + 1,
243 Bound::Excluded(ref k) => {
244 match self.lookup_index_for(k) {
245 Ok(index) | Err(index) => index,
248 Bound::Unbounded => self.data.len(),
255 pub fn contains_key<Q>(&self, key: &Q) -> bool
259 self.get(key).is_some()
263 impl<K: Ord, V> IntoIterator for SortedMap<K, V> {
265 type IntoIter = ::std::vec::IntoIter<(K, V)>;
267 fn into_iter(self) -> Self::IntoIter {
268 self.data.into_iter()
272 impl<'a, K, Q, V> Index<&'a Q> for SortedMap<K, V>
273 where K: Ord + Borrow<Q>,
278 fn index(&self, key: &Q) -> &Self::Output {
279 self.get(key).expect("no entry found for key")
283 impl<'a, K, Q, V> IndexMut<&'a Q> for SortedMap<K, V>
284 where K: Ord + Borrow<Q>,
287 fn index_mut(&mut self, key: &Q) -> &mut Self::Output {
288 self.get_mut(key).expect("no entry found for key")
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();
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
309 use super::SortedMap;
312 fn test_insert_and_iter() {
313 let mut map = SortedMap::new();
314 let mut expected = Vec::new();
317 assert_eq!(map.iter().cloned().collect::<Vec<_>>(), expected);
319 let x = 1000 - x * 2;
321 expected.insert(0, (x, x));
326 fn test_get_and_index() {
327 let mut map = SortedMap::new();
328 let mut expected = Vec::new();
338 for mut x in expected {
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);
345 assert_eq!(map.get(&x), None);
346 assert_eq!(map.get_mut(&x), None);
353 let mut map = SortedMap::new();
359 let keys = |s: &[(_, _)]| {
360 s.into_iter().map(|e| e.0).collect::<Vec<u32>>()
363 for start in 0 .. 11 {
369 let mut expected = vec![1, 3, 6, 9];
370 expected.retain(|&x| x >= start && x < end);
372 assert_eq!(keys(map.range(start..end)), expected, "range = {}..{}", start, end);
379 fn test_offset_keys() {
380 let mut map = SortedMap::new();
385 map.offset_keys(|k| *k += 1);
387 let mut expected = SortedMap::new();
388 expected.insert(2, 1);
389 expected.insert(4, 3);
390 expected.insert(7, 6);
392 assert_eq!(map, expected);
395 fn keys(s: SortedMap<u32, u32>) -> Vec<u32> {
396 s.into_iter().map(|(k, _)| k).collect::<Vec<u32>>()
399 fn elements(s: SortedMap<u32, u32>) -> Vec<(u32, u32)> {
400 s.into_iter().collect::<Vec<(u32, u32)>>()
404 fn test_remove_range() {
405 let mut map = SortedMap::new();
411 for start in 0 .. 11 {
417 let mut expected = vec![1, 3, 6, 9];
418 expected.retain(|&x| x < start || x >= end);
420 let mut map = map.clone();
421 map.remove_range(start .. end);
423 assert_eq!(keys(map), expected, "range = {}..{}", start, end);
430 let mut map = SortedMap::new();
431 let mut expected = Vec::new();
435 expected.push((x, x));
439 let mut map = map.clone();
440 let mut expected = expected.clone();
442 assert_eq!(map.remove(&x), Some(x));
443 expected.remove(x as usize);
445 assert_eq!(map.iter().cloned().collect::<Vec<_>>(), expected);
450 fn test_insert_presorted_non_overlapping() {
451 let mut map = SortedMap::new();
455 map.insert_presorted(vec![(3, 0), (7, 0)]);
457 let expected = vec![2, 3, 7, 8];
458 assert_eq!(keys(map), expected);
462 fn test_insert_presorted_first_elem_equal() {
463 let mut map = SortedMap::new();
467 map.insert_presorted(vec![(2, 0), (7, 7)]);
469 let expected = vec![(2, 0), (7, 7), (8, 8)];
470 assert_eq!(elements(map), expected);
474 fn test_insert_presorted_last_elem_equal() {
475 let mut map = SortedMap::new();
479 map.insert_presorted(vec![(3, 3), (8, 0)]);
481 let expected = vec![(2, 2), (3, 3), (8, 0)];
482 assert_eq!(elements(map), expected);
486 fn test_insert_presorted_shuffle() {
487 let mut map = SortedMap::new();
491 map.insert_presorted(vec![(1, 1), (3, 3), (8, 8)]);
493 let expected = vec![(1, 1), (2, 2), (3, 3), (7, 7), (8, 8)];
494 assert_eq!(elements(map), expected);
498 fn test_insert_presorted_at_end() {
499 let mut map = SortedMap::new();
503 map.insert_presorted(vec![(3, 3), (8, 8)]);
505 let expected = vec![(1, 1), (2, 2), (3, 3), (8, 8)];
506 assert_eq!(elements(map), expected);