1 use std::borrow::Borrow; 2 use std::cmp::Ordering; 3 use std::iter::FromIterator; 4 use std::mem; 5 use std::ops::{Bound, Index, IndexMut, RangeBounds}; 6 7 mod index_map; 8 9 pub use index_map::SortedIndexMultiMap; 10 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. 14 /// 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, Default, Debug, Encodable, Decodable)] 20 pub struct SortedMap<K: Ord, V> { 21 data: Vec<(K, V)>, 22 } 23 24 impl<K: Ord, V> SortedMap<K, V> { 25 #[inline] new() -> SortedMap<K, V>26 pub fn new() -> SortedMap<K, V> { 27 SortedMap { data: vec![] } 28 } 29 30 /// Construct a `SortedMap` from a presorted set of elements. This is faster 31 /// than creating an empty map and then inserting the elements individually. 32 /// 33 /// It is up to the caller to make sure that the elements are sorted by key 34 /// and that there are no duplicates. 35 #[inline] from_presorted_elements(elements: Vec<(K, V)>) -> SortedMap<K, V>36 pub fn from_presorted_elements(elements: Vec<(K, V)>) -> SortedMap<K, V> { 37 debug_assert!(elements.array_windows().all(|[fst, snd]| fst.0 < snd.0)); 38 39 SortedMap { data: elements } 40 } 41 42 #[inline] insert(&mut self, key: K, mut value: V) -> Option<V>43 pub fn insert(&mut self, key: K, mut value: V) -> Option<V> { 44 match self.lookup_index_for(&key) { 45 Ok(index) => { 46 let slot = unsafe { self.data.get_unchecked_mut(index) }; 47 mem::swap(&mut slot.1, &mut value); 48 Some(value) 49 } 50 Err(index) => { 51 self.data.insert(index, (key, value)); 52 None 53 } 54 } 55 } 56 57 #[inline] remove(&mut self, key: &K) -> Option<V>58 pub fn remove(&mut self, key: &K) -> Option<V> { 59 match self.lookup_index_for(key) { 60 Ok(index) => Some(self.data.remove(index).1), 61 Err(_) => None, 62 } 63 } 64 65 #[inline] get<Q>(&self, key: &Q) -> Option<&V> where K: Borrow<Q>, Q: Ord + ?Sized,66 pub fn get<Q>(&self, key: &Q) -> Option<&V> 67 where 68 K: Borrow<Q>, 69 Q: Ord + ?Sized, 70 { 71 match self.lookup_index_for(key) { 72 Ok(index) => unsafe { Some(&self.data.get_unchecked(index).1) }, 73 Err(_) => None, 74 } 75 } 76 77 #[inline] get_mut<Q>(&mut self, key: &Q) -> Option<&mut V> where K: Borrow<Q>, Q: Ord + ?Sized,78 pub fn get_mut<Q>(&mut self, key: &Q) -> Option<&mut V> 79 where 80 K: Borrow<Q>, 81 Q: Ord + ?Sized, 82 { 83 match self.lookup_index_for(key) { 84 Ok(index) => unsafe { Some(&mut self.data.get_unchecked_mut(index).1) }, 85 Err(_) => None, 86 } 87 } 88 89 #[inline] clear(&mut self)90 pub fn clear(&mut self) { 91 self.data.clear(); 92 } 93 94 /// Iterate over elements, sorted by key 95 #[inline] iter(&self) -> std::slice::Iter<'_, (K, V)>96 pub fn iter(&self) -> std::slice::Iter<'_, (K, V)> { 97 self.data.iter() 98 } 99 100 /// Iterate over the keys, sorted 101 #[inline] keys(&self) -> impl Iterator<Item = &K> + ExactSizeIterator + DoubleEndedIterator102 pub fn keys(&self) -> impl Iterator<Item = &K> + ExactSizeIterator + DoubleEndedIterator { 103 self.data.iter().map(|&(ref k, _)| k) 104 } 105 106 /// Iterate over values, sorted by key 107 #[inline] values(&self) -> impl Iterator<Item = &V> + ExactSizeIterator + DoubleEndedIterator108 pub fn values(&self) -> impl Iterator<Item = &V> + ExactSizeIterator + DoubleEndedIterator { 109 self.data.iter().map(|&(_, ref v)| v) 110 } 111 112 #[inline] len(&self) -> usize113 pub fn len(&self) -> usize { 114 self.data.len() 115 } 116 117 #[inline] is_empty(&self) -> bool118 pub fn is_empty(&self) -> bool { 119 self.len() == 0 120 } 121 122 #[inline] range<R>(&self, range: R) -> &[(K, V)] where R: RangeBounds<K>,123 pub fn range<R>(&self, range: R) -> &[(K, V)] 124 where 125 R: RangeBounds<K>, 126 { 127 let (start, end) = self.range_slice_indices(range); 128 &self.data[start..end] 129 } 130 131 #[inline] remove_range<R>(&mut self, range: R) where R: RangeBounds<K>,132 pub fn remove_range<R>(&mut self, range: R) 133 where 134 R: RangeBounds<K>, 135 { 136 let (start, end) = self.range_slice_indices(range); 137 self.data.splice(start..end, std::iter::empty()); 138 } 139 140 /// Mutate all keys with the given function `f`. This mutation must not 141 /// change the sort-order of keys. 142 #[inline] offset_keys<F>(&mut self, f: F) where F: Fn(&mut K),143 pub fn offset_keys<F>(&mut self, f: F) 144 where 145 F: Fn(&mut K), 146 { 147 self.data.iter_mut().map(|&mut (ref mut k, _)| k).for_each(f); 148 } 149 150 /// Inserts a presorted range of elements into the map. If the range can be 151 /// inserted as a whole in between to existing elements of the map, this 152 /// will be faster than inserting the elements individually. 153 /// 154 /// It is up to the caller to make sure that the elements are sorted by key 155 /// and that there are no duplicates. 156 #[inline] insert_presorted(&mut self, mut elements: Vec<(K, V)>)157 pub fn insert_presorted(&mut self, mut elements: Vec<(K, V)>) { 158 if elements.is_empty() { 159 return; 160 } 161 162 debug_assert!(elements.array_windows().all(|[fst, snd]| fst.0 < snd.0)); 163 164 let start_index = self.lookup_index_for(&elements[0].0); 165 166 let drain = match start_index { 167 Ok(index) => { 168 let mut drain = elements.drain(..); 169 self.data[index] = drain.next().unwrap(); 170 drain 171 } 172 Err(index) => { 173 if index == self.data.len() || elements.last().unwrap().0 < self.data[index].0 { 174 // We can copy the whole range without having to mix with 175 // existing elements. 176 self.data.splice(index..index, elements.drain(..)); 177 return; 178 } 179 180 let mut drain = elements.drain(..); 181 self.data.insert(index, drain.next().unwrap()); 182 drain 183 } 184 }; 185 186 // Insert the rest 187 for (k, v) in drain { 188 self.insert(k, v); 189 } 190 } 191 192 /// Looks up the key in `self.data` via `slice::binary_search()`. 193 #[inline(always)] lookup_index_for<Q>(&self, key: &Q) -> Result<usize, usize> where K: Borrow<Q>, Q: Ord + ?Sized,194 fn lookup_index_for<Q>(&self, key: &Q) -> Result<usize, usize> 195 where 196 K: Borrow<Q>, 197 Q: Ord + ?Sized, 198 { 199 self.data.binary_search_by(|&(ref x, _)| x.borrow().cmp(key)) 200 } 201 202 #[inline] range_slice_indices<R>(&self, range: R) -> (usize, usize) where R: RangeBounds<K>,203 fn range_slice_indices<R>(&self, range: R) -> (usize, usize) 204 where 205 R: RangeBounds<K>, 206 { 207 let start = match range.start_bound() { 208 Bound::Included(ref k) => match self.lookup_index_for(k) { 209 Ok(index) | Err(index) => index, 210 }, 211 Bound::Excluded(ref k) => match self.lookup_index_for(k) { 212 Ok(index) => index + 1, 213 Err(index) => index, 214 }, 215 Bound::Unbounded => 0, 216 }; 217 218 let end = match range.end_bound() { 219 Bound::Included(ref k) => match self.lookup_index_for(k) { 220 Ok(index) => index + 1, 221 Err(index) => index, 222 }, 223 Bound::Excluded(ref k) => match self.lookup_index_for(k) { 224 Ok(index) | Err(index) => index, 225 }, 226 Bound::Unbounded => self.data.len(), 227 }; 228 229 (start, end) 230 } 231 232 #[inline] contains_key<Q>(&self, key: &Q) -> bool where K: Borrow<Q>, Q: Ord + ?Sized,233 pub fn contains_key<Q>(&self, key: &Q) -> bool 234 where 235 K: Borrow<Q>, 236 Q: Ord + ?Sized, 237 { 238 self.get(key).is_some() 239 } 240 } 241 242 impl<K: Ord, V> IntoIterator for SortedMap<K, V> { 243 type Item = (K, V); 244 type IntoIter = std::vec::IntoIter<(K, V)>; 245 into_iter(self) -> Self::IntoIter246 fn into_iter(self) -> Self::IntoIter { 247 self.data.into_iter() 248 } 249 } 250 251 impl<'a, K, Q, V> Index<&'a Q> for SortedMap<K, V> 252 where 253 K: Ord + Borrow<Q>, 254 Q: Ord + ?Sized, 255 { 256 type Output = V; 257 index(&self, key: &Q) -> &Self::Output258 fn index(&self, key: &Q) -> &Self::Output { 259 self.get(key).expect("no entry found for key") 260 } 261 } 262 263 impl<'a, K, Q, V> IndexMut<&'a Q> for SortedMap<K, V> 264 where 265 K: Ord + Borrow<Q>, 266 Q: Ord + ?Sized, 267 { index_mut(&mut self, key: &Q) -> &mut Self::Output268 fn index_mut(&mut self, key: &Q) -> &mut Self::Output { 269 self.get_mut(key).expect("no entry found for key") 270 } 271 } 272 273 impl<K: Ord, V> FromIterator<(K, V)> for SortedMap<K, V> { from_iter<T: IntoIterator<Item = (K, V)>>(iter: T) -> Self274 fn from_iter<T: IntoIterator<Item = (K, V)>>(iter: T) -> Self { 275 let mut data: Vec<(K, V)> = iter.into_iter().collect(); 276 277 data.sort_unstable_by(|&(ref k1, _), &(ref k2, _)| k1.cmp(k2)); 278 data.dedup_by(|&mut (ref k1, _), &mut (ref k2, _)| k1.cmp(k2) == Ordering::Equal); 279 280 SortedMap { data } 281 } 282 } 283 284 #[cfg(test)] 285 mod tests; 286