1 use crate::stable_hasher::{HashStable, StableHasher, StableOrd};
2 use std::borrow::Borrow;
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 is *O*(log(*n*)), insertion and removal
13 /// are *O*(*n*) but 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.
19 #[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Encodable, Decodable)]
20 pub struct SortedMap<K, V> {
24 impl<K, V> Default for SortedMap<K, V> {
26 fn default() -> SortedMap<K, V> {
27 SortedMap { data: Vec::new() }
31 impl<K, V> SortedMap<K, V> {
33 pub const fn new() -> SortedMap<K, V> {
34 SortedMap { data: Vec::new() }
38 impl<K: Ord, V> 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> {
46 debug_assert!(elements.array_windows().all(|[fst, snd]| fst.0 < snd.0));
48 SortedMap { data: elements }
52 pub fn insert(&mut self, key: K, mut value: V) -> Option<V> {
53 match self.lookup_index_for(&key) {
55 let slot = unsafe { self.data.get_unchecked_mut(index) };
56 mem::swap(&mut slot.1, &mut value);
60 self.data.insert(index, (key, value));
67 pub fn remove(&mut self, key: &K) -> Option<V> {
68 match self.lookup_index_for(key) {
69 Ok(index) => Some(self.data.remove(index).1),
75 pub fn get<Q>(&self, key: &Q) -> Option<&V>
80 match self.lookup_index_for(key) {
81 Ok(index) => unsafe { Some(&self.data.get_unchecked(index).1) },
87 pub fn get_mut<Q>(&mut self, key: &Q) -> Option<&mut V>
92 match self.lookup_index_for(key) {
93 Ok(index) => unsafe { Some(&mut self.data.get_unchecked_mut(index).1) },
98 /// Gets a mutable reference to the value in the entry, or insert a new one.
100 pub fn get_mut_or_insert_default(&mut self, key: K) -> &mut V
105 let index = match self.lookup_index_for(&key) {
108 self.data.insert(index, (key, V::default()));
112 unsafe { &mut self.data.get_unchecked_mut(index).1 }
116 pub fn clear(&mut self) {
120 /// Iterate over elements, sorted by key
122 pub fn iter(&self) -> std::slice::Iter<'_, (K, V)> {
126 /// Iterate over the keys, sorted
128 pub fn keys(&self) -> impl Iterator<Item = &K> + ExactSizeIterator + DoubleEndedIterator {
129 self.data.iter().map(|(k, _)| k)
132 /// Iterate over values, sorted by key
134 pub fn values(&self) -> impl Iterator<Item = &V> + ExactSizeIterator + DoubleEndedIterator {
135 self.data.iter().map(|(_, v)| v)
139 pub fn len(&self) -> usize {
144 pub fn is_empty(&self) -> bool {
149 pub fn range<R>(&self, range: R) -> &[(K, V)]
153 let (start, end) = self.range_slice_indices(range);
154 &self.data[start..end]
158 pub fn remove_range<R>(&mut self, range: R)
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)
173 self.data.iter_mut().map(|(k, _)| k).for_each(f);
176 /// Inserts a presorted range of elements into the map. If the range can be
177 /// inserted as a whole in between to existing elements of the map, this
178 /// will be faster than inserting the elements individually.
180 /// It is up to the caller to make sure that the elements are sorted by key
181 /// and that there are no duplicates.
183 pub fn insert_presorted(&mut self, elements: Vec<(K, V)>) {
184 if elements.is_empty() {
188 debug_assert!(elements.array_windows().all(|[fst, snd]| fst.0 < snd.0));
190 let start_index = self.lookup_index_for(&elements[0].0);
192 let elements = match start_index {
194 let mut elements = elements.into_iter();
195 self.data[index] = elements.next().unwrap();
199 if index == self.data.len() || 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.into_iter());
206 let mut elements = elements.into_iter();
207 self.data.insert(index, elements.next().unwrap());
213 for (k, v) in elements {
218 /// Looks up the key in `self.data` via `slice::binary_search()`.
220 fn lookup_index_for<Q>(&self, key: &Q) -> Result<usize, usize>
225 self.data.binary_search_by(|(x, _)| x.borrow().cmp(key))
229 fn range_slice_indices<R>(&self, range: R) -> (usize, usize)
233 let start = match range.start_bound() {
234 Bound::Included(k) => match self.lookup_index_for(k) {
235 Ok(index) | Err(index) => index,
237 Bound::Excluded(k) => match self.lookup_index_for(k) {
238 Ok(index) => index + 1,
241 Bound::Unbounded => 0,
244 let end = match range.end_bound() {
245 Bound::Included(k) => match self.lookup_index_for(k) {
246 Ok(index) => index + 1,
249 Bound::Excluded(k) => match self.lookup_index_for(k) {
250 Ok(index) | Err(index) => index,
252 Bound::Unbounded => self.data.len(),
259 pub fn contains_key<Q>(&self, key: &Q) -> bool
264 self.get(key).is_some()
268 impl<K: Ord, V> IntoIterator for SortedMap<K, V> {
270 type IntoIter = std::vec::IntoIter<(K, V)>;
272 fn into_iter(self) -> Self::IntoIter {
273 self.data.into_iter()
277 impl<'a, K, Q, V> Index<&'a Q> for SortedMap<K, V>
284 fn index(&self, key: &Q) -> &Self::Output {
285 self.get(key).expect("no entry found for key")
289 impl<'a, K, Q, V> IndexMut<&'a Q> for SortedMap<K, V>
294 fn index_mut(&mut self, key: &Q) -> &mut Self::Output {
295 self.get_mut(key).expect("no entry found for key")
299 impl<K: Ord, V> FromIterator<(K, V)> for SortedMap<K, V> {
300 fn from_iter<T: IntoIterator<Item = (K, V)>>(iter: T) -> Self {
301 let mut data: Vec<(K, V)> = iter.into_iter().collect();
303 data.sort_unstable_by(|(k1, _), (k2, _)| k1.cmp(k2));
304 data.dedup_by(|(k1, _), (k2, _)| k1 == k2);
310 impl<K: HashStable<CTX> + StableOrd, V: HashStable<CTX>, CTX> HashStable<CTX> for SortedMap<K, V> {
312 fn hash_stable(&self, ctx: &mut CTX, hasher: &mut StableHasher) {
313 self.data.hash_stable(ctx, hasher);
317 impl<K: Debug, V: Debug> Debug for SortedMap<K, V> {
318 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
319 f.debug_map().entries(self.data.iter().map(|(a, b)| (a, b))).finish()