Line data Source code
1 : use std::{alloc::Layout, cmp::Ordering, ops::RangeBounds};
2 :
3 : #[derive(Clone, Copy, Debug, PartialEq, Eq)]
4 : pub enum VecMapOrdering {
5 : Greater,
6 : GreaterOrEqual,
7 : }
8 :
9 : /// Ordered map datastructure implemented in a Vec.
10 : ///
11 : /// Append only - can only add keys that are larger than the
12 : /// current max key.
13 : /// Ordering can be adjusted using [`VecMapOrdering`]
14 : /// during `VecMap` construction.
15 : #[derive(Clone, Debug)]
16 : pub struct VecMap<K, V> {
17 : data: Vec<(K, V)>,
18 : ordering: VecMapOrdering,
19 : }
20 :
21 : impl<K, V> Default for VecMap<K, V> {
22 4748849 : fn default() -> Self {
23 4748849 : VecMap {
24 4748849 : data: Default::default(),
25 4748849 : ordering: VecMapOrdering::Greater,
26 4748849 : }
27 4748849 : }
28 : }
29 :
30 0 : #[derive(thiserror::Error, Debug)]
31 : pub enum VecMapError {
32 : #[error("Key violates ordering constraint")]
33 : InvalidKey,
34 : #[error("Mismatched ordering constraints")]
35 : ExtendOrderingError,
36 : }
37 :
38 : impl<K: Ord, V> VecMap<K, V> {
39 4 : pub fn new(ordering: VecMapOrdering) -> Self {
40 4 : Self {
41 4 : data: Vec::new(),
42 4 : ordering,
43 4 : }
44 4 : }
45 :
46 4 : pub fn with_capacity(capacity: usize, ordering: VecMapOrdering) -> Self {
47 4 : Self {
48 4 : data: Vec::with_capacity(capacity),
49 4 : ordering,
50 4 : }
51 4 : }
52 :
53 0 : pub fn is_empty(&self) -> bool {
54 0 : self.data.is_empty()
55 0 : }
56 :
57 8722909 : pub fn as_slice(&self) -> &[(K, V)] {
58 8722909 : self.data.as_slice()
59 8722909 : }
60 :
61 : /// This function may panic if given a range where the lower bound is
62 : /// greater than the upper bound.
63 499440 : pub fn slice_range<R: RangeBounds<K>>(&self, range: R) -> &[(K, V)] {
64 : use std::ops::Bound::*;
65 :
66 998834 : let binary_search = |k: &K| self.data.binary_search_by_key(&k, extract_key);
67 :
68 499440 : let start_idx = match range.start_bound() {
69 27 : Unbounded => 0,
70 499330 : Included(k) => binary_search(k).unwrap_or_else(std::convert::identity),
71 83 : Excluded(k) => match binary_search(k) {
72 36 : Ok(idx) => idx + 1,
73 47 : Err(idx) => idx,
74 : },
75 : };
76 :
77 499440 : let end_idx = match range.end_bound() {
78 19 : Unbounded => self.data.len(),
79 99 : Included(k) => match binary_search(k) {
80 44 : Ok(idx) => idx + 1,
81 55 : Err(idx) => idx,
82 : },
83 499322 : Excluded(k) => binary_search(k).unwrap_or_else(std::convert::identity),
84 : };
85 :
86 499440 : &self.data[start_idx..end_idx]
87 499440 : }
88 :
89 : /// Add a key value pair to the map.
90 : /// If `key` is not respective of the `self` ordering the
91 : /// pair will not be added and `InvalidKey` error will be returned.
92 2655159 : pub fn append(&mut self, key: K, value: V) -> Result<usize, VecMapError> {
93 2655159 : self.validate_key_order(&key)?;
94 :
95 2655157 : let delta_size = self.instrument_vec_op(|vec| vec.push((key, value)));
96 2655157 : Ok(delta_size)
97 2655159 : }
98 :
99 : /// Update the maximum key value pair or add a new key value pair to the map.
100 : /// If `key` is not respective of the `self` ordering no updates or additions
101 : /// will occur and `InvalidKey` error will be returned.
102 5090488 : pub fn append_or_update_last(
103 5090488 : &mut self,
104 5090488 : key: K,
105 5090488 : mut value: V,
106 5090488 : ) -> Result<(Option<V>, usize), VecMapError> {
107 5090488 : if let Some((last_key, last_value)) = self.data.last_mut() {
108 553169 : match key.cmp(last_key) {
109 0 : Ordering::Less => return Err(VecMapError::InvalidKey),
110 : Ordering::Equal => {
111 0 : std::mem::swap(last_value, &mut value);
112 : const DELTA_SIZE: usize = 0;
113 0 : return Ok((Some(value), DELTA_SIZE));
114 : }
115 553169 : Ordering::Greater => {}
116 : }
117 4537319 : }
118 :
119 5090488 : let delta_size = self.instrument_vec_op(|vec| vec.push((key, value)));
120 5090488 : Ok((None, delta_size))
121 5090488 : }
122 :
123 : /// Move items from `other` to the end of `self`, leaving `other` empty.
124 : /// If the `other` ordering is different from `self` ordering
125 : /// `ExtendOrderingError` error will be returned.
126 : /// If any keys in `other` is not respective of the ordering defined in
127 : /// `self`, `InvalidKey` error will be returned and no mutation will occur.
128 7 : pub fn extend(&mut self, other: &mut Self) -> Result<usize, VecMapError> {
129 7 : if self.ordering != other.ordering {
130 2 : return Err(VecMapError::ExtendOrderingError);
131 5 : }
132 5 :
133 5 : let other_first_opt = other.data.last().map(extract_key);
134 5 : if let Some(other_first) = other_first_opt {
135 4 : self.validate_key_order(other_first)?;
136 1 : }
137 :
138 3 : let delta_size = self.instrument_vec_op(|vec| vec.append(&mut other.data));
139 3 : Ok(delta_size)
140 7 : }
141 :
142 : /// Validate the current last key in `self` and key being
143 : /// inserted against the order defined in `self`.
144 2655163 : fn validate_key_order(&self, key: &K) -> Result<(), VecMapError> {
145 2655163 : if let Some(last_key) = self.data.last().map(extract_key) {
146 2443626 : match (&self.ordering, &key.cmp(last_key)) {
147 : (VecMapOrdering::Greater, Ordering::Less | Ordering::Equal) => {
148 3 : return Err(VecMapError::InvalidKey);
149 : }
150 2443611 : (VecMapOrdering::Greater, Ordering::Greater) => {}
151 : (VecMapOrdering::GreaterOrEqual, Ordering::Less) => {
152 1 : return Err(VecMapError::InvalidKey);
153 : }
154 11 : (VecMapOrdering::GreaterOrEqual, Ordering::Equal | Ordering::Greater) => {}
155 : }
156 211537 : }
157 :
158 2655159 : Ok(())
159 2655163 : }
160 :
161 : /// Instrument an operation on the underlying [`Vec`].
162 : /// Will panic if the operation decreases capacity.
163 : /// Returns the increase in memory usage caused by the op.
164 7745648 : fn instrument_vec_op(&mut self, op: impl FnOnce(&mut Vec<(K, V)>)) -> usize {
165 7745648 : let old_cap = self.data.capacity();
166 7745648 : op(&mut self.data);
167 7745648 : let new_cap = self.data.capacity();
168 7745648 :
169 7745648 : match old_cap.cmp(&new_cap) {
170 : Ordering::Less => {
171 4786013 : let old_size = Layout::array::<(K, V)>(old_cap).unwrap().size();
172 4786013 : let new_size = Layout::array::<(K, V)>(new_cap).unwrap().size();
173 4786013 : new_size - old_size
174 : }
175 2959635 : Ordering::Equal => 0,
176 0 : Ordering::Greater => panic!("VecMap capacity shouldn't ever decrease"),
177 : }
178 7745648 : }
179 :
180 : /// Similar to `from_iter` defined in `FromIter` trait except
181 : /// that it accepts an [`VecMapOrdering`]
182 4 : pub fn from_iter<I: IntoIterator<Item = (K, V)>>(iter: I, ordering: VecMapOrdering) -> Self {
183 4 : let iter = iter.into_iter();
184 4 : let initial_capacity = {
185 4 : match iter.size_hint() {
186 0 : (lower_bound, None) => lower_bound,
187 4 : (_, Some(upper_bound)) => upper_bound,
188 : }
189 : };
190 :
191 4 : let mut vec_map = VecMap::with_capacity(initial_capacity, ordering);
192 22 : for (key, value) in iter {
193 18 : vec_map
194 18 : .append(key, value)
195 18 : .expect("The passed collection needs to be sorted!");
196 18 : }
197 :
198 4 : vec_map
199 4 : }
200 : }
201 :
202 : impl<K: Ord, V> IntoIterator for VecMap<K, V> {
203 : type Item = (K, V);
204 : type IntoIter = std::vec::IntoIter<(K, V)>;
205 :
206 0 : fn into_iter(self) -> Self::IntoIter {
207 0 : self.data.into_iter()
208 0 : }
209 : }
210 :
211 12333184 : fn extract_key<K, V>(entry: &(K, V)) -> &K {
212 12333184 : &entry.0
213 12333184 : }
214 :
215 : #[cfg(test)]
216 : mod tests {
217 : use std::{collections::BTreeMap, ops::Bound};
218 :
219 : use super::{VecMap, VecMapOrdering};
220 :
221 : #[test]
222 1 : fn unbounded_range() {
223 1 : let mut vec = VecMap::default();
224 1 : vec.append(0, ()).unwrap();
225 1 :
226 1 : assert_eq!(vec.slice_range(0..0), &[]);
227 1 : }
228 :
229 : #[test]
230 : #[should_panic]
231 1 : fn invalid_ordering_range() {
232 1 : let mut vec = VecMap::default();
233 1 : vec.append(0, ()).unwrap();
234 1 :
235 1 : #[allow(clippy::reversed_empty_ranges)]
236 1 : vec.slice_range(1..0);
237 1 : }
238 :
239 : #[test]
240 1 : fn range_tests() {
241 1 : let mut vec = VecMap::default();
242 1 : vec.append(0, ()).unwrap();
243 1 : vec.append(2, ()).unwrap();
244 1 : vec.append(4, ()).unwrap();
245 1 :
246 1 : assert_eq!(vec.slice_range(0..0), &[]);
247 1 : assert_eq!(vec.slice_range(0..1), &[(0, ())]);
248 1 : assert_eq!(vec.slice_range(0..2), &[(0, ())]);
249 1 : assert_eq!(vec.slice_range(0..3), &[(0, ()), (2, ())]);
250 :
251 1 : assert_eq!(vec.slice_range(..0), &[]);
252 1 : assert_eq!(vec.slice_range(..1), &[(0, ())]);
253 :
254 1 : assert_eq!(vec.slice_range(..3), &[(0, ()), (2, ())]);
255 1 : assert_eq!(vec.slice_range(..3), &[(0, ()), (2, ())]);
256 :
257 1 : assert_eq!(vec.slice_range(0..=0), &[(0, ())]);
258 1 : assert_eq!(vec.slice_range(0..=1), &[(0, ())]);
259 1 : assert_eq!(vec.slice_range(0..=2), &[(0, ()), (2, ())]);
260 1 : assert_eq!(vec.slice_range(0..=3), &[(0, ()), (2, ())]);
261 :
262 1 : assert_eq!(vec.slice_range(..=0), &[(0, ())]);
263 1 : assert_eq!(vec.slice_range(..=1), &[(0, ())]);
264 1 : assert_eq!(vec.slice_range(..=2), &[(0, ()), (2, ())]);
265 1 : assert_eq!(vec.slice_range(..=3), &[(0, ()), (2, ())]);
266 1 : }
267 :
268 : struct BoundIter {
269 : min: i32,
270 : max: i32,
271 :
272 : next: Option<Bound<i32>>,
273 : }
274 :
275 : impl BoundIter {
276 20 : fn new(min: i32, max: i32) -> Self {
277 20 : Self {
278 20 : min,
279 20 : max,
280 20 :
281 20 : next: Some(Bound::Unbounded),
282 20 : }
283 20 : }
284 : }
285 :
286 : impl Iterator for BoundIter {
287 : type Item = Bound<i32>;
288 :
289 240 : fn next(&mut self) -> Option<Self::Item> {
290 240 : let cur = self.next?;
291 :
292 220 : self.next = match &cur {
293 20 : Bound::Unbounded => Some(Bound::Included(self.min)),
294 100 : Bound::Included(x) => {
295 100 : if *x >= self.max {
296 20 : Some(Bound::Excluded(self.min))
297 : } else {
298 80 : Some(Bound::Included(x + 1))
299 : }
300 : }
301 100 : Bound::Excluded(x) => {
302 100 : if *x >= self.max {
303 20 : None
304 : } else {
305 80 : Some(Bound::Excluded(x + 1))
306 : }
307 : }
308 : };
309 :
310 220 : Some(cur)
311 240 : }
312 : }
313 :
314 : #[test]
315 1 : fn range_exhaustive() {
316 4 : let map: BTreeMap<i32, ()> = (1..=7).step_by(2).map(|x| (x, ())).collect();
317 1 : let mut vec = VecMap::default();
318 4 : for &key in map.keys() {
319 4 : vec.append(key, ()).unwrap();
320 4 : }
321 :
322 : const RANGE_MIN: i32 = 0;
323 : const RANGE_MAX: i32 = 8;
324 20 : for lower_bound in BoundIter::new(RANGE_MIN, RANGE_MAX) {
325 19 : let ub_min = match lower_bound {
326 1 : Bound::Unbounded => RANGE_MIN,
327 9 : Bound::Included(x) => x,
328 9 : Bound::Excluded(x) => x + 1,
329 : };
330 201 : for upper_bound in BoundIter::new(ub_min, RANGE_MAX) {
331 201 : let map_range: Vec<(i32, ())> = map
332 201 : .range((lower_bound, upper_bound))
333 320 : .map(|(&x, _)| (x, ()))
334 201 : .collect();
335 201 : let vec_slice = vec.slice_range((lower_bound, upper_bound));
336 201 :
337 201 : assert_eq!(map_range, vec_slice);
338 : }
339 : }
340 1 : }
341 :
342 : #[test]
343 1 : fn extend() {
344 1 : let mut left = VecMap::default();
345 1 : left.append(0, ()).unwrap();
346 1 : assert_eq!(left.as_slice(), &[(0, ())]);
347 :
348 1 : let mut empty = VecMap::default();
349 1 : left.extend(&mut empty).unwrap();
350 1 : assert_eq!(left.as_slice(), &[(0, ())]);
351 1 : assert_eq!(empty.as_slice(), &[]);
352 :
353 1 : let mut right = VecMap::default();
354 1 : right.append(1, ()).unwrap();
355 1 :
356 1 : left.extend(&mut right).unwrap();
357 1 :
358 1 : assert_eq!(left.as_slice(), &[(0, ()), (1, ())]);
359 1 : assert_eq!(right.as_slice(), &[]);
360 :
361 1 : let mut zero_map = VecMap::default();
362 1 : zero_map.append(0, ()).unwrap();
363 1 :
364 1 : left.extend(&mut zero_map).unwrap_err();
365 1 : assert_eq!(left.as_slice(), &[(0, ()), (1, ())]);
366 1 : assert_eq!(zero_map.as_slice(), &[(0, ())]);
367 :
368 1 : let mut one_map = VecMap::default();
369 1 : one_map.append(1, ()).unwrap();
370 1 :
371 1 : left.extend(&mut one_map).unwrap_err();
372 1 : assert_eq!(left.as_slice(), &[(0, ()), (1, ())]);
373 1 : assert_eq!(one_map.as_slice(), &[(1, ())]);
374 :
375 1 : let mut map_greater_or_equal = VecMap::new(VecMapOrdering::GreaterOrEqual);
376 1 : map_greater_or_equal.append(2, ()).unwrap();
377 1 : map_greater_or_equal.append(2, ()).unwrap();
378 1 :
379 1 : left.extend(&mut map_greater_or_equal).unwrap_err();
380 1 : assert_eq!(left.as_slice(), &[(0, ()), (1, ())]);
381 1 : assert_eq!(map_greater_or_equal.as_slice(), &[(2, ()), (2, ())]);
382 1 : }
383 :
384 : #[test]
385 1 : fn extend_with_ordering() {
386 1 : let mut left = VecMap::new(VecMapOrdering::GreaterOrEqual);
387 1 : left.append(0, ()).unwrap();
388 1 : assert_eq!(left.as_slice(), &[(0, ())]);
389 :
390 1 : let mut greater_right = VecMap::new(VecMapOrdering::Greater);
391 1 : greater_right.append(0, ()).unwrap();
392 1 : left.extend(&mut greater_right).unwrap_err();
393 1 : assert_eq!(left.as_slice(), &[(0, ())]);
394 :
395 1 : let mut greater_or_equal_right = VecMap::new(VecMapOrdering::GreaterOrEqual);
396 1 : greater_or_equal_right.append(2, ()).unwrap();
397 1 : greater_or_equal_right.append(2, ()).unwrap();
398 1 : left.extend(&mut greater_or_equal_right).unwrap();
399 1 : assert_eq!(left.as_slice(), &[(0, ()), (2, ()), (2, ())]);
400 1 : }
401 :
402 : #[test]
403 1 : fn vec_map_from_sorted() {
404 1 : let vec = vec![(1, ()), (2, ()), (3, ()), (6, ())];
405 1 : let vec_map = VecMap::from_iter(vec, VecMapOrdering::Greater);
406 1 : assert_eq!(vec_map.as_slice(), &[(1, ()), (2, ()), (3, ()), (6, ())]);
407 :
408 1 : let vec = vec![(1, ()), (2, ()), (3, ()), (3, ()), (6, ()), (6, ())];
409 1 : let vec_map = VecMap::from_iter(vec, VecMapOrdering::GreaterOrEqual);
410 1 : assert_eq!(
411 1 : vec_map.as_slice(),
412 1 : &[(1, ()), (2, ()), (3, ()), (3, ()), (6, ()), (6, ())]
413 1 : );
414 1 : }
415 :
416 : #[test]
417 : #[should_panic]
418 1 : fn vec_map_from_unsorted_greater() {
419 1 : let vec = vec![(1, ()), (2, ()), (2, ()), (3, ()), (6, ())];
420 1 : let _ = VecMap::from_iter(vec, VecMapOrdering::Greater);
421 1 : }
422 :
423 : #[test]
424 : #[should_panic]
425 1 : fn vec_map_from_unsorted_greater_or_equal() {
426 1 : let vec = vec![(1, ()), (2, ()), (3, ()), (6, ()), (5, ())];
427 1 : let _ = VecMap::from_iter(vec, VecMapOrdering::GreaterOrEqual);
428 1 : }
429 : }
|