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 14234961 : fn default() -> Self {
23 14234961 : VecMap {
24 14234961 : data: Default::default(),
25 14234961 : ordering: VecMapOrdering::Greater,
26 14234961 : }
27 14234961 : }
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 26155399 : pub fn as_slice(&self) -> &[(K, V)] {
58 26155399 : self.data.as_slice()
59 26155399 : }
60 :
61 : /// This function may panic if given a range where the lower bound is
62 : /// greater than the upper bound.
63 1496348 : pub fn slice_range<R: RangeBounds<K>>(&self, range: R) -> &[(K, V)] {
64 : use std::ops::Bound::*;
65 :
66 2992650 : let binary_search = |k: &K| self.data.binary_search_by_key(&k, extract_key);
67 :
68 1496348 : let start_idx = match range.start_bound() {
69 27 : Unbounded => 0,
70 1496238 : 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 1496348 : 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 1496230 : Excluded(k) => binary_search(k).unwrap_or_else(std::convert::identity),
84 : };
85 :
86 1496348 : &self.data[start_idx..end_idx]
87 1496348 : }
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 7858452 : pub fn append(&mut self, key: K, value: V) -> Result<usize, VecMapError> {
93 7858452 : self.validate_key_order(&key)?;
94 :
95 7858450 : let delta_size = self.instrument_vec_op(|vec| vec.push((key, value)));
96 7858450 : Ok(delta_size)
97 7858452 : }
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 15271446 : pub fn append_or_update_last(
103 15271446 : &mut self,
104 15271446 : key: K,
105 15271446 : mut value: V,
106 15271446 : ) -> Result<(Option<V>, usize), VecMapError> {
107 15271446 : if let Some((last_key, last_value)) = self.data.last_mut() {
108 1661139 : 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 1661139 : Ordering::Greater => {}
116 : }
117 13610307 : }
118 :
119 15271446 : let delta_size = self.instrument_vec_op(|vec| vec.push((key, value)));
120 15271446 : Ok((None, delta_size))
121 15271446 : }
122 :
123 : /// Split the map into two.
124 : ///
125 : /// The left map contains everything before `cutoff` (exclusive).
126 : /// Right map contains `cutoff` and everything after (inclusive).
127 0 : pub fn split_at(&self, cutoff: &K) -> (Self, Self)
128 0 : where
129 0 : K: Clone,
130 0 : V: Clone,
131 0 : {
132 0 : let split_idx = self
133 0 : .data
134 0 : .binary_search_by_key(&cutoff, extract_key)
135 0 : .unwrap_or_else(std::convert::identity);
136 0 :
137 0 : (
138 0 : VecMap {
139 0 : data: self.data[..split_idx].to_vec(),
140 0 : ordering: self.ordering,
141 0 : },
142 0 : VecMap {
143 0 : data: self.data[split_idx..].to_vec(),
144 0 : ordering: self.ordering,
145 0 : },
146 0 : )
147 0 : }
148 :
149 : /// Move items from `other` to the end of `self`, leaving `other` empty.
150 : /// If the `other` ordering is different from `self` ordering
151 : /// `ExtendOrderingError` error will be returned.
152 : /// If any keys in `other` is not respective of the ordering defined in
153 : /// `self`, `InvalidKey` error will be returned and no mutation will occur.
154 7 : pub fn extend(&mut self, other: &mut Self) -> Result<usize, VecMapError> {
155 7 : if self.ordering != other.ordering {
156 2 : return Err(VecMapError::ExtendOrderingError);
157 5 : }
158 5 :
159 5 : let other_first_opt = other.data.last().map(extract_key);
160 5 : if let Some(other_first) = other_first_opt {
161 4 : self.validate_key_order(other_first)?;
162 1 : }
163 :
164 3 : let delta_size = self.instrument_vec_op(|vec| vec.append(&mut other.data));
165 3 : Ok(delta_size)
166 7 : }
167 :
168 : /// Validate the current last key in `self` and key being
169 : /// inserted against the order defined in `self`.
170 7858456 : fn validate_key_order(&self, key: &K) -> Result<(), VecMapError> {
171 7858456 : if let Some(last_key) = self.data.last().map(extract_key) {
172 7233795 : match (&self.ordering, &key.cmp(last_key)) {
173 : (VecMapOrdering::Greater, Ordering::Less | Ordering::Equal) => {
174 3 : return Err(VecMapError::InvalidKey);
175 : }
176 7233780 : (VecMapOrdering::Greater, Ordering::Greater) => {}
177 : (VecMapOrdering::GreaterOrEqual, Ordering::Less) => {
178 1 : return Err(VecMapError::InvalidKey);
179 : }
180 11 : (VecMapOrdering::GreaterOrEqual, Ordering::Equal | Ordering::Greater) => {}
181 : }
182 624661 : }
183 :
184 7858452 : Ok(())
185 7858456 : }
186 :
187 : /// Instrument an operation on the underlying [`Vec`].
188 : /// Will panic if the operation decreases capacity.
189 : /// Returns the increase in memory usage caused by the op.
190 23129899 : fn instrument_vec_op(&mut self, op: impl FnOnce(&mut Vec<(K, V)>)) -> usize {
191 23129899 : let old_cap = self.data.capacity();
192 23129899 : op(&mut self.data);
193 23129899 : let new_cap = self.data.capacity();
194 23129899 :
195 23129899 : match old_cap.cmp(&new_cap) {
196 : Ordering::Less => {
197 14346008 : let old_size = Layout::array::<(K, V)>(old_cap).unwrap().size();
198 14346008 : let new_size = Layout::array::<(K, V)>(new_cap).unwrap().size();
199 14346008 : new_size - old_size
200 : }
201 8783891 : Ordering::Equal => 0,
202 0 : Ordering::Greater => panic!("VecMap capacity shouldn't ever decrease"),
203 : }
204 23129899 : }
205 :
206 : /// Similar to `from_iter` defined in `FromIter` trait except
207 : /// that it accepts an [`VecMapOrdering`]
208 4 : pub fn from_iter<I: IntoIterator<Item = (K, V)>>(iter: I, ordering: VecMapOrdering) -> Self {
209 4 : let iter = iter.into_iter();
210 4 : let initial_capacity = {
211 4 : match iter.size_hint() {
212 0 : (lower_bound, None) => lower_bound,
213 4 : (_, Some(upper_bound)) => upper_bound,
214 : }
215 : };
216 :
217 4 : let mut vec_map = VecMap::with_capacity(initial_capacity, ordering);
218 22 : for (key, value) in iter {
219 18 : vec_map
220 18 : .append(key, value)
221 18 : .expect("The passed collection needs to be sorted!");
222 18 : }
223 :
224 4 : vec_map
225 4 : }
226 : }
227 :
228 : impl<K: Ord, V> IntoIterator for VecMap<K, V> {
229 : type Item = (K, V);
230 : type IntoIter = std::vec::IntoIter<(K, V)>;
231 :
232 0 : fn into_iter(self) -> Self::IntoIter {
233 0 : self.data.into_iter()
234 0 : }
235 : }
236 :
237 34132239 : fn extract_key<K, V>(entry: &(K, V)) -> &K {
238 34132239 : &entry.0
239 34132239 : }
240 :
241 : #[cfg(test)]
242 : mod tests {
243 : use std::{collections::BTreeMap, ops::Bound};
244 :
245 : use super::{VecMap, VecMapOrdering};
246 :
247 : #[test]
248 1 : fn unbounded_range() {
249 1 : let mut vec = VecMap::default();
250 1 : vec.append(0, ()).unwrap();
251 1 :
252 1 : assert_eq!(vec.slice_range(0..0), &[]);
253 1 : }
254 :
255 : #[test]
256 : #[should_panic]
257 1 : fn invalid_ordering_range() {
258 1 : let mut vec = VecMap::default();
259 1 : vec.append(0, ()).unwrap();
260 1 :
261 1 : #[allow(clippy::reversed_empty_ranges)]
262 1 : vec.slice_range(1..0);
263 1 : }
264 :
265 : #[test]
266 1 : fn range_tests() {
267 1 : let mut vec = VecMap::default();
268 1 : vec.append(0, ()).unwrap();
269 1 : vec.append(2, ()).unwrap();
270 1 : vec.append(4, ()).unwrap();
271 1 :
272 1 : assert_eq!(vec.slice_range(0..0), &[]);
273 1 : assert_eq!(vec.slice_range(0..1), &[(0, ())]);
274 1 : assert_eq!(vec.slice_range(0..2), &[(0, ())]);
275 1 : assert_eq!(vec.slice_range(0..3), &[(0, ()), (2, ())]);
276 :
277 1 : assert_eq!(vec.slice_range(..0), &[]);
278 1 : assert_eq!(vec.slice_range(..1), &[(0, ())]);
279 :
280 1 : assert_eq!(vec.slice_range(..3), &[(0, ()), (2, ())]);
281 1 : assert_eq!(vec.slice_range(..3), &[(0, ()), (2, ())]);
282 :
283 1 : assert_eq!(vec.slice_range(0..=0), &[(0, ())]);
284 1 : assert_eq!(vec.slice_range(0..=1), &[(0, ())]);
285 1 : assert_eq!(vec.slice_range(0..=2), &[(0, ()), (2, ())]);
286 1 : assert_eq!(vec.slice_range(0..=3), &[(0, ()), (2, ())]);
287 :
288 1 : assert_eq!(vec.slice_range(..=0), &[(0, ())]);
289 1 : assert_eq!(vec.slice_range(..=1), &[(0, ())]);
290 1 : assert_eq!(vec.slice_range(..=2), &[(0, ()), (2, ())]);
291 1 : assert_eq!(vec.slice_range(..=3), &[(0, ()), (2, ())]);
292 1 : }
293 :
294 : struct BoundIter {
295 : min: i32,
296 : max: i32,
297 :
298 : next: Option<Bound<i32>>,
299 : }
300 :
301 : impl BoundIter {
302 20 : fn new(min: i32, max: i32) -> Self {
303 20 : Self {
304 20 : min,
305 20 : max,
306 20 :
307 20 : next: Some(Bound::Unbounded),
308 20 : }
309 20 : }
310 : }
311 :
312 : impl Iterator for BoundIter {
313 : type Item = Bound<i32>;
314 :
315 240 : fn next(&mut self) -> Option<Self::Item> {
316 240 : let cur = self.next?;
317 :
318 220 : self.next = match &cur {
319 20 : Bound::Unbounded => Some(Bound::Included(self.min)),
320 100 : Bound::Included(x) => {
321 100 : if *x >= self.max {
322 20 : Some(Bound::Excluded(self.min))
323 : } else {
324 80 : Some(Bound::Included(x + 1))
325 : }
326 : }
327 100 : Bound::Excluded(x) => {
328 100 : if *x >= self.max {
329 20 : None
330 : } else {
331 80 : Some(Bound::Excluded(x + 1))
332 : }
333 : }
334 : };
335 :
336 220 : Some(cur)
337 240 : }
338 : }
339 :
340 : #[test]
341 1 : fn range_exhaustive() {
342 4 : let map: BTreeMap<i32, ()> = (1..=7).step_by(2).map(|x| (x, ())).collect();
343 1 : let mut vec = VecMap::default();
344 4 : for &key in map.keys() {
345 4 : vec.append(key, ()).unwrap();
346 4 : }
347 :
348 : const RANGE_MIN: i32 = 0;
349 : const RANGE_MAX: i32 = 8;
350 20 : for lower_bound in BoundIter::new(RANGE_MIN, RANGE_MAX) {
351 19 : let ub_min = match lower_bound {
352 1 : Bound::Unbounded => RANGE_MIN,
353 9 : Bound::Included(x) => x,
354 9 : Bound::Excluded(x) => x + 1,
355 : };
356 201 : for upper_bound in BoundIter::new(ub_min, RANGE_MAX) {
357 201 : let map_range: Vec<(i32, ())> = map
358 201 : .range((lower_bound, upper_bound))
359 320 : .map(|(&x, _)| (x, ()))
360 201 : .collect();
361 201 : let vec_slice = vec.slice_range((lower_bound, upper_bound));
362 201 :
363 201 : assert_eq!(map_range, vec_slice);
364 : }
365 : }
366 1 : }
367 :
368 : #[test]
369 1 : fn extend() {
370 1 : let mut left = VecMap::default();
371 1 : left.append(0, ()).unwrap();
372 1 : assert_eq!(left.as_slice(), &[(0, ())]);
373 :
374 1 : let mut empty = VecMap::default();
375 1 : left.extend(&mut empty).unwrap();
376 1 : assert_eq!(left.as_slice(), &[(0, ())]);
377 1 : assert_eq!(empty.as_slice(), &[]);
378 :
379 1 : let mut right = VecMap::default();
380 1 : right.append(1, ()).unwrap();
381 1 :
382 1 : left.extend(&mut right).unwrap();
383 1 :
384 1 : assert_eq!(left.as_slice(), &[(0, ()), (1, ())]);
385 1 : assert_eq!(right.as_slice(), &[]);
386 :
387 1 : let mut zero_map = VecMap::default();
388 1 : zero_map.append(0, ()).unwrap();
389 1 :
390 1 : left.extend(&mut zero_map).unwrap_err();
391 1 : assert_eq!(left.as_slice(), &[(0, ()), (1, ())]);
392 1 : assert_eq!(zero_map.as_slice(), &[(0, ())]);
393 :
394 1 : let mut one_map = VecMap::default();
395 1 : one_map.append(1, ()).unwrap();
396 1 :
397 1 : left.extend(&mut one_map).unwrap_err();
398 1 : assert_eq!(left.as_slice(), &[(0, ()), (1, ())]);
399 1 : assert_eq!(one_map.as_slice(), &[(1, ())]);
400 :
401 1 : let mut map_greater_or_equal = VecMap::new(VecMapOrdering::GreaterOrEqual);
402 1 : map_greater_or_equal.append(2, ()).unwrap();
403 1 : map_greater_or_equal.append(2, ()).unwrap();
404 1 :
405 1 : left.extend(&mut map_greater_or_equal).unwrap_err();
406 1 : assert_eq!(left.as_slice(), &[(0, ()), (1, ())]);
407 1 : assert_eq!(map_greater_or_equal.as_slice(), &[(2, ()), (2, ())]);
408 1 : }
409 :
410 : #[test]
411 1 : fn extend_with_ordering() {
412 1 : let mut left = VecMap::new(VecMapOrdering::GreaterOrEqual);
413 1 : left.append(0, ()).unwrap();
414 1 : assert_eq!(left.as_slice(), &[(0, ())]);
415 :
416 1 : let mut greater_right = VecMap::new(VecMapOrdering::Greater);
417 1 : greater_right.append(0, ()).unwrap();
418 1 : left.extend(&mut greater_right).unwrap_err();
419 1 : assert_eq!(left.as_slice(), &[(0, ())]);
420 :
421 1 : let mut greater_or_equal_right = VecMap::new(VecMapOrdering::GreaterOrEqual);
422 1 : greater_or_equal_right.append(2, ()).unwrap();
423 1 : greater_or_equal_right.append(2, ()).unwrap();
424 1 : left.extend(&mut greater_or_equal_right).unwrap();
425 1 : assert_eq!(left.as_slice(), &[(0, ()), (2, ()), (2, ())]);
426 1 : }
427 :
428 : #[test]
429 1 : fn vec_map_from_sorted() {
430 1 : let vec = vec![(1, ()), (2, ()), (3, ()), (6, ())];
431 1 : let vec_map = VecMap::from_iter(vec, VecMapOrdering::Greater);
432 1 : assert_eq!(vec_map.as_slice(), &[(1, ()), (2, ()), (3, ()), (6, ())]);
433 :
434 1 : let vec = vec![(1, ()), (2, ()), (3, ()), (3, ()), (6, ()), (6, ())];
435 1 : let vec_map = VecMap::from_iter(vec, VecMapOrdering::GreaterOrEqual);
436 1 : assert_eq!(
437 1 : vec_map.as_slice(),
438 1 : &[(1, ()), (2, ()), (3, ()), (3, ()), (6, ()), (6, ())]
439 1 : );
440 1 : }
441 :
442 : #[test]
443 : #[should_panic]
444 1 : fn vec_map_from_unsorted_greater() {
445 1 : let vec = vec![(1, ()), (2, ()), (2, ()), (3, ()), (6, ())];
446 1 : let _ = VecMap::from_iter(vec, VecMapOrdering::Greater);
447 1 : }
448 :
449 : #[test]
450 : #[should_panic]
451 1 : fn vec_map_from_unsorted_greater_or_equal() {
452 1 : let vec = vec![(1, ()), (2, ()), (3, ()), (6, ()), (5, ())];
453 1 : let _ = VecMap::from_iter(vec, VecMapOrdering::GreaterOrEqual);
454 1 : }
455 : }
|