1use std::borrow::Borrow;
4use std::cmp::Ordering;
5use std::collections::{hash_map, HashMap as Inner};
6use std::fmt;
7use std::hash::Hash;
8use std::sync::Arc;
9
10use get_size::GetSize;
11use get_size_derive::*;
12
13use super::set::OrdHashSet;
14
15pub struct Drain<'a, K, V> {
17 inner: &'a mut Inner<Arc<K>, V>,
18 order: super::set::Drain<'a, Arc<K>>,
19}
20
21impl<'a, K: Eq + Hash + fmt::Debug, V> Drain<'a, K, V> {
22 fn next_entry(&mut self, key: Arc<K>) -> Option<(K, V)> {
23 let value = self.inner.remove(&*key).expect("value");
24 let key = Arc::try_unwrap(key).expect("key");
25 Some((key, value))
26 }
27}
28
29impl<'a, K: Eq + Hash + fmt::Debug, V> Iterator for Drain<'a, K, V> {
30 type Item = (K, V);
31
32 fn next(&mut self) -> Option<Self::Item> {
33 let key = self.order.next()?;
34 self.next_entry(key)
35 }
36
37 fn size_hint(&self) -> (usize, Option<usize>) {
38 self.order.size_hint()
39 }
40}
41
42impl<'a, K: Eq + Hash + fmt::Debug, V> DoubleEndedIterator for Drain<'a, K, V> {
43 fn next_back(&mut self) -> Option<Self::Item> {
44 let key = self.order.next_back()?;
45 self.next_entry(key)
46 }
47}
48
49pub struct DrainWhile<'a, K, V, Cond> {
51 inner: &'a mut Inner<Arc<K>, V>,
52 order: &'a mut OrdHashSet<Arc<K>>,
53 cond: Cond,
54}
55
56impl<'a, K, V, Cond> Iterator for DrainWhile<'a, K, V, Cond>
57where
58 K: Eq + Hash + Ord + fmt::Debug,
59 Cond: Fn((&K, &V)) -> bool,
60{
61 type Item = (K, V);
62
63 fn next(&mut self) -> Option<Self::Item> {
64 if {
65 let key = self.order.first()?;
66 let value = self.inner.get(&**key).expect("value");
67 (self.cond)((&**key, value))
68 } {
69 let key = self.order.pop_first().expect("key");
70 let value = self.inner.remove(&*key).expect("value");
71 let key = Arc::try_unwrap(key).expect("key");
72 Some((key, value))
73 } else {
74 None
75 }
76 }
77
78 fn size_hint(&self) -> (usize, Option<usize>) {
79 (0, Some(self.inner.len()))
80 }
81}
82
83pub struct IntoIter<K, V> {
85 inner: Inner<Arc<K>, V>,
86 order: super::set::IntoIter<Arc<K>>,
87}
88
89impl<K: Eq + Hash + fmt::Debug, V> IntoIter<K, V> {
90 fn next_entry(&mut self, key: Arc<K>) -> Option<(K, V)> {
91 let value = self.inner.remove(&*key).expect("value");
92 let key = Arc::try_unwrap(key).expect("key");
93 Some((key, value))
94 }
95}
96
97impl<K: Eq + Hash + fmt::Debug, V> Iterator for IntoIter<K, V> {
98 type Item = (K, V);
99
100 fn next(&mut self) -> Option<Self::Item> {
101 let key = self.order.next()?;
102 self.next_entry(key)
103 }
104
105 fn size_hint(&self) -> (usize, Option<usize>) {
106 self.order.size_hint()
107 }
108}
109
110impl<K: Eq + Hash + fmt::Debug, V> DoubleEndedIterator for IntoIter<K, V> {
111 fn next_back(&mut self) -> Option<Self::Item> {
112 let key = self.order.next_back()?;
113 self.next_entry(key)
114 }
115}
116
117pub struct Iter<'a, K, V> {
119 inner: &'a Inner<Arc<K>, V>,
120 order: super::set::Iter<'a, Arc<K>>,
121}
122
123impl<'a, K: Eq + Hash + fmt::Debug, V> Iterator for Iter<'a, K, V> {
124 type Item = (&'a K, &'a V);
125
126 fn next(&mut self) -> Option<Self::Item> {
127 let key = &**self.order.next()?;
128 let (key, value) = self.inner.get_key_value(key).expect("entry");
129 Some((&**key, value))
130 }
131
132 fn size_hint(&self) -> (usize, Option<usize>) {
133 self.order.size_hint()
134 }
135}
136
137impl<'a, K: Eq + Hash + fmt::Debug, V> DoubleEndedIterator for Iter<'a, K, V> {
138 fn next_back(&mut self) -> Option<Self::Item> {
139 let key = self.order.next_back()?;
140 let (key, value) = self.inner.get_key_value(&**key).expect("entry");
141 Some((&**key, value))
142 }
143}
144
145pub struct Keys<'a, K, V> {
147 inner: &'a Inner<Arc<K>, V>,
148 order: super::set::Iter<'a, Arc<K>>,
149}
150
151impl<'a, K: Eq + Hash + fmt::Debug, V> Iterator for Keys<'a, K, V> {
152 type Item = &'a K;
153
154 fn next(&mut self) -> Option<Self::Item> {
155 let key = self.order.next()?;
156 let (key, _) = self.inner.get_key_value(&**key).expect("entry");
157 Some(&**key)
158 }
159
160 fn size_hint(&self) -> (usize, Option<usize>) {
161 self.order.size_hint()
162 }
163}
164
165impl<'a, K: Eq + Hash + fmt::Debug, V> DoubleEndedIterator for Keys<'a, K, V> {
166 fn next_back(&mut self) -> Option<Self::Item> {
167 let key = self.order.next_back()?;
168 let (key, _) = self.inner.get_key_value(&**key).expect("entry");
169 Some(&**key)
170 }
171}
172
173pub struct IntoValues<K, V> {
175 inner: Inner<Arc<K>, V>,
176 order: super::set::IntoIter<Arc<K>>,
177}
178
179impl<K: Eq + Hash + fmt::Debug, V> Iterator for IntoValues<K, V> {
180 type Item = V;
181
182 fn next(&mut self) -> Option<Self::Item> {
183 let key = self.order.next()?;
184 self.inner.remove(&*key)
185 }
186
187 fn size_hint(&self) -> (usize, Option<usize>) {
188 self.order.size_hint()
189 }
190}
191
192impl<K: Eq + Hash + fmt::Debug, V> DoubleEndedIterator for IntoValues<K, V> {
193 fn next_back(&mut self) -> Option<Self::Item> {
194 let key = self.order.next_back()?;
195 self.inner.remove(&*key)
196 }
197}
198
199pub struct Values<'a, K, V> {
201 inner: &'a Inner<Arc<K>, V>,
202 order: super::set::Iter<'a, Arc<K>>,
203}
204
205impl<'a, K: Eq + Hash + fmt::Debug, V> Iterator for Values<'a, K, V> {
206 type Item = &'a V;
207
208 fn next(&mut self) -> Option<Self::Item> {
209 let key = self.order.next()?;
210 self.inner.get(&**key)
211 }
212
213 fn size_hint(&self) -> (usize, Option<usize>) {
214 self.order.size_hint()
215 }
216}
217
218impl<'a, K: Eq + Hash + fmt::Debug, V> DoubleEndedIterator for Values<'a, K, V> {
219 fn next_back(&mut self) -> Option<Self::Item> {
220 let key = self.order.next_back()?;
221 self.inner.get(&**key)
222 }
223}
224
225pub type OccupiedEntry<'a, K, V> = hash_map::OccupiedEntry<'a, Arc<K>, V>;
227
228pub struct VacantEntry<'a, K, V> {
230 order: &'a mut OrdHashSet<Arc<K>>,
231 entry: hash_map::VacantEntry<'a, Arc<K>, V>,
232}
233
234impl<'a, K, V> VacantEntry<'a, K, V>
235where
236 K: Eq + Hash + Ord,
237{
238 pub fn insert(self, value: V) -> &'a mut V {
240 self.order.insert(self.entry.key().clone());
241 self.entry.insert(value)
242 }
243}
244
245pub enum Entry<'a, K, V> {
247 Occupied(hash_map::OccupiedEntry<'a, Arc<K>, V>),
248 Vacant(VacantEntry<'a, K, V>),
249}
250
251#[derive(GetSize)]
253pub struct OrdHashMap<K, V> {
254 inner: Inner<Arc<K>, V>,
255 order: OrdHashSet<Arc<K>>,
256}
257
258impl<K: Clone + Eq + Hash + Ord + fmt::Debug, V: Clone> Clone for OrdHashMap<K, V> {
259 fn clone(&self) -> Self {
260 let mut inner = Inner::with_capacity(self.inner.capacity());
261 let mut order = OrdHashSet::<Arc<K>>::with_capacity(inner.capacity());
262
263 for key in &self.order {
264 let key = Arc::new(K::clone(&**key));
265 let value = self.inner.get(&key).cloned().expect("value");
266 inner.insert(key.clone(), value);
267 order.insert(key);
268 }
269
270 Self { inner, order }
271 }
272}
273impl<K, V> OrdHashMap<K, V> {
274 pub fn new() -> Self {
276 Self {
277 inner: Inner::new(),
278 order: OrdHashSet::new(),
279 }
280 }
281 pub fn with_capacity(capacity: usize) -> Self {
283 Self {
284 inner: Inner::with_capacity(capacity),
285 order: OrdHashSet::with_capacity(capacity),
286 }
287 }
288
289 pub fn len(&self) -> usize {
291 self.inner.len()
292 }
293}
294
295impl<K: Eq + Hash + Ord, V> OrdHashMap<K, V> {
296 pub fn bisect<Cmp>(&self, cmp: Cmp) -> Option<&V>
302 where
303 Cmp: Fn(&K) -> Option<Ordering> + Copy,
304 {
305 self.order
306 .bisect(|key| cmp(&*key))
307 .map(|key| self.get(&**key).expect("value"))
308 }
309
310 pub fn clear(&mut self) {
312 self.inner.clear();
313 self.order.clear();
314 }
315
316 pub fn contains_key<Q>(&self, key: &Q) -> bool
318 where
319 Arc<K>: Borrow<Q>,
320 Q: Eq + Hash + ?Sized,
321 {
322 self.inner.contains_key(key)
323 }
324
325 pub fn drain(&mut self) -> Drain<K, V> {
327 Drain {
328 inner: &mut self.inner,
329 order: self.order.drain(),
330 }
331 }
332
333 pub fn drain_while<Cond>(&mut self, cond: Cond) -> DrainWhile<K, V, Cond>
335 where
336 Cond: Fn((&K, &V)) -> bool,
337 {
338 DrainWhile {
339 inner: &mut self.inner,
340 order: &mut self.order,
341 cond,
342 }
343 }
344
345 pub fn entry<Q: Into<Arc<K>>>(&mut self, key: Q) -> Entry<K, V> {
347 match self.inner.entry(key.into()) {
348 hash_map::Entry::Occupied(entry) => Entry::Occupied(entry),
349 hash_map::Entry::Vacant(entry) => Entry::Vacant(VacantEntry {
350 entry,
351 order: &mut self.order,
352 }),
353 }
354 }
355
356 pub fn extend<I: IntoIterator<Item = (K, V)>>(&mut self, iter: I) {
358 for (key, value) in iter {
359 self.insert(key, value);
360 }
361 }
362
363 pub fn first(&self) -> Option<&V> {
365 let key = self.order.first()?;
366 self.inner.get(&**key)
367 }
368
369 pub fn first_mut(&mut self) -> Option<&mut V> {
371 let key = self.order.first()?;
372 self.inner.get_mut(&**key)
373 }
374
375 pub fn get<Q>(&self, key: &Q) -> Option<&V>
377 where
378 Arc<K>: Borrow<Q>,
379 Q: Eq + Hash + ?Sized,
380 {
381 self.inner.get(key)
382 }
383
384 pub fn get_key_value<Q>(&self, key: &Q) -> Option<(&K, &V)>
386 where
387 Arc<K>: Borrow<Q>,
388 Q: Eq + Hash + ?Sized,
389 {
390 self.inner
391 .get_key_value(key)
392 .map(|(key, value)| (&**key, value))
393 }
394
395 pub fn get_mut<Q>(&mut self, key: &Q) -> Option<&mut V>
397 where
398 Arc<K>: Borrow<Q>,
399 Q: Eq + Hash + ?Sized,
400 {
401 self.inner.get_mut(key)
402 }
403
404 pub fn insert(&mut self, key: K, value: V) -> Option<V> {
406 let key = Arc::new(key);
407 self.order.insert(key.clone());
408 self.inner.insert(key, value)
409 }
410
411 pub fn iter(&self) -> Iter<K, V> {
413 Iter {
414 inner: &self.inner,
415 order: self.order.iter(),
416 }
417 }
418
419 pub fn is_empty(&self) -> bool {
421 self.inner.is_empty()
422 }
423
424 pub fn keys(&self) -> Keys<K, V> {
426 Keys {
427 inner: &self.inner,
428 order: self.order.iter(),
429 }
430 }
431
432 pub fn last(&self) -> Option<&V> {
434 let key = self.order.last()?;
435 self.inner.get(&**key)
436 }
437
438 pub fn last_mut(&mut self) -> Option<&mut V> {
440 let key = self.order.last()?;
441 self.inner.get_mut(&**key)
442 }
443
444 pub fn pop_first(&mut self) -> Option<V>
446 where
447 K: fmt::Debug,
448 {
449 let key = self.order.pop_first()?;
450 self.inner.remove(&*key)
451 }
452
453 pub fn pop_last(&mut self) -> Option<V>
455 where
456 K: fmt::Debug,
457 {
458 let key = self.order.pop_last()?;
459 self.inner.remove(&*key)
460 }
461
462 pub fn remove<Q>(&mut self, key: &Q) -> Option<V>
464 where
465 Arc<K>: Borrow<Q>,
466 Q: Hash + Eq + ?Sized,
467 {
468 let (key, value) = self.inner.remove_entry(key)?;
469 assert!(self.order.remove(&key));
470 Some(value)
471 }
472
473 pub fn starts_with<'a, I: IntoIterator<Item = (&'a K, &'a V)>>(&self, other: I) -> bool
475 where
476 K: fmt::Debug + 'a,
477 V: PartialEq + 'a,
478 {
479 let mut this = self.iter();
480 let mut that = other.into_iter();
481
482 while let Some(item) = that.next() {
483 if this.next() != Some(item) {
484 return false;
485 }
486 }
487
488 true
489 }
490
491 pub fn into_values(self) -> IntoValues<K, V>
493 where
494 K: fmt::Debug,
495 {
496 IntoValues {
497 inner: self.inner,
498 order: self.order.into_iter(),
499 }
500 }
501
502 pub fn values(&self) -> Values<K, V> {
504 Values {
505 inner: &self.inner,
506 order: self.order.iter(),
507 }
508 }
509}
510
511impl<K: Eq + Hash + Ord + fmt::Debug, V> OrdHashMap<K, V> {
512 pub fn bisect_and_remove<Cmp>(&mut self, cmp: Cmp) -> Option<(K, V)>
518 where
519 Cmp: Fn(&K) -> Option<Ordering> + Copy,
520 {
521 let key = self.order.bisect_and_remove(|key| cmp(&*key))?;
522 let value = self.inner.remove(&*key).expect("value");
523 let key = Arc::try_unwrap(key).expect("key");
524 Some((key, value))
525 }
526
527 pub fn pop_first_entry(&mut self) -> Option<(K, V)> {
529 let key = self.order.pop_first()?;
530 let (key, value) = self.inner.remove_entry(&*key).expect("entry");
531 let key = Arc::try_unwrap(key).expect("key");
532 Some((key, value))
533 }
534
535 pub fn pop_last_entry(&mut self) -> Option<(K, V)> {
537 let key = self.order.pop_last()?;
538 let (key, value) = self.inner.remove_entry(&*key).expect("entry");
539 let key = Arc::try_unwrap(key).expect("key");
540 Some((key, value))
541 }
542
543 pub fn remove_entry<Q>(&mut self, key: &Q) -> Option<(K, V)>
545 where
546 Arc<K>: Borrow<Q>,
547 Q: Hash + Eq + ?Sized,
548 {
549 let (key, value) = self.inner.remove_entry(key)?;
550 assert!(self.order.remove(&key));
551 let key = Arc::try_unwrap(key).expect("key");
552 Some((key, value))
553 }
554}
555
556impl<K: Eq + Hash + Ord + fmt::Debug, V: PartialEq> PartialEq<Self> for OrdHashMap<K, V> {
557 fn eq(&self, other: &Self) -> bool {
558 if self.len() != other.len() {
559 return false;
560 }
561
562 self.iter()
563 .zip(other)
564 .all(|((lk, lv), (rk, rv))| lk == rk && lv == rv)
565 }
566}
567
568impl<K: Eq + Hash + Ord + fmt::Debug, V: Eq> Eq for OrdHashMap<K, V> {}
569
570impl<K: Eq + Hash + Ord + fmt::Debug, V> FromIterator<(K, V)> for OrdHashMap<K, V> {
571 fn from_iter<T: IntoIterator<Item = (K, V)>>(iter: T) -> Self {
572 let iter = iter.into_iter();
573 let mut map = match iter.size_hint() {
574 (_, Some(max)) => Self::with_capacity(max),
575 (min, None) if min > 0 => Self::with_capacity(min),
576 _ => Self::new(),
577 };
578
579 map.extend(iter);
580 map
581 }
582}
583
584impl<K: Eq + Hash + fmt::Debug, V> IntoIterator for OrdHashMap<K, V> {
585 type Item = (K, V);
586 type IntoIter = IntoIter<K, V>;
587
588 fn into_iter(self) -> Self::IntoIter {
589 IntoIter {
590 inner: self.inner,
591 order: self.order.into_iter(),
592 }
593 }
594}
595
596impl<'a, K: Ord + Eq + Hash + fmt::Debug, V> IntoIterator for &'a OrdHashMap<K, V> {
597 type Item = (&'a K, &'a V);
598 type IntoIter = Iter<'a, K, V>;
599
600 fn into_iter(self) -> Self::IntoIter {
601 self.iter()
602 }
603}
604
605impl<K: Eq + Hash + Ord + fmt::Debug, V: fmt::Debug> fmt::Debug for OrdHashMap<K, V> {
606 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
607 f.write_str("{")?;
608
609 for (key, value) in self.iter() {
610 write!(f, "{:?}: {:?}, ", key, value)?;
611 }
612
613 f.write_str("}")
614 }
615}
616
617#[cfg(test)]
618mod tests {
619 use super::*;
620
621 #[test]
622 fn test_find() {
623 let mut map = OrdHashMap::<&str, u32>::new();
624 map.insert(".tmp", 0);
625 map.insert("1", 1);
626 map.insert("2", 2);
627 map.insert("3", 3);
628 map.insert("five", 5);
629 map.insert("six", 6);
630
631 assert_eq!(
632 map.bisect(|key| key.parse().ok().map(|key: u32| 1.cmp(&key))),
633 Some(&1)
634 );
635
636 assert_eq!(
637 map.bisect(|key| key.parse().ok().map(|key: u32| 2.cmp(&key))),
638 Some(&2)
639 );
640
641 assert_eq!(
642 map.bisect(|key| key.parse().ok().map(|key: u32| 3.cmp(&key))),
643 Some(&3)
644 );
645 }
646
647 #[test]
648 fn test_drain() {
649 let mut map: OrdHashMap<u32, String> =
650 (0..10).into_iter().map(|i| (i, i.to_string())).collect();
651
652 assert_eq!(
653 map.drain().collect::<Vec<_>>(),
654 (0..10)
655 .into_iter()
656 .map(|i| (i, i.to_string()))
657 .collect::<Vec<_>>()
658 );
659 }
660
661 #[test]
662 fn test_drain_partial() {
663 let mut map: OrdHashMap<u32, String> =
664 (0..10).into_iter().map(|i| (i, i.to_string())).collect();
665
666 assert_eq!(
667 map.drain().take_while(|(k, _v)| *k < 5).collect::<Vec<_>>(),
668 (0..5)
669 .into_iter()
670 .map(|i| (i, i.to_string()))
671 .collect::<Vec<_>>()
672 );
673
674 assert_eq!(
675 map,
676 (6..10).into_iter().map(|i| (i, i.to_string())).collect()
677 );
678 }
679
680 #[test]
681 fn test_drain_while() {
682 let mut map: OrdHashMap<u32, String> =
683 (0..10).into_iter().map(|i| (i, i.to_string())).collect();
684
685 assert_eq!(
686 map.drain_while(|(k, _v)| *k < 5).collect::<Vec<_>>(),
687 (0..5)
688 .into_iter()
689 .map(|i| (i, i.to_string()))
690 .collect::<Vec<_>>()
691 );
692
693 assert_eq!(
694 map,
695 (5..10).into_iter().map(|i| (i, i.to_string())).collect()
696 );
697 }
698}