]>
Commit | Line | Data |
---|---|---|
f20569fa XL |
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] | |
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] | |
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] | |
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] | |
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] | |
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] | |
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] | |
90 | pub fn clear(&mut self) { | |
91 | self.data.clear(); | |
92 | } | |
93 | ||
94 | /// Iterate over elements, sorted by key | |
95 | #[inline] | |
96 | pub fn iter(&self) -> std::slice::Iter<'_, (K, V)> { | |
97 | self.data.iter() | |
98 | } | |
99 | ||
100 | /// Iterate over the keys, sorted | |
101 | #[inline] | |
102 | 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] | |
108 | pub fn values(&self) -> impl Iterator<Item = &V> + ExactSizeIterator + DoubleEndedIterator { | |
109 | self.data.iter().map(|&(_, ref v)| v) | |
110 | } | |
111 | ||
112 | #[inline] | |
113 | pub fn len(&self) -> usize { | |
114 | self.data.len() | |
115 | } | |
116 | ||
117 | #[inline] | |
118 | pub fn is_empty(&self) -> bool { | |
119 | self.len() == 0 | |
120 | } | |
121 | ||
122 | #[inline] | |
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] | |
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] | |
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] | |
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)] | |
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] | |
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] | |
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 | ||
246 | 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 | ||
258 | 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 | { | |
268 | 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> { | |
274 | 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; |