1use crate::vecset::VecSet;
2
3use core::borrow::Borrow;
4use core::cmp::Ordering;
5use core::fmt;
6use core::mem;
7use core::ptr;
8use core::slice;
9
10use alloc::vec;
11use alloc::vec::Vec;
12
13#[derive(Clone, PartialEq, Eq, Hash)]
14pub struct VecMap<K, V>(Vec<(K, V)>);
15
16impl<K, V> VecMap<K, V> {
17 #[inline]
18 #[must_use]
19 pub const fn new() -> Self {
20 Self(Vec::new())
21 }
22
23 #[inline]
24 #[must_use]
25 pub fn from_single(key: K, value: V) -> Self {
26 Self(vec![(key, value)])
27 }
28
29 #[inline]
30 #[must_use]
31 pub fn with_capacity(cap: usize) -> Self {
32 Self(Vec::with_capacity(cap))
33 }
34
35 #[inline]
36 #[must_use]
37 pub fn len(&self) -> usize {
38 self.0.len()
39 }
40
41 #[inline]
42 #[must_use]
43 pub fn is_empty(&self) -> bool {
44 self.0.is_empty()
45 }
46
47 #[inline]
48 #[must_use]
49 pub fn iter(&self) -> Iter<'_, K, V> {
50 Iter(self.0.as_slice().iter())
51 }
52
53 #[inline]
54 #[must_use]
55 pub fn iter_mut(&mut self) -> IterMut<'_, K, V> {
56 IterMut(self.0.as_mut_slice().iter_mut())
57 }
58
59 unsafe fn at_unchecked(&self, idx: usize) -> &(K, V) {
60 self.0.get_unchecked(idx)
61 }
62
63 unsafe fn at_unchecked_mut(&mut self, idx: usize) -> &mut (K, V) {
64 self.0.get_unchecked_mut(idx)
65 }
66}
67
68impl<K: Ord, V> VecMap<K, V> {
69 #[inline]
70 #[must_use]
71 pub fn from_vec(mut v: Vec<(K, V)>) -> Self {
72 v.sort_unstable_by(|lhs, rhs| lhs.0.cmp(&rhs.0));
73 v.dedup_by(|x, first| x.0 == first.0);
74 Self(v)
75 }
76
77 fn search<Q>(&self, key: &Q) -> Result<usize, usize>
78 where
79 K: Borrow<Q>,
80 Q: Ord + ?Sized,
81 {
82 self.0.binary_search_by(|probe| probe.0.borrow().cmp(key))
83 }
84
85 #[inline]
86 #[must_use]
87 pub fn contains_key<Q>(&self, key: &Q) -> bool
88 where
89 K: Borrow<Q>,
90 Q: Ord + ?Sized,
91 {
92 self.search(key).is_ok()
93 }
94
95 #[inline]
96 #[must_use]
97 pub fn get<Q>(&self, key: &Q) -> Option<&V>
98 where
99 K: Borrow<Q>,
100 Q: Ord + ?Sized,
101 {
102 let idx = self.search(key).ok()?;
103 let entry = unsafe { self.at_unchecked(idx) };
104 Some(&entry.1)
105 }
106
107 #[inline]
108 #[must_use]
109 pub fn get_mut<Q>(&mut self, key: &Q) -> Option<&mut V>
110 where
111 K: Borrow<Q>,
112 Q: Ord + ?Sized,
113 {
114 let idx = self.search(key).ok()?;
115 let entry = unsafe { self.at_unchecked_mut(idx) };
116 Some(&mut entry.1)
117 }
118
119 #[inline]
120 pub fn insert(&mut self, key: K, value: V) -> Option<V> {
121 match self.search(&key) {
122 Ok(idx) => {
123 let entry = unsafe { self.at_unchecked_mut(idx) };
124 Some(mem::replace(&mut entry.1, value))
125 }
126 Err(idx) => {
127 self.0.insert(idx, (key, value));
128 None
129 }
130 }
131 }
132
133 #[inline]
134 pub fn remove<Q>(&mut self, key: &Q) -> Option<V>
135 where
136 K: Borrow<Q>,
137 Q: Ord + ?Sized,
138 {
139 let idx = self.search(key).ok()?;
140 let entry = self.0.remove(idx);
141 Some(entry.1)
142 }
143
144 #[inline]
145 pub fn entry(&mut self, key: K) -> Entry<'_, K, V> {
146 match self.search(&key) {
147 Ok(idx) => Entry::Occupied(OccupiedEntry { map: self, idx }),
148 Err(idx) => Entry::Vacant(VacantEntry {
149 map: self,
150 idx,
151 key,
152 }),
153 }
154 }
155
156 #[inline]
157 pub fn merge_copied_with(&mut self, other: &Self, mut f: impl FnMut(V, V) -> V)
158 where
159 K: Copy,
160 V: Copy,
161 {
162 let lhs = &mut self.0;
163 let rhs = &other.0;
164
165 let ans_cap = lhs.len().checked_add(rhs.len()).unwrap();
166 lhs.reserve(ans_cap);
167
168 unsafe {
169 let mut p1 = lhs.as_ptr();
170 let mut p2 = rhs.as_ptr();
171 let mut p3 = lhs.as_mut_ptr().add(lhs.len());
172 let e1 = p1.add(lhs.len());
173 let e2 = p2.add(rhs.len());
174
175 while p1 < e1 && p2 < e2 {
176 let (k1, v1) = &*p1;
177 let (k2, v2) = &*p2;
178 match Ord::cmp(k1, k2) {
179 Ordering::Less => {
180 ptr::copy_nonoverlapping(p1, p3, 1);
181 p1 = p1.add(1);
182 }
183 Ordering::Greater => {
184 ptr::copy_nonoverlapping(p2, p3, 1);
185 p2 = p2.add(1);
186 }
187 Ordering::Equal => {
188 let v = f(*v1, *v2);
189 p3.write((*k1, v));
190 p1 = p1.add(1);
191 p2 = p2.add(1);
192 }
193 }
194 p3 = p3.add(1);
195 }
196 if p1 < e1 {
197 let cnt = e1.offset_from(p1) as usize;
198 ptr::copy_nonoverlapping(p1, p3, cnt);
199 p3 = p3.add(cnt);
200 }
201 if p2 < e2 {
202 let cnt = e2.offset_from(p2) as usize;
203 ptr::copy_nonoverlapping(p2, p3, cnt);
204 p3 = p3.add(cnt);
205 }
206 {
207 let dst = lhs.as_mut_ptr();
208 let src = dst.add(lhs.len());
209 let cnt = p3.offset_from(src) as usize;
210 ptr::copy(src, dst, cnt);
211 lhs.set_len(cnt)
212 }
213 }
214 }
215
216 #[inline]
217 pub fn remove_less_than<Q>(&mut self, key: &Q)
218 where
219 K: Borrow<Q>,
220 Q: Ord + ?Sized,
221 {
222 struct Guard<'a, K, V> {
223 v: &'a mut Vec<(K, V)>,
224 remove_cnt: usize,
225 }
226
227 impl<K, V> Drop for Guard<'_, K, V> {
228 fn drop(&mut self) {
229 let v = &mut *self.v;
230 let remove_cnt = self.remove_cnt;
231 let remain_cnt = v.len().wrapping_sub(remove_cnt);
232 unsafe {
233 let dst = v.as_mut_ptr();
234 let src = dst.add(remove_cnt);
235 ptr::copy(src, dst, remain_cnt);
236 v.set_len(remain_cnt)
237 }
238 }
239 }
240
241 let remove_cnt = match self.search(key) {
242 Ok(idx) => idx,
243 Err(idx) => idx,
244 };
245 if remove_cnt == 0 || remove_cnt >= self.0.len() {
246 return;
247 }
248 let guard = Guard {
249 remove_cnt,
250 v: &mut self.0,
251 };
252 unsafe {
253 let entries: *mut [(K, V)] = guard.v.get_unchecked_mut(..remove_cnt);
254 ptr::drop_in_place(entries);
255 }
256 drop(guard);
257 }
258
259 #[inline]
260 #[must_use]
261 pub fn remove_max(&mut self) -> Option<(K, V)> {
262 self.0.pop()
263 }
264
265 #[inline]
266 pub fn apply(&self, keys: &VecSet<K>, mut f: impl FnMut(&V)) {
267 unsafe {
268 let mut p1 = self.0.as_ptr();
269 let e1 = p1.add(self.len());
270 let mut p2 = keys.as_slice().as_ptr();
271 let e2 = p2.add(keys.len());
272
273 while p1 < e1 && p2 < e2 {
274 let (k1, v) = &*p1;
275 let k2 = &*p2;
276 match Ord::cmp(k1, k2) {
277 Ordering::Less => {
278 p1 = p1.add(1);
279 }
280 Ordering::Greater => {
281 p2 = p2.add(1);
282 }
283 Ordering::Equal => {
284 f(v);
285 p1 = p1.add(1);
286 p2 = p2.add(1);
287 }
288 }
289 }
290 }
291 }
292}
293
294impl<K: Ord, V> From<Vec<(K, V)>> for VecMap<K, V> {
295 #[inline]
296 fn from(v: Vec<(K, V)>) -> Self {
297 Self::from_vec(v)
298 }
299}
300
301impl<K: Ord, V> FromIterator<(K, V)> for VecMap<K, V> {
302 #[inline]
303 fn from_iter<I: IntoIterator<Item = (K, V)>>(iter: I) -> Self {
304 Self::from_vec(iter.into_iter().collect())
305 }
306}
307
308impl<K, V> Default for VecMap<K, V> {
309 #[inline]
310 #[must_use]
311 fn default() -> Self {
312 Self::new()
313 }
314}
315
316impl<K, V> fmt::Debug for VecMap<K, V>
317where
318 K: fmt::Debug,
319 V: fmt::Debug,
320{
321 #[inline]
322 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
323 let entries = self.0.iter().map(|(k, v)| (k, v));
324 f.debug_map().entries(entries).finish()
325 }
326}
327
328pub struct Iter<'a, K, V>(slice::Iter<'a, (K, V)>);
329
330impl<'a, K, V> Iterator for Iter<'a, K, V> {
331 type Item = &'a (K, V);
332
333 #[inline]
334 fn next(&mut self) -> Option<Self::Item> {
335 self.0.next()
336 }
337
338 #[inline]
339 fn size_hint(&self) -> (usize, Option<usize>) {
340 self.0.size_hint()
341 }
342}
343impl<'a, K, V> IntoIterator for &'a VecMap<K, V> {
344 type Item = &'a (K, V);
345
346 type IntoIter = Iter<'a, K, V>;
347
348 #[inline]
349 fn into_iter(self) -> Self::IntoIter {
350 self.iter()
351 }
352}
353
354pub struct IterMut<'a, K, V>(slice::IterMut<'a, (K, V)>);
355
356impl<'a, K, V> Iterator for IterMut<'a, K, V> {
357 type Item = &'a mut (K, V);
358
359 #[inline]
360 fn next(&mut self) -> Option<Self::Item> {
361 self.0.next()
362 }
363
364 #[inline]
365 fn size_hint(&self) -> (usize, Option<usize>) {
366 self.0.size_hint()
367 }
368}
369
370impl<'a, K, V> IntoIterator for &'a mut VecMap<K, V> {
371 type Item = &'a mut (K, V);
372
373 type IntoIter = IterMut<'a, K, V>;
374
375 #[inline]
376 fn into_iter(self) -> Self::IntoIter {
377 self.iter_mut()
378 }
379}
380
381pub struct IntoIter<K, V>(vec::IntoIter<(K, V)>);
382
383impl<K, V> IntoIterator for VecMap<K, V> {
384 type Item = (K, V);
385
386 type IntoIter = IntoIter<K, V>;
387
388 #[inline]
389 fn into_iter(self) -> Self::IntoIter {
390 IntoIter(self.0.into_iter())
391 }
392}
393
394impl<K, V> Iterator for IntoIter<K, V> {
395 type Item = (K, V);
396
397 #[inline]
398 fn next(&mut self) -> Option<Self::Item> {
399 self.0.next()
400 }
401
402 #[inline]
403 fn size_hint(&self) -> (usize, Option<usize>) {
404 self.0.size_hint()
405 }
406}
407
408#[must_use]
409pub enum Entry<'a, K, V>
410where
411 K: 'a,
412 V: 'a,
413{
414 Vacant(VacantEntry<'a, K, V>),
415 Occupied(OccupiedEntry<'a, K, V>),
416}
417
418#[must_use]
419pub struct VacantEntry<'a, K, V> {
420 map: &'a mut VecMap<K, V>,
421 idx: usize,
422 key: K,
423}
424
425#[must_use]
426pub struct OccupiedEntry<'a, K, V> {
427 map: &'a mut VecMap<K, V>,
428 idx: usize,
429}
430
431impl<'a, K, V> Entry<'a, K, V> {
432 #[inline]
433 pub fn and_modify(mut self, f: impl FnOnce(&mut V)) -> Self {
434 if let Entry::Occupied(ref mut e) = self {
435 f(e.get_mut())
436 }
437 self
438 }
439
440 #[inline]
441 pub fn key(&self) -> &K {
442 match self {
443 Entry::Vacant(e) => e.key(),
444 Entry::Occupied(e) => e.key(),
445 }
446 }
447
448 #[inline]
449 pub fn or_default(self) -> &'a mut V
450 where
451 V: Default,
452 {
453 self.or_insert_with(V::default)
454 }
455
456 #[inline]
457 pub fn or_insert(self, default: V) -> &'a mut V {
458 match self {
459 Entry::Vacant(e) => e.insert(default),
460 Entry::Occupied(e) => e.into_mut(),
461 }
462 }
463
464 #[inline]
465 pub fn or_insert_with(self, default: impl FnOnce() -> V) -> &'a mut V {
466 match self {
467 Entry::Vacant(e) => e.insert(default()),
468 Entry::Occupied(e) => e.into_mut(),
469 }
470 }
471
472 #[inline]
473 pub fn or_insert_with_key(self, default: impl FnOnce(&K) -> V) -> &'a mut V {
474 match self {
475 Entry::Vacant(e) => {
476 let val = default(e.key());
477 e.insert(val)
478 }
479 Entry::Occupied(e) => e.into_mut(),
480 }
481 }
482}
483
484impl<'a, K, V> VacantEntry<'a, K, V> {
485 #[inline]
486 #[must_use]
487 pub fn key(&self) -> &K {
488 &self.key
489 }
490
491 #[inline]
492 #[must_use]
493 pub fn into_key(self) -> K {
494 self.key
495 }
496
497 #[inline]
498 pub fn insert(self, value: V) -> &'a mut V {
499 self.map.0.insert(self.idx, (self.key, value));
500 let entry = unsafe { self.map.at_unchecked_mut(self.idx) };
501 &mut entry.1
502 }
503}
504
505impl<'a, K, V> OccupiedEntry<'a, K, V> {
506 #[inline]
507 #[must_use]
508 pub fn get(&self) -> &V {
509 let entry = unsafe { self.map.at_unchecked(self.idx) };
510 &entry.1
511 }
512
513 #[inline]
514 #[must_use]
515 pub fn get_mut(&mut self) -> &mut V {
516 let entry = unsafe { self.map.at_unchecked_mut(self.idx) };
517 &mut entry.1
518 }
519
520 #[inline]
521 pub fn insert(&mut self, value: V) -> V {
522 mem::replace(self.get_mut(), value)
523 }
524
525 #[inline]
526 #[must_use]
527 pub fn into_mut(self) -> &'a mut V {
528 let entry = unsafe { self.map.at_unchecked_mut(self.idx) };
529 &mut entry.1
530 }
531
532 #[inline]
533 #[must_use]
534 pub fn key(&self) -> &K {
535 let entry = unsafe { self.map.at_unchecked(self.idx) };
536 &entry.0
537 }
538
539 #[inline]
540 #[must_use]
541 pub fn remove(self) -> V {
542 self.remove_entry().1
543 }
544
545 #[inline]
546 #[must_use]
547 pub fn remove_entry(self) -> (K, V) {
548 self.map.0.remove(self.idx)
549 }
550}
551
552#[cfg(feature = "serde")]
553mod serde_impl {
554 use super::*;
555
556 use serde::{Deserialize, Serialize};
557
558 impl<'de, K, V> Deserialize<'de> for VecMap<K, V>
559 where
560 K: Ord + Deserialize<'de>,
561 V: Deserialize<'de>,
562 {
563 #[inline]
564 fn deserialize<D>(deserializer: D) -> Result<VecMap<K, V>, D::Error>
565 where
566 D: ::serde::de::Deserializer<'de>,
567 {
568 <Vec<(K, V)>>::deserialize(deserializer).map(VecMap::from_vec)
569 }
570 }
571
572 impl<K, V> Serialize for VecMap<K, V>
573 where
574 K: Serialize,
575 V: Serialize,
576 {
577 #[inline]
578 fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
579 where
580 S: ::serde::ser::Serializer,
581 {
582 <[(K, V)]>::serialize(self.0.as_slice(), serializer)
583 }
584 }
585}
586
587#[cfg(test)]
588mod tests {
589 use super::*;
590
591 use alloc::string::String;
592 use alloc::string::ToString;
593
594 #[test]
595 fn from_vec() {
596 let m: VecMap<u8, u8> =
597 VecMap::from_vec(vec![(4, 1), (2, 3), (5, 7), (2, 9), (4, 6), (7, 8)]);
598 assert!([1, 6].contains(m.get(&4).unwrap()));
599 assert!([3, 9].contains(m.get(&2).unwrap()));
600 assert_eq!(*m.get(&5).unwrap(), 7);
601 assert_eq!(*m.get(&7).unwrap(), 8);
602 }
603
604 #[test]
605 fn merge_max() {
606 let mut m1: VecMap<u8, u8> = VecMap::from_vec(vec![(1, 1), (3, 3), (5, 5)]);
607 let m2: VecMap<u8, u8> = VecMap::from_vec(vec![(1, 1), (2, 2), (3, 2), (4, 4), (5, 6)]);
608 m1.merge_copied_with(&m2, |v1, v2| v1.max(v2));
609 assert_eq!(*m1.get(&1).unwrap(), 1);
610 assert_eq!(*m1.get(&2).unwrap(), 2);
611 assert_eq!(*m1.get(&3).unwrap(), 3);
612 assert_eq!(*m1.get(&4).unwrap(), 4);
613 assert_eq!(*m1.get(&5).unwrap(), 6);
614 }
615
616 #[test]
617 fn remove_less_than() {
618 let mut m: VecMap<u8, String> = VecMap::from_vec(vec![
619 (4, 1.to_string()),
620 (2, 3.to_string()),
621 (5, 7.to_string()),
622 (2, 9.to_string()),
623 (4, 6.to_string()),
624 (7, 8.to_string()),
625 ]);
626 m.remove_less_than(&5);
627 assert!(m.get(&2).is_none());
628 assert!(m.get(&4).is_none());
629 assert!(m.get(&5).is_some());
630 assert!(m.get(&7).is_some());
631 }
632
633 #[test]
634 fn apply() {
635 let map = VecMap::from_iter([(1, 2), (3, 4), (5, 6)]);
636
637 {
638 let keys = VecSet::new();
639 let mut ans = Vec::new();
640 map.apply(&keys, |&v| ans.push(v));
641 assert!(ans.is_empty());
642 }
643 {
644 let keys = VecSet::from_single(3);
645 let mut ans = Vec::new();
646 map.apply(&keys, |&v| ans.push(v));
647 assert_eq!(ans, [4]);
648 }
649 {
650 let keys = VecSet::from_iter([0, 1, 2, 3, 4, 5, 6]);
651 let mut ans = Vec::new();
652 map.apply(&keys, |&v| ans.push(v));
653 assert_eq!(ans, [2, 4, 6]);
654 }
655 }
656}