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