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