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