1 use std::borrow::Borrow;
2 use std::cmp::Ordering;
3 use std::iter::FromIterator;
5 use std::ops::{Bound, Index, IndexMut, RangeBounds};
9 pub use index_map::SortedIndexMultiMap;
11 /// `SortedMap` is a data structure with similar characteristics as BTreeMap but
12 /// slightly different trade-offs: lookup, insertion, and removal are O(log(N))
13 /// and elements can be iterated in order cheaply.
15 /// `SortedMap` can be faster than a `BTreeMap` for small sizes (<50) since it
16 /// stores data in a more compact way. It also supports accessing contiguous
17 /// ranges of elements as a slice, and slices of already sorted elements can be
18 /// inserted efficiently.
31 pub struct SortedMap<K: Ord, V> {
35 impl<K: Ord, V> SortedMap<K, V> {
37 pub fn new() -> SortedMap<K, V> {
38 SortedMap { data: vec![] }
41 /// Construct a `SortedMap` from a presorted set of elements. This is faster
42 /// than creating an empty map and then inserting the elements individually.
44 /// It is up to the caller to make sure that the elements are sorted by key
45 /// and that there are no duplicates.
47 pub fn from_presorted_elements(elements: Vec<(K, V)>) -> SortedMap<K, V> {
48 debug_assert!(elements.windows(2).all(|w| w[0].0 < w[1].0));
50 SortedMap { data: elements }
54 pub fn insert(&mut self, key: K, mut value: V) -> Option<V> {
55 match self.lookup_index_for(&key) {
57 let slot = unsafe { self.data.get_unchecked_mut(index) };
58 mem::swap(&mut slot.1, &mut value);
62 self.data.insert(index, (key, value));
69 pub fn remove(&mut self, key: &K) -> Option<V> {
70 match self.lookup_index_for(key) {
71 Ok(index) => Some(self.data.remove(index).1),
77 pub fn get<Q>(&self, key: &Q) -> Option<&V>
82 match self.lookup_index_for(key) {
83 Ok(index) => unsafe { Some(&self.data.get_unchecked(index).1) },
89 pub fn get_mut<Q>(&mut self, key: &Q) -> Option<&mut V>
94 match self.lookup_index_for(key) {
95 Ok(index) => unsafe { Some(&mut self.data.get_unchecked_mut(index).1) },
101 pub fn clear(&mut self) {
105 /// Iterate over elements, sorted by key
107 pub fn iter(&self) -> ::std::slice::Iter<'_, (K, V)> {
111 /// Iterate over the keys, sorted
113 pub fn keys(&self) -> impl Iterator<Item = &K> + ExactSizeIterator {
114 self.data.iter().map(|&(ref k, _)| k)
117 /// Iterate over values, sorted by key
119 pub fn values(&self) -> impl Iterator<Item = &V> + ExactSizeIterator {
120 self.data.iter().map(|&(_, ref v)| v)
124 pub fn len(&self) -> usize {
129 pub fn is_empty(&self) -> bool {
134 pub fn range<R>(&self, range: R) -> &[(K, V)]
138 let (start, end) = self.range_slice_indices(range);
139 &self.data[start..end]
143 pub fn remove_range<R>(&mut self, range: R)
147 let (start, end) = self.range_slice_indices(range);
148 self.data.splice(start..end, ::std::iter::empty());
151 /// Mutate all keys with the given function `f`. This mutation must not
152 /// change the sort-order of keys.
154 pub fn offset_keys<F>(&mut self, f: F)
158 self.data.iter_mut().map(|&mut (ref mut k, _)| k).for_each(f);
161 /// Inserts a presorted range of elements into the map. If the range can be
162 /// inserted as a whole in between to existing elements of the map, this
163 /// will be faster than inserting the elements individually.
165 /// It is up to the caller to make sure that the elements are sorted by key
166 /// and that there are no duplicates.
168 pub fn insert_presorted(&mut self, mut elements: Vec<(K, V)>) {
169 if elements.is_empty() {
173 debug_assert!(elements.windows(2).all(|w| w[0].0 < w[1].0));
175 let start_index = self.lookup_index_for(&elements[0].0);
177 let drain = match start_index {
179 let mut drain = elements.drain(..);
180 self.data[index] = drain.next().unwrap();
184 if index == self.data.len() || elements.last().unwrap().0 < self.data[index].0 {
185 // We can copy the whole range without having to mix with
186 // existing elements.
187 self.data.splice(index..index, elements.drain(..));
191 let mut drain = elements.drain(..);
192 self.data.insert(index, drain.next().unwrap());
198 for (k, v) in drain {
203 /// Looks up the key in `self.data` via `slice::binary_search()`.
205 fn lookup_index_for<Q>(&self, key: &Q) -> Result<usize, usize>
210 self.data.binary_search_by(|&(ref x, _)| x.borrow().cmp(key))
214 fn range_slice_indices<R>(&self, range: R) -> (usize, usize)
218 let start = match range.start_bound() {
219 Bound::Included(ref k) => match self.lookup_index_for(k) {
220 Ok(index) | Err(index) => index,
222 Bound::Excluded(ref k) => match self.lookup_index_for(k) {
223 Ok(index) => index + 1,
226 Bound::Unbounded => 0,
229 let end = match range.end_bound() {
230 Bound::Included(ref k) => match self.lookup_index_for(k) {
231 Ok(index) => index + 1,
234 Bound::Excluded(ref k) => match self.lookup_index_for(k) {
235 Ok(index) | Err(index) => index,
237 Bound::Unbounded => self.data.len(),
244 pub fn contains_key<Q>(&self, key: &Q) -> bool
249 self.get(key).is_some()
253 impl<K: Ord, V> IntoIterator for SortedMap<K, V> {
255 type IntoIter = ::std::vec::IntoIter<(K, V)>;
257 fn into_iter(self) -> Self::IntoIter {
258 self.data.into_iter()
262 impl<'a, K, Q, V> Index<&'a Q> for SortedMap<K, V>
269 fn index(&self, key: &Q) -> &Self::Output {
270 self.get(key).expect("no entry found for key")
274 impl<'a, K, Q, V> IndexMut<&'a Q> for SortedMap<K, V>
279 fn index_mut(&mut self, key: &Q) -> &mut Self::Output {
280 self.get_mut(key).expect("no entry found for key")
284 impl<K: Ord, V> FromIterator<(K, V)> for SortedMap<K, V> {
285 fn from_iter<T: IntoIterator<Item = (K, V)>>(iter: T) -> Self {
286 let mut data: Vec<(K, V)> = iter.into_iter().collect();
288 data.sort_unstable_by(|&(ref k1, _), &(ref k2, _)| k1.cmp(k2));
289 data.dedup_by(|&mut (ref k1, _), &mut (ref k2, _)| k1.cmp(k2) == Ordering::Equal);