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 14234725 : fn default() -> Self {
22 14234725 : VecMap {
23 14234725 : data: Default::default(),
24 14234725 : ordering: VecMapOrdering::Greater,
25 14234725 : }
26 14234725 : }
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 24 : pub fn new(ordering: VecMapOrdering) -> Self {
39 24 : Self {
40 24 : data: Vec::new(),
41 24 : ordering,
42 24 : }
43 24 : }
44 :
45 24 : pub fn with_capacity(capacity: usize, ordering: VecMapOrdering) -> Self {
46 24 : Self {
47 24 : data: Vec::with_capacity(capacity),
48 24 : ordering,
49 24 : }
50 24 : }
51 :
52 0 : pub fn is_empty(&self) -> bool {
53 0 : self.data.is_empty()
54 0 : }
55 :
56 26155804 : pub fn as_slice(&self) -> &[(K, V)] {
57 26155804 : self.data.as_slice()
58 26155804 : }
59 :
60 : /// This function may panic if given a range where the lower bound is
61 : /// greater than the upper bound.
62 1497910 : pub fn slice_range<R: RangeBounds<K>>(&self, range: R) -> &[(K, V)] {
63 1497910 : use std::ops::Bound::*;
64 1497910 :
65 2995544 : let binary_search = |k: &K| self.data.binary_search_by_key(&k, extract_key);
66 :
67 1497910 : let start_idx = match range.start_bound() {
68 162 : Unbounded => 0,
69 1497250 : Included(k) => binary_search(k).unwrap_or_else(std::convert::identity),
70 498 : Excluded(k) => match binary_search(k) {
71 216 : Ok(idx) => idx + 1,
72 282 : Err(idx) => idx,
73 : },
74 : };
75 :
76 1497910 : let end_idx = match range.end_bound() {
77 114 : Unbounded => self.data.len(),
78 594 : Included(k) => match binary_search(k) {
79 264 : Ok(idx) => idx + 1,
80 330 : Err(idx) => idx,
81 : },
82 1497202 : Excluded(k) => binary_search(k).unwrap_or_else(std::convert::identity),
83 : };
84 :
85 1497910 : &self.data[start_idx..end_idx]
86 1497910 : }
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 7858250 : pub fn append(&mut self, key: K, value: V) -> Result<usize, VecMapError> {
92 7858250 : self.validate_key_order(&key)?;
93 :
94 7858238 : let delta_size = self.instrument_vec_op(|vec| vec.push((key, value)));
95 7858238 : Ok(delta_size)
96 7858250 : }
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 15271446 : pub fn append_or_update_last(
102 15271446 : &mut self,
103 15271446 : key: K,
104 15271446 : mut value: V,
105 15271446 : ) -> Result<(Option<V>, usize), VecMapError> {
106 15271446 : if let Some((last_key, last_value)) = self.data.last_mut() {
107 1660449 : 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 1660449 : Ordering::Greater => {}
115 : }
116 13610997 : }
117 :
118 15271446 : let delta_size = self.instrument_vec_op(|vec| vec.push((key, value)));
119 15271446 : Ok((None, delta_size))
120 15271446 : }
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 42 : pub fn extend(&mut self, other: &mut Self) -> Result<usize, VecMapError> {
154 42 : if self.ordering != other.ordering {
155 12 : return Err(VecMapError::ExtendOrderingError);
156 30 : }
157 30 :
158 30 : let other_first_opt = other.data.last().map(extract_key);
159 30 : if let Some(other_first) = other_first_opt {
160 24 : self.validate_key_order(other_first)?;
161 6 : }
162 :
163 18 : let delta_size = self.instrument_vec_op(|vec| vec.append(&mut other.data));
164 18 : Ok(delta_size)
165 42 : }
166 :
167 : /// Validate the current last key in `self` and key being
168 : /// inserted against the order defined in `self`.
169 7858274 : fn validate_key_order(&self, key: &K) -> Result<(), VecMapError> {
170 7858274 : if let Some(last_key) = self.data.last().map(extract_key) {
171 7234504 : match (&self.ordering, &key.cmp(last_key)) {
172 : (VecMapOrdering::Greater, Ordering::Less | Ordering::Equal) => {
173 18 : return Err(VecMapError::InvalidKey);
174 : }
175 7234414 : (VecMapOrdering::Greater, Ordering::Greater) => {}
176 : (VecMapOrdering::GreaterOrEqual, Ordering::Less) => {
177 6 : return Err(VecMapError::InvalidKey);
178 : }
179 66 : (VecMapOrdering::GreaterOrEqual, Ordering::Equal | Ordering::Greater) => {}
180 : }
181 623770 : }
182 :
183 7858250 : Ok(())
184 7858274 : }
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 23129702 : fn instrument_vec_op(&mut self, op: impl FnOnce(&mut Vec<(K, V)>)) -> usize {
190 23129702 : let old_cap = self.data.capacity();
191 23129702 : op(&mut self.data);
192 23129702 : let new_cap = self.data.capacity();
193 23129702 :
194 23129702 : match old_cap.cmp(&new_cap) {
195 : Ordering::Less => {
196 14345829 : let old_size = Layout::array::<(K, V)>(old_cap).unwrap().size();
197 14345829 : let new_size = Layout::array::<(K, V)>(new_cap).unwrap().size();
198 14345829 : new_size - old_size
199 : }
200 8783873 : Ordering::Equal => 0,
201 0 : Ordering::Greater => panic!("VecMap capacity shouldn't ever decrease"),
202 : }
203 23129702 : }
204 :
205 : /// Similar to `from_iter` defined in `FromIter` trait except
206 : /// that it accepts an [`VecMapOrdering`]
207 24 : pub fn from_iter<I: IntoIterator<Item = (K, V)>>(iter: I, ordering: VecMapOrdering) -> Self {
208 24 : let iter = iter.into_iter();
209 24 : let initial_capacity = {
210 24 : match iter.size_hint() {
211 0 : (lower_bound, None) => lower_bound,
212 24 : (_, Some(upper_bound)) => upper_bound,
213 : }
214 : };
215 :
216 24 : let mut vec_map = VecMap::with_capacity(initial_capacity, ordering);
217 132 : for (key, value) in iter {
218 108 : vec_map
219 108 : .append(key, value)
220 108 : .expect("The passed collection needs to be sorted!");
221 108 : }
222 :
223 24 : vec_map
224 24 : }
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 0 : fn into_iter(self) -> Self::IntoIter {
232 0 : self.data.into_iter()
233 0 : }
234 : }
235 :
236 34137653 : fn extract_key<K, V>(entry: &(K, V)) -> &K {
237 34137653 : &entry.0
238 34137653 : }
239 :
240 : #[cfg(test)]
241 : mod tests {
242 : use std::{collections::BTreeMap, ops::Bound};
243 :
244 : use super::{VecMap, VecMapOrdering};
245 :
246 : #[test]
247 6 : fn unbounded_range() {
248 6 : let mut vec = VecMap::default();
249 6 : vec.append(0, ()).unwrap();
250 6 :
251 6 : assert_eq!(vec.slice_range(0..0), &[]);
252 6 : }
253 :
254 : #[test]
255 : #[should_panic]
256 6 : fn invalid_ordering_range() {
257 6 : let mut vec = VecMap::default();
258 6 : vec.append(0, ()).unwrap();
259 6 :
260 6 : #[allow(clippy::reversed_empty_ranges)]
261 6 : vec.slice_range(1..0);
262 6 : }
263 :
264 : #[test]
265 6 : fn range_tests() {
266 6 : let mut vec = VecMap::default();
267 6 : vec.append(0, ()).unwrap();
268 6 : vec.append(2, ()).unwrap();
269 6 : vec.append(4, ()).unwrap();
270 6 :
271 6 : assert_eq!(vec.slice_range(0..0), &[]);
272 6 : assert_eq!(vec.slice_range(0..1), &[(0, ())]);
273 6 : assert_eq!(vec.slice_range(0..2), &[(0, ())]);
274 6 : assert_eq!(vec.slice_range(0..3), &[(0, ()), (2, ())]);
275 :
276 6 : assert_eq!(vec.slice_range(..0), &[]);
277 6 : assert_eq!(vec.slice_range(..1), &[(0, ())]);
278 :
279 6 : assert_eq!(vec.slice_range(..3), &[(0, ()), (2, ())]);
280 6 : assert_eq!(vec.slice_range(..3), &[(0, ()), (2, ())]);
281 :
282 6 : assert_eq!(vec.slice_range(0..=0), &[(0, ())]);
283 6 : assert_eq!(vec.slice_range(0..=1), &[(0, ())]);
284 6 : assert_eq!(vec.slice_range(0..=2), &[(0, ()), (2, ())]);
285 6 : assert_eq!(vec.slice_range(0..=3), &[(0, ()), (2, ())]);
286 :
287 6 : assert_eq!(vec.slice_range(..=0), &[(0, ())]);
288 6 : assert_eq!(vec.slice_range(..=1), &[(0, ())]);
289 6 : assert_eq!(vec.slice_range(..=2), &[(0, ()), (2, ())]);
290 6 : assert_eq!(vec.slice_range(..=3), &[(0, ()), (2, ())]);
291 6 : }
292 :
293 : struct BoundIter {
294 : min: i32,
295 : max: i32,
296 :
297 : next: Option<Bound<i32>>,
298 : }
299 :
300 : impl BoundIter {
301 120 : fn new(min: i32, max: i32) -> Self {
302 120 : Self {
303 120 : min,
304 120 : max,
305 120 :
306 120 : next: Some(Bound::Unbounded),
307 120 : }
308 120 : }
309 : }
310 :
311 : impl Iterator for BoundIter {
312 : type Item = Bound<i32>;
313 :
314 1440 : fn next(&mut self) -> Option<Self::Item> {
315 1440 : let cur = self.next?;
316 :
317 1320 : self.next = match &cur {
318 120 : Bound::Unbounded => Some(Bound::Included(self.min)),
319 600 : Bound::Included(x) => {
320 600 : if *x >= self.max {
321 120 : Some(Bound::Excluded(self.min))
322 : } else {
323 480 : Some(Bound::Included(x + 1))
324 : }
325 : }
326 600 : Bound::Excluded(x) => {
327 600 : if *x >= self.max {
328 120 : None
329 : } else {
330 480 : Some(Bound::Excluded(x + 1))
331 : }
332 : }
333 : };
334 :
335 1320 : Some(cur)
336 1440 : }
337 : }
338 :
339 : #[test]
340 6 : fn range_exhaustive() {
341 24 : let map: BTreeMap<i32, ()> = (1..=7).step_by(2).map(|x| (x, ())).collect();
342 6 : let mut vec = VecMap::default();
343 24 : for &key in map.keys() {
344 24 : vec.append(key, ()).unwrap();
345 24 : }
346 :
347 : const RANGE_MIN: i32 = 0;
348 : const RANGE_MAX: i32 = 8;
349 120 : for lower_bound in BoundIter::new(RANGE_MIN, RANGE_MAX) {
350 114 : let ub_min = match lower_bound {
351 6 : Bound::Unbounded => RANGE_MIN,
352 54 : Bound::Included(x) => x,
353 54 : Bound::Excluded(x) => x + 1,
354 : };
355 1206 : for upper_bound in BoundIter::new(ub_min, RANGE_MAX) {
356 1206 : let map_range: Vec<(i32, ())> = map
357 1206 : .range((lower_bound, upper_bound))
358 1920 : .map(|(&x, _)| (x, ()))
359 1206 : .collect();
360 1206 : let vec_slice = vec.slice_range((lower_bound, upper_bound));
361 1206 :
362 1206 : assert_eq!(map_range, vec_slice);
363 : }
364 : }
365 6 : }
366 :
367 : #[test]
368 6 : fn extend() {
369 6 : let mut left = VecMap::default();
370 6 : left.append(0, ()).unwrap();
371 6 : assert_eq!(left.as_slice(), &[(0, ())]);
372 :
373 6 : let mut empty = VecMap::default();
374 6 : left.extend(&mut empty).unwrap();
375 6 : assert_eq!(left.as_slice(), &[(0, ())]);
376 6 : assert_eq!(empty.as_slice(), &[]);
377 :
378 6 : let mut right = VecMap::default();
379 6 : right.append(1, ()).unwrap();
380 6 :
381 6 : left.extend(&mut right).unwrap();
382 6 :
383 6 : assert_eq!(left.as_slice(), &[(0, ()), (1, ())]);
384 6 : assert_eq!(right.as_slice(), &[]);
385 :
386 6 : let mut zero_map = VecMap::default();
387 6 : zero_map.append(0, ()).unwrap();
388 6 :
389 6 : left.extend(&mut zero_map).unwrap_err();
390 6 : assert_eq!(left.as_slice(), &[(0, ()), (1, ())]);
391 6 : assert_eq!(zero_map.as_slice(), &[(0, ())]);
392 :
393 6 : let mut one_map = VecMap::default();
394 6 : one_map.append(1, ()).unwrap();
395 6 :
396 6 : left.extend(&mut one_map).unwrap_err();
397 6 : assert_eq!(left.as_slice(), &[(0, ()), (1, ())]);
398 6 : assert_eq!(one_map.as_slice(), &[(1, ())]);
399 :
400 6 : let mut map_greater_or_equal = VecMap::new(VecMapOrdering::GreaterOrEqual);
401 6 : map_greater_or_equal.append(2, ()).unwrap();
402 6 : map_greater_or_equal.append(2, ()).unwrap();
403 6 :
404 6 : left.extend(&mut map_greater_or_equal).unwrap_err();
405 6 : assert_eq!(left.as_slice(), &[(0, ()), (1, ())]);
406 6 : assert_eq!(map_greater_or_equal.as_slice(), &[(2, ()), (2, ())]);
407 6 : }
408 :
409 : #[test]
410 6 : fn extend_with_ordering() {
411 6 : let mut left = VecMap::new(VecMapOrdering::GreaterOrEqual);
412 6 : left.append(0, ()).unwrap();
413 6 : assert_eq!(left.as_slice(), &[(0, ())]);
414 :
415 6 : let mut greater_right = VecMap::new(VecMapOrdering::Greater);
416 6 : greater_right.append(0, ()).unwrap();
417 6 : left.extend(&mut greater_right).unwrap_err();
418 6 : assert_eq!(left.as_slice(), &[(0, ())]);
419 :
420 6 : let mut greater_or_equal_right = VecMap::new(VecMapOrdering::GreaterOrEqual);
421 6 : greater_or_equal_right.append(2, ()).unwrap();
422 6 : greater_or_equal_right.append(2, ()).unwrap();
423 6 : left.extend(&mut greater_or_equal_right).unwrap();
424 6 : assert_eq!(left.as_slice(), &[(0, ()), (2, ()), (2, ())]);
425 6 : }
426 :
427 : #[test]
428 6 : fn vec_map_from_sorted() {
429 6 : let vec = vec![(1, ()), (2, ()), (3, ()), (6, ())];
430 6 : let vec_map = VecMap::from_iter(vec, VecMapOrdering::Greater);
431 6 : assert_eq!(vec_map.as_slice(), &[(1, ()), (2, ()), (3, ()), (6, ())]);
432 :
433 6 : let vec = vec![(1, ()), (2, ()), (3, ()), (3, ()), (6, ()), (6, ())];
434 6 : let vec_map = VecMap::from_iter(vec, VecMapOrdering::GreaterOrEqual);
435 6 : assert_eq!(
436 6 : vec_map.as_slice(),
437 6 : &[(1, ()), (2, ()), (3, ()), (3, ()), (6, ()), (6, ())]
438 6 : );
439 6 : }
440 :
441 : #[test]
442 : #[should_panic]
443 6 : fn vec_map_from_unsorted_greater() {
444 6 : let vec = vec![(1, ()), (2, ()), (2, ()), (3, ()), (6, ())];
445 6 : let _ = VecMap::from_iter(vec, VecMapOrdering::Greater);
446 6 : }
447 :
448 : #[test]
449 : #[should_panic]
450 6 : fn vec_map_from_unsorted_greater_or_equal() {
451 6 : let vec = vec![(1, ()), (2, ()), (3, ()), (6, ()), (5, ())];
452 6 : let _ = VecMap::from_iter(vec, VecMapOrdering::GreaterOrEqual);
453 6 : }
454 : }
|