1 use crate::stable_hasher::{HashStable, StableHasher};
2 use std::borrow::Borrow;
3 use std::cmp::Ordering;
4 use std::iter::FromIterator;
5 use std::mem;
6 use std::ops::{Bound, Index, IndexMut, RangeBounds};
7 
8 mod index_map;
9 
10 pub use index_map::SortedIndexMultiMap;
11 
12 /// `SortedMap` is a data structure with similar characteristics as BTreeMap but
13 /// slightly different trade-offs: lookup, insertion, and removal are *O*(log(*n*))
14 /// and elements can be iterated in order cheaply.
15 ///
16 /// `SortedMap` can be faster than a `BTreeMap` for small sizes (<50) since it
17 /// stores data in a more compact way. It also supports accessing contiguous
18 /// ranges of elements as a slice, and slices of already sorted elements can be
19 /// inserted efficiently.
20 #[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Debug, Encodable, Decodable)]
21 pub struct SortedMap<K, V> {
22     data: Vec<(K, V)>,
23 }
24 
25 impl<K, V> Default for SortedMap<K, V> {
26     #[inline]
default() -> SortedMap<K, V>27     fn default() -> SortedMap<K, V> {
28         SortedMap { data: Vec::new() }
29     }
30 }
31 
32 impl<K, V> SortedMap<K, V> {
33     #[inline]
new() -> SortedMap<K, V>34     pub const fn new() -> SortedMap<K, V> {
35         SortedMap { data: Vec::new() }
36     }
37 }
38 
39 impl<K: Ord, V> SortedMap<K, V> {
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]
from_presorted_elements(elements: Vec<(K, V)>) -> SortedMap<K, V>46     pub fn from_presorted_elements(elements: Vec<(K, V)>) -> SortedMap<K, V> {
47         debug_assert!(elements.array_windows().all(|[fst, snd]| fst.0 < snd.0));
48 
49         SortedMap { data: elements }
50     }
51 
52     #[inline]
insert(&mut self, key: K, mut value: V) -> Option<V>53     pub fn insert(&mut self, key: K, mut value: V) -> Option<V> {
54         match self.lookup_index_for(&key) {
55             Ok(index) => {
56                 let slot = unsafe { self.data.get_unchecked_mut(index) };
57                 mem::swap(&mut slot.1, &mut value);
58                 Some(value)
59             }
60             Err(index) => {
61                 self.data.insert(index, (key, value));
62                 None
63             }
64         }
65     }
66 
67     #[inline]
remove(&mut self, key: &K) -> Option<V>68     pub fn remove(&mut self, key: &K) -> Option<V> {
69         match self.lookup_index_for(key) {
70             Ok(index) => Some(self.data.remove(index).1),
71             Err(_) => None,
72         }
73     }
74 
75     #[inline]
get<Q>(&self, key: &Q) -> Option<&V> where K: Borrow<Q>, Q: Ord + ?Sized,76     pub fn get<Q>(&self, key: &Q) -> Option<&V>
77     where
78         K: Borrow<Q>,
79         Q: Ord + ?Sized,
80     {
81         match self.lookup_index_for(key) {
82             Ok(index) => unsafe { Some(&self.data.get_unchecked(index).1) },
83             Err(_) => None,
84         }
85     }
86 
87     #[inline]
get_mut<Q>(&mut self, key: &Q) -> Option<&mut V> where K: Borrow<Q>, Q: Ord + ?Sized,88     pub fn get_mut<Q>(&mut self, key: &Q) -> Option<&mut V>
89     where
90         K: Borrow<Q>,
91         Q: Ord + ?Sized,
92     {
93         match self.lookup_index_for(key) {
94             Ok(index) => unsafe { Some(&mut self.data.get_unchecked_mut(index).1) },
95             Err(_) => None,
96         }
97     }
98 
99     #[inline]
clear(&mut self)100     pub fn clear(&mut self) {
101         self.data.clear();
102     }
103 
104     /// Iterate over elements, sorted by key
105     #[inline]
iter(&self) -> std::slice::Iter<'_, (K, V)>106     pub fn iter(&self) -> std::slice::Iter<'_, (K, V)> {
107         self.data.iter()
108     }
109 
110     /// Iterate over the keys, sorted
111     #[inline]
keys(&self) -> impl Iterator<Item = &K> + ExactSizeIterator + DoubleEndedIterator112     pub fn keys(&self) -> impl Iterator<Item = &K> + ExactSizeIterator + DoubleEndedIterator {
113         self.data.iter().map(|&(ref k, _)| k)
114     }
115 
116     /// Iterate over values, sorted by key
117     #[inline]
values(&self) -> impl Iterator<Item = &V> + ExactSizeIterator + DoubleEndedIterator118     pub fn values(&self) -> impl Iterator<Item = &V> + ExactSizeIterator + DoubleEndedIterator {
119         self.data.iter().map(|&(_, ref v)| v)
120     }
121 
122     #[inline]
len(&self) -> usize123     pub fn len(&self) -> usize {
124         self.data.len()
125     }
126 
127     #[inline]
is_empty(&self) -> bool128     pub fn is_empty(&self) -> bool {
129         self.len() == 0
130     }
131 
132     #[inline]
range<R>(&self, range: R) -> &[(K, V)] where R: RangeBounds<K>,133     pub fn range<R>(&self, range: R) -> &[(K, V)]
134     where
135         R: RangeBounds<K>,
136     {
137         let (start, end) = self.range_slice_indices(range);
138         &self.data[start..end]
139     }
140 
141     #[inline]
remove_range<R>(&mut self, range: R) where R: RangeBounds<K>,142     pub fn remove_range<R>(&mut self, range: R)
143     where
144         R: RangeBounds<K>,
145     {
146         let (start, end) = self.range_slice_indices(range);
147         self.data.splice(start..end, std::iter::empty());
148     }
149 
150     /// Mutate all keys with the given function `f`. This mutation must not
151     /// change the sort-order of keys.
152     #[inline]
offset_keys<F>(&mut self, f: F) where F: Fn(&mut K),153     pub fn offset_keys<F>(&mut self, f: F)
154     where
155         F: Fn(&mut K),
156     {
157         self.data.iter_mut().map(|&mut (ref mut k, _)| k).for_each(f);
158     }
159 
160     /// Inserts a presorted range of elements into the map. If the range can be
161     /// inserted as a whole in between to existing elements of the map, this
162     /// will be faster than inserting the elements individually.
163     ///
164     /// It is up to the caller to make sure that the elements are sorted by key
165     /// and that there are no duplicates.
166     #[inline]
insert_presorted(&mut self, mut elements: Vec<(K, V)>)167     pub fn insert_presorted(&mut self, mut elements: Vec<(K, V)>) {
168         if elements.is_empty() {
169             return;
170         }
171 
172         debug_assert!(elements.array_windows().all(|[fst, snd]| fst.0 < snd.0));
173 
174         let start_index = self.lookup_index_for(&elements[0].0);
175 
176         let drain = match start_index {
177             Ok(index) => {
178                 let mut drain = elements.drain(..);
179                 self.data[index] = drain.next().unwrap();
180                 drain
181             }
182             Err(index) => {
183                 if index == self.data.len() || elements.last().unwrap().0 < self.data[index].0 {
184                     // We can copy the whole range without having to mix with
185                     // existing elements.
186                     self.data.splice(index..index, elements.drain(..));
187                     return;
188                 }
189 
190                 let mut drain = elements.drain(..);
191                 self.data.insert(index, drain.next().unwrap());
192                 drain
193             }
194         };
195 
196         // Insert the rest
197         for (k, v) in drain {
198             self.insert(k, v);
199         }
200     }
201 
202     /// Looks up the key in `self.data` via `slice::binary_search()`.
203     #[inline(always)]
lookup_index_for<Q>(&self, key: &Q) -> Result<usize, usize> where K: Borrow<Q>, Q: Ord + ?Sized,204     fn lookup_index_for<Q>(&self, key: &Q) -> Result<usize, usize>
205     where
206         K: Borrow<Q>,
207         Q: Ord + ?Sized,
208     {
209         self.data.binary_search_by(|&(ref x, _)| x.borrow().cmp(key))
210     }
211 
212     #[inline]
range_slice_indices<R>(&self, range: R) -> (usize, usize) where R: RangeBounds<K>,213     fn range_slice_indices<R>(&self, range: R) -> (usize, usize)
214     where
215         R: RangeBounds<K>,
216     {
217         let start = match range.start_bound() {
218             Bound::Included(ref k) => match self.lookup_index_for(k) {
219                 Ok(index) | Err(index) => index,
220             },
221             Bound::Excluded(ref k) => match self.lookup_index_for(k) {
222                 Ok(index) => index + 1,
223                 Err(index) => index,
224             },
225             Bound::Unbounded => 0,
226         };
227 
228         let end = match range.end_bound() {
229             Bound::Included(ref k) => match self.lookup_index_for(k) {
230                 Ok(index) => index + 1,
231                 Err(index) => index,
232             },
233             Bound::Excluded(ref k) => match self.lookup_index_for(k) {
234                 Ok(index) | Err(index) => index,
235             },
236             Bound::Unbounded => self.data.len(),
237         };
238 
239         (start, end)
240     }
241 
242     #[inline]
contains_key<Q>(&self, key: &Q) -> bool where K: Borrow<Q>, Q: Ord + ?Sized,243     pub fn contains_key<Q>(&self, key: &Q) -> bool
244     where
245         K: Borrow<Q>,
246         Q: Ord + ?Sized,
247     {
248         self.get(key).is_some()
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 
into_iter(self) -> Self::IntoIter256     fn into_iter(self) -> Self::IntoIter {
257         self.data.into_iter()
258     }
259 }
260 
261 impl<'a, K, Q, V> Index<&'a Q> for SortedMap<K, V>
262 where
263     K: Ord + Borrow<Q>,
264     Q: Ord + ?Sized,
265 {
266     type Output = V;
267 
index(&self, key: &Q) -> &Self::Output268     fn index(&self, key: &Q) -> &Self::Output {
269         self.get(key).expect("no entry found for key")
270     }
271 }
272 
273 impl<'a, K, Q, V> IndexMut<&'a Q> for SortedMap<K, V>
274 where
275     K: Ord + Borrow<Q>,
276     Q: Ord + ?Sized,
277 {
index_mut(&mut self, key: &Q) -> &mut Self::Output278     fn index_mut(&mut self, key: &Q) -> &mut Self::Output {
279         self.get_mut(key).expect("no entry found for key")
280     }
281 }
282 
283 impl<K: Ord, V> FromIterator<(K, V)> for SortedMap<K, V> {
from_iter<T: IntoIterator<Item = (K, V)>>(iter: T) -> Self284     fn from_iter<T: IntoIterator<Item = (K, V)>>(iter: T) -> Self {
285         let mut data: Vec<(K, V)> = iter.into_iter().collect();
286 
287         data.sort_unstable_by(|&(ref k1, _), &(ref k2, _)| k1.cmp(k2));
288         data.dedup_by(|&mut (ref k1, _), &mut (ref k2, _)| k1.cmp(k2) == Ordering::Equal);
289 
290         SortedMap { data }
291     }
292 }
293 
294 impl<K: HashStable<CTX>, V: HashStable<CTX>, CTX> HashStable<CTX> for SortedMap<K, V> {
295     #[inline]
hash_stable(&self, ctx: &mut CTX, hasher: &mut StableHasher)296     fn hash_stable(&self, ctx: &mut CTX, hasher: &mut StableHasher) {
297         self.data.hash_stable(ctx, hasher);
298     }
299 }
300 
301 #[cfg(test)]
302 mod tests;
303