TLA Line data Source code
1 : use std::{alloc::Layout, cmp::Ordering, ops::RangeBounds};
2 :
3 : /// Ordered map datastructure implemented in a Vec.
4 : /// Append only - can only add keys that are larger than the
5 : /// current max key.
6 UBC 0 : #[derive(Clone, Debug)]
7 : pub struct VecMap<K, V>(Vec<(K, V)>);
8 :
9 : impl<K, V> Default for VecMap<K, V> {
10 CBC 4999133 : fn default() -> Self {
11 4999133 : VecMap(Default::default())
12 4999133 : }
13 : }
14 :
15 UBC 0 : #[derive(Debug)]
16 : pub struct InvalidKey;
17 :
18 : impl<K: Ord, V> VecMap<K, V> {
19 0 : pub fn is_empty(&self) -> bool {
20 0 : self.0.is_empty()
21 0 : }
22 :
23 CBC 9204007 : pub fn as_slice(&self) -> &[(K, V)] {
24 9204007 : self.0.as_slice()
25 9204007 : }
26 :
27 : /// This function may panic if given a range where the lower bound is
28 : /// greater than the upper bound.
29 3157216 : pub fn slice_range<R: RangeBounds<K>>(&self, range: R) -> &[(K, V)] {
30 3157216 : use std::ops::Bound::*;
31 3157216 :
32 6314386 : let binary_search = |k: &K| self.0.binary_search_by_key(&k, extract_key);
33 :
34 3157216 : let start_idx = match range.start_bound() {
35 27 : Unbounded => 0,
36 3157106 : Included(k) => binary_search(k).unwrap_or_else(std::convert::identity),
37 83 : Excluded(k) => match binary_search(k) {
38 36 : Ok(idx) => idx + 1,
39 47 : Err(idx) => idx,
40 : },
41 : };
42 :
43 3157216 : let end_idx = match range.end_bound() {
44 19 : Unbounded => self.0.len(),
45 99 : Included(k) => match binary_search(k) {
46 44 : Ok(idx) => idx + 1,
47 55 : Err(idx) => idx,
48 : },
49 3157098 : Excluded(k) => binary_search(k).unwrap_or_else(std::convert::identity),
50 : };
51 :
52 3157216 : &self.0[start_idx..end_idx]
53 3157216 : }
54 :
55 : /// Add a key value pair to the map.
56 : /// If `key` is less than or equal to the current maximum key
57 : /// the pair will not be added and InvalidKey error will be returned.
58 : pub fn append(&mut self, key: K, value: V) -> Result<usize, InvalidKey> {
59 13 : if let Some((last_key, _last_value)) = self.0.last() {
60 5 : if &key <= last_key {
61 UBC 0 : return Err(InvalidKey);
62 CBC 5 : }
63 8 : }
64 :
65 13 : let delta_size = self.instrument_vec_op(|vec| vec.push((key, value)));
66 13 : Ok(delta_size)
67 13 : }
68 :
69 : /// Update the maximum key value pair or add a new key value pair to the map.
70 : /// If `key` is less than the current maximum key no updates or additions
71 : /// will occur and InvalidKey error will be returned.
72 : pub fn append_or_update_last(
73 : &mut self,
74 : key: K,
75 : mut value: V,
76 : ) -> Result<(Option<V>, usize), InvalidKey> {
77 76621875 : if let Some((last_key, last_value)) = self.0.last_mut() {
78 71622751 : match key.cmp(last_key) {
79 UBC 0 : Ordering::Less => return Err(InvalidKey),
80 : Ordering::Equal => {
81 0 : std::mem::swap(last_value, &mut value);
82 0 : const DELTA_SIZE: usize = 0;
83 0 : return Ok((Some(value), DELTA_SIZE));
84 : }
85 CBC 71622751 : Ordering::Greater => {}
86 : }
87 4999124 : }
88 :
89 76621875 : let delta_size = self.instrument_vec_op(|vec| vec.push((key, value)));
90 76621875 : Ok((None, delta_size))
91 76621875 : }
92 :
93 : /// Split the map into two.
94 : ///
95 : /// The left map contains everything before `cutoff` (exclusive).
96 : /// Right map contains `cutoff` and everything after (inclusive).
97 UBC 0 : pub fn split_at(&self, cutoff: &K) -> (Self, Self)
98 0 : where
99 0 : K: Clone,
100 0 : V: Clone,
101 0 : {
102 0 : let split_idx = self
103 0 : .0
104 0 : .binary_search_by_key(&cutoff, extract_key)
105 0 : .unwrap_or_else(std::convert::identity);
106 0 :
107 0 : (
108 0 : VecMap(self.0[..split_idx].to_vec()),
109 0 : VecMap(self.0[split_idx..].to_vec()),
110 0 : )
111 0 : }
112 :
113 : /// Move items from `other` to the end of `self`, leaving `other` empty.
114 : /// If any keys in `other` is less than or equal to any key in `self`,
115 : /// `InvalidKey` error will be returned and no mutation will occur.
116 CBC 4 : pub fn extend(&mut self, other: &mut Self) -> Result<usize, InvalidKey> {
117 4 : let self_last_opt = self.0.last().map(extract_key);
118 4 : let other_first_opt = other.0.last().map(extract_key);
119 :
120 4 : if let (Some(self_last), Some(other_first)) = (self_last_opt, other_first_opt) {
121 3 : if self_last >= other_first {
122 2 : return Err(InvalidKey);
123 1 : }
124 1 : }
125 :
126 2 : let delta_size = self.instrument_vec_op(|vec| vec.append(&mut other.0));
127 2 : Ok(delta_size)
128 4 : }
129 :
130 : /// Instrument an operation on the underlying [`Vec`].
131 : /// Will panic if the operation decreases capacity.
132 : /// Returns the increase in memory usage caused by the op.
133 76621890 : fn instrument_vec_op(&mut self, op: impl FnOnce(&mut Vec<(K, V)>)) -> usize {
134 76621890 : let old_cap = self.0.capacity();
135 76621890 : op(&mut self.0);
136 76621890 : let new_cap = self.0.capacity();
137 76621890 :
138 76621890 : match old_cap.cmp(&new_cap) {
139 : Ordering::Less => {
140 7599205 : let old_size = Layout::array::<(K, V)>(old_cap).unwrap().size();
141 7599205 : let new_size = Layout::array::<(K, V)>(new_cap).unwrap().size();
142 7599205 : new_size - old_size
143 : }
144 69022685 : Ordering::Equal => 0,
145 UBC 0 : Ordering::Greater => panic!("VecMap capacity shouldn't ever decrease"),
146 : }
147 CBC 76621890 : }
148 : }
149 :
150 32250894 : fn extract_key<K, V>(entry: &(K, V)) -> &K {
151 32250894 : &entry.0
152 32250894 : }
153 :
154 : #[cfg(test)]
155 : mod tests {
156 : use std::{collections::BTreeMap, ops::Bound};
157 :
158 : use super::VecMap;
159 :
160 1 : #[test]
161 1 : fn unbounded_range() {
162 1 : let mut vec = VecMap::default();
163 1 : vec.append(0, ()).unwrap();
164 1 :
165 1 : assert_eq!(vec.slice_range(0..0), &[]);
166 1 : }
167 :
168 1 : #[test]
169 : #[should_panic]
170 1 : fn invalid_ordering_range() {
171 1 : let mut vec = VecMap::default();
172 1 : vec.append(0, ()).unwrap();
173 1 :
174 1 : #[allow(clippy::reversed_empty_ranges)]
175 1 : vec.slice_range(1..0);
176 1 : }
177 :
178 1 : #[test]
179 1 : fn range_tests() {
180 1 : let mut vec = VecMap::default();
181 1 : vec.append(0, ()).unwrap();
182 1 : vec.append(2, ()).unwrap();
183 1 : vec.append(4, ()).unwrap();
184 1 :
185 1 : assert_eq!(vec.slice_range(0..0), &[]);
186 1 : assert_eq!(vec.slice_range(0..1), &[(0, ())]);
187 1 : assert_eq!(vec.slice_range(0..2), &[(0, ())]);
188 1 : assert_eq!(vec.slice_range(0..3), &[(0, ()), (2, ())]);
189 :
190 1 : assert_eq!(vec.slice_range(..0), &[]);
191 1 : assert_eq!(vec.slice_range(..1), &[(0, ())]);
192 :
193 1 : assert_eq!(vec.slice_range(..3), &[(0, ()), (2, ())]);
194 1 : assert_eq!(vec.slice_range(..3), &[(0, ()), (2, ())]);
195 :
196 1 : assert_eq!(vec.slice_range(0..=0), &[(0, ())]);
197 1 : assert_eq!(vec.slice_range(0..=1), &[(0, ())]);
198 1 : assert_eq!(vec.slice_range(0..=2), &[(0, ()), (2, ())]);
199 1 : assert_eq!(vec.slice_range(0..=3), &[(0, ()), (2, ())]);
200 :
201 1 : assert_eq!(vec.slice_range(..=0), &[(0, ())]);
202 1 : assert_eq!(vec.slice_range(..=1), &[(0, ())]);
203 1 : assert_eq!(vec.slice_range(..=2), &[(0, ()), (2, ())]);
204 1 : assert_eq!(vec.slice_range(..=3), &[(0, ()), (2, ())]);
205 1 : }
206 :
207 : struct BoundIter {
208 : min: i32,
209 : max: i32,
210 :
211 : next: Option<Bound<i32>>,
212 : }
213 :
214 : impl BoundIter {
215 20 : fn new(min: i32, max: i32) -> Self {
216 20 : Self {
217 20 : min,
218 20 : max,
219 20 :
220 20 : next: Some(Bound::Unbounded),
221 20 : }
222 20 : }
223 : }
224 :
225 : impl Iterator for BoundIter {
226 : type Item = Bound<i32>;
227 :
228 240 : fn next(&mut self) -> Option<Self::Item> {
229 240 : let cur = self.next?;
230 :
231 220 : self.next = match &cur {
232 20 : Bound::Unbounded => Some(Bound::Included(self.min)),
233 100 : Bound::Included(x) => {
234 100 : if *x >= self.max {
235 20 : Some(Bound::Excluded(self.min))
236 : } else {
237 80 : Some(Bound::Included(x + 1))
238 : }
239 : }
240 100 : Bound::Excluded(x) => {
241 100 : if *x >= self.max {
242 20 : None
243 : } else {
244 80 : Some(Bound::Excluded(x + 1))
245 : }
246 : }
247 : };
248 :
249 220 : Some(cur)
250 240 : }
251 : }
252 :
253 1 : #[test]
254 1 : fn range_exhaustive() {
255 4 : let map: BTreeMap<i32, ()> = (1..=7).step_by(2).map(|x| (x, ())).collect();
256 1 : let mut vec = VecMap::default();
257 4 : for &key in map.keys() {
258 4 : vec.append(key, ()).unwrap();
259 4 : }
260 :
261 : const RANGE_MIN: i32 = 0;
262 : const RANGE_MAX: i32 = 8;
263 20 : for lower_bound in BoundIter::new(RANGE_MIN, RANGE_MAX) {
264 19 : let ub_min = match lower_bound {
265 1 : Bound::Unbounded => RANGE_MIN,
266 9 : Bound::Included(x) => x,
267 9 : Bound::Excluded(x) => x + 1,
268 : };
269 201 : for upper_bound in BoundIter::new(ub_min, RANGE_MAX) {
270 201 : let map_range: Vec<(i32, ())> = map
271 201 : .range((lower_bound, upper_bound))
272 320 : .map(|(&x, _)| (x, ()))
273 201 : .collect();
274 201 : let vec_slice = vec.slice_range((lower_bound, upper_bound));
275 201 :
276 201 : assert_eq!(map_range, vec_slice);
277 : }
278 : }
279 1 : }
280 :
281 1 : #[test]
282 1 : fn extend() {
283 1 : let mut left = VecMap::default();
284 1 : left.append(0, ()).unwrap();
285 1 : assert_eq!(left.as_slice(), &[(0, ())]);
286 :
287 1 : let mut empty = VecMap::default();
288 1 : left.extend(&mut empty).unwrap();
289 1 : assert_eq!(left.as_slice(), &[(0, ())]);
290 1 : assert_eq!(empty.as_slice(), &[]);
291 :
292 1 : let mut right = VecMap::default();
293 1 : right.append(1, ()).unwrap();
294 1 :
295 1 : left.extend(&mut right).unwrap();
296 1 :
297 1 : assert_eq!(left.as_slice(), &[(0, ()), (1, ())]);
298 1 : assert_eq!(right.as_slice(), &[]);
299 :
300 1 : let mut zero_map = VecMap::default();
301 1 : zero_map.append(0, ()).unwrap();
302 1 :
303 1 : left.extend(&mut zero_map).unwrap_err();
304 1 : assert_eq!(left.as_slice(), &[(0, ()), (1, ())]);
305 1 : assert_eq!(zero_map.as_slice(), &[(0, ())]);
306 :
307 1 : let mut one_map = VecMap::default();
308 1 : one_map.append(1, ()).unwrap();
309 1 :
310 1 : left.extend(&mut one_map).unwrap_err();
311 1 : assert_eq!(left.as_slice(), &[(0, ()), (1, ())]);
312 1 : assert_eq!(one_map.as_slice(), &[(1, ())]);
313 1 : }
314 : }
|
