1#![deny(missing_docs)]
2#![doc = include_str!("../README.md")]
3
4use std::{
5 cmp::Reverse, default::Default, fmt, iter::FromIterator, iter::FusedIterator, num::Wrapping,
6};
7
8type Bucket<K, V> = Vec<(K, V)>;
9
10#[derive(Clone)]
21pub struct RadixHeapMap<K, V> {
22 len: usize,
23
24 top: Option<K>,
26
27 buckets: Vec<Bucket<K, V>>,
32
33 initial: Bucket<K, V>,
35}
36
37impl<K: Radix + Ord + Copy, V> RadixHeapMap<K, V> {
38 pub fn new() -> RadixHeapMap<K, V> {
40 RadixHeapMap {
41 len: 0,
42 top: None,
43 buckets: (0..=K::RADIX_BITS).map(|_| Bucket::default()).collect(),
44 initial: Bucket::default(),
45 }
46 }
47
48 pub fn new_at(top: K) -> RadixHeapMap<K, V> {
54 RadixHeapMap {
55 len: 0,
56 top: Some(top),
57 buckets: (0..=K::RADIX_BITS).map(|_| Bucket::default()).collect(),
58 initial: Bucket::default(),
59 }
60 }
61
62 pub fn clear(&mut self) {
64 self.len = 0;
65 self.top = None;
66 self.initial.clear();
67
68 for bucket in &mut self.buckets {
69 bucket.clear();
70 }
71 }
72
73 pub fn clear_to(&mut self, top: K) {
79 self.clear();
80 self.top = Some(top);
81 }
82
83 pub fn constrain(&mut self) {
85 let repush = if self.top.is_some() {
86 let index = self
87 .buckets
88 .iter()
89 .enumerate()
90 .find(|&(_, bucket)| !bucket.is_empty())
91 .map(|(i, _)| i);
92
93 match index {
94 None | Some(0) => None,
95 Some(index) => {
96 let (buckets, rest) = self.buckets.split_at_mut(index);
97 Some((buckets, &mut rest[0]))
98 }
99 }
100 } else if !self.initial.is_empty() {
101 Some((&mut self.buckets[..], &mut self.initial))
102 } else {
103 None
104 };
105
106 if let Some((buckets, repush)) = repush {
107 let top = *repush
108 .iter()
109 .map(|(k, _)| k)
110 .max()
111 .expect("Expected non-empty bucket");
112
113 self.top = Some(top);
114
115 repush.drain(..).for_each(|(key, value)| {
116 buckets[key.radix_distance(&top) as usize].push((key, value))
117 });
118 }
119 }
120
121 #[inline]
127 pub fn push(&mut self, key: K, value: V) {
128 let bucket = if let Some(top) = self.top {
129 assert!(key <= top, "Key must be lower or equal to current top key");
130 &mut self.buckets[key.radix_distance(&top) as usize]
131 } else {
132 &mut self.initial
133 };
134
135 bucket.push((key, value));
136 self.len += 1;
137 }
138
139 #[inline]
147 pub fn pop(&mut self) -> Option<(K, V)> {
148 let mut constrained = false;
149
150 loop {
151 let pop = self.buckets[0].pop();
152
153 if pop.is_some() {
154 self.len -= 1;
155 return pop;
156 } else if constrained {
157 return pop;
158 } else {
159 constrained = true;
160 self.constrain()
161 }
162 }
163 }
164
165 #[inline]
167 pub fn len(&self) -> usize {
168 self.len
169 }
170
171 #[inline]
173 pub fn is_empty(&self) -> bool {
174 self.len() == 0
175 }
176
177 #[inline]
179 pub fn top(&self) -> Option<K> {
180 self.top
181 }
182
183 pub fn shrink_to_fit(&mut self) {
185 self.initial.shrink_to_fit();
186
187 for bucket in &mut self.buckets {
188 bucket.shrink_to_fit();
189 }
190 }
191
192 pub fn iter(&self) -> Iter<K, V> {
194 Iter {
195 cur_bucket: self.initial.iter(),
196 buckets: self.buckets.iter(),
197 size: self.len,
198 }
199 }
200
201 pub fn keys(&self) -> Keys<K, V> {
203 Keys(self.iter())
204 }
205
206 pub fn values(&self) -> Values<K, V> {
208 Values(self.iter())
209 }
210}
211
212impl<K: Radix + Ord + Copy, V> Default for RadixHeapMap<K, V> {
213 fn default() -> RadixHeapMap<K, V> {
214 RadixHeapMap::new()
215 }
216}
217
218impl<K: Radix + Ord + Copy, V> FromIterator<(K, V)> for RadixHeapMap<K, V> {
219 fn from_iter<I>(iter: I) -> RadixHeapMap<K, V>
220 where
221 I: IntoIterator<Item = (K, V)>,
222 {
223 let mut heap = RadixHeapMap::new();
224
225 for (k, v) in iter {
226 heap.push(k, v);
227 }
228
229 heap
230 }
231}
232
233impl<K: Radix + Ord + Copy, V> Extend<(K, V)> for RadixHeapMap<K, V> {
234 fn extend<I>(&mut self, iter: I)
235 where
236 I: IntoIterator<Item = (K, V)>,
237 {
238 for (k, v) in iter {
239 self.push(k, v);
240 }
241 }
242}
243
244impl<'a, K: Radix + Ord + Copy + 'a, V: Copy + 'a> Extend<&'a (K, V)> for RadixHeapMap<K, V> {
245 fn extend<I>(&mut self, iter: I)
246 where
247 I: IntoIterator<Item = &'a (K, V)>,
248 {
249 for &(k, v) in iter {
250 self.push(k, v);
251 }
252 }
253}
254
255impl<K: Radix + Ord + Copy + fmt::Debug, V: fmt::Debug> fmt::Debug for RadixHeapMap<K, V> {
256 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
257 f.debug_list().entries(self.iter()).finish()
258 }
259}
260
261#[derive(Clone)]
263pub struct IntoIter<K, V> {
264 cur_bucket: std::vec::IntoIter<(K, V)>,
265 buckets: std::vec::IntoIter<Bucket<K, V>>,
266 size: usize,
267}
268
269impl<K, V> Iterator for IntoIter<K, V> {
270 type Item = (K, V);
271
272 fn next(&mut self) -> Option<Self::Item> {
273 loop {
274 if let pair @ Some(_) = self.cur_bucket.next() {
275 self.size -= 1;
276 return pair;
277 } else {
278 self.cur_bucket = self.buckets.next()?.into_iter();
279 }
280 }
281 }
282
283 fn size_hint(&self) -> (usize, Option<usize>) {
284 (self.size, Some(self.size))
285 }
286
287 fn for_each<F>(self, mut f: F)
288 where
289 F: FnMut(Self::Item),
290 {
291 self.cur_bucket.for_each(&mut f);
292 self.buckets.for_each(|b| b.into_iter().for_each(&mut f));
293 }
294}
295
296impl<K, V> ExactSizeIterator for IntoIter<K, V> {}
297
298impl<K, V> FusedIterator for IntoIter<K, V> {}
299
300#[derive(Clone)]
302pub struct Iter<'a, K, V> {
303 cur_bucket: std::slice::Iter<'a, (K, V)>,
304 buckets: std::slice::Iter<'a, Bucket<K, V>>,
305 size: usize,
306}
307
308impl<'a, K, V> Iterator for Iter<'a, K, V> {
309 type Item = &'a (K, V);
310
311 fn next(&mut self) -> Option<Self::Item> {
312 loop {
313 if let pair @ Some(_) = self.cur_bucket.next() {
314 self.size -= 1;
315 return pair;
316 } else {
317 self.cur_bucket = self.buckets.next()?.iter();
318 }
319 }
320 }
321
322 fn size_hint(&self) -> (usize, Option<usize>) {
323 (self.size, Some(self.size))
324 }
325
326 fn for_each<F>(self, mut f: F)
327 where
328 F: FnMut(Self::Item),
329 {
330 self.cur_bucket.for_each(&mut f);
331 self.buckets.for_each(|b| b.iter().for_each(&mut f));
332 }
333}
334
335impl<'a, K, V> ExactSizeIterator for Iter<'a, K, V> {}
336
337impl<'a, K, V> FusedIterator for Iter<'a, K, V> {}
338
339#[derive(Clone)]
341pub struct Keys<'a, K, V>(Iter<'a, K, V>);
342
343impl<'a, K, V> Iterator for Keys<'a, K, V> {
344 type Item = &'a K;
345
346 fn next(&mut self) -> Option<Self::Item> {
347 self.0.next().map(|(k, _)| k)
348 }
349
350 fn size_hint(&self) -> (usize, Option<usize>) {
351 self.0.size_hint()
352 }
353
354 fn for_each<F>(self, mut f: F)
355 where
356 F: FnMut(Self::Item),
357 {
358 self.0.for_each(|(k, _)| f(k))
359 }
360}
361
362impl<'a, K, V> ExactSizeIterator for Keys<'a, K, V> {}
363
364impl<'a, K, V> FusedIterator for Keys<'a, K, V> {}
365
366#[derive(Clone)]
368pub struct Values<'a, K, V>(Iter<'a, K, V>);
369
370impl<'a, K, V> Iterator for Values<'a, K, V> {
371 type Item = &'a V;
372
373 fn next(&mut self) -> Option<Self::Item> {
374 self.0.next().map(|(_, v)| v)
375 }
376
377 fn size_hint(&self) -> (usize, Option<usize>) {
378 self.0.size_hint()
379 }
380
381 fn for_each<F>(self, mut f: F)
382 where
383 F: FnMut(Self::Item),
384 {
385 self.0.for_each(|(_, v)| f(v))
386 }
387}
388
389impl<'a, K, V> ExactSizeIterator for Values<'a, K, V> {}
390
391impl<'a, K, V> FusedIterator for Values<'a, K, V> {}
392
393impl<K: Radix + Ord + Copy, V> IntoIterator for RadixHeapMap<K, V> {
394 type Item = (K, V);
395 type IntoIter = IntoIter<K, V>;
396
397 fn into_iter(self) -> Self::IntoIter {
398 IntoIter {
399 cur_bucket: self.initial.into_iter(),
400 buckets: self.buckets.into_iter(),
401 size: self.len,
402 }
403 }
404}
405
406impl<'a, K: Radix + Ord + Copy, V> IntoIterator for &'a RadixHeapMap<K, V> {
407 type Item = &'a (K, V);
408 type IntoIter = Iter<'a, K, V>;
409
410 fn into_iter(self) -> Self::IntoIter {
411 self.iter()
412 }
413}
414
415pub trait Radix {
417 fn radix_similarity(&self, other: &Self) -> u32;
422
423 fn radix_distance(&self, other: &Self) -> u32 {
426 Self::RADIX_BITS - self.radix_similarity(other)
427 }
428
429 const RADIX_BITS: u32;
431}
432
433macro_rules! radix_wrapper_impl {
434 ($t:ident) => {
435 impl<T: Radix> Radix for $t<T> {
436 #[inline]
437 fn radix_similarity(&self, other: &$t<T>) -> u32 {
438 self.0.radix_similarity(&other.0)
439 }
440
441 const RADIX_BITS: u32 = T::RADIX_BITS;
442 }
443 };
444}
445
446radix_wrapper_impl!(Reverse);
447radix_wrapper_impl!(Wrapping);
448
449macro_rules! radix_int_impl {
450 ($t:ty) => {
451 impl Radix for $t {
452 #[inline]
453 fn radix_similarity(&self, other: &$t) -> u32 {
454 (self ^ other).leading_zeros()
455 }
456
457 const RADIX_BITS: u32 = (std::mem::size_of::<$t>() * 8) as u32;
458 }
459 };
460}
461
462radix_int_impl!(i8);
463radix_int_impl!(i16);
464radix_int_impl!(i32);
465radix_int_impl!(i64);
466radix_int_impl!(i128);
467radix_int_impl!(isize);
468
469radix_int_impl!(u8);
470radix_int_impl!(u16);
471radix_int_impl!(u32);
472radix_int_impl!(u64);
473radix_int_impl!(u128);
474radix_int_impl!(usize);
475
476#[cfg(feature = "ordered-float")]
477macro_rules! radix_float_impl {
478 ($t:ty, $bits:ty, $wrapper:path) => {
479 impl Radix for $wrapper {
480 #[inline]
481 fn radix_similarity(&self, other: &$wrapper) -> u32 {
482 let self_bits: $bits = self.to_bits();
483 let other_bits: $bits = other.to_bits();
484 self_bits.radix_similarity(&other_bits)
485 }
486
487 const RADIX_BITS: u32 = <$bits>::RADIX_BITS;
488 }
489 };
490}
491
492#[cfg(feature = "ordered-float")]
493radix_float_impl!(f32, u32, ordered_float::NotNan<f32>);
494
495#[cfg(feature = "ordered-float")]
496radix_float_impl!(f64, u64, ordered_float::NotNan<f64>);
497
498impl Radix for () {
499 #[inline]
500 fn radix_similarity(&self, _: &()) -> u32 {
501 0
502 }
503 const RADIX_BITS: u32 = 0;
504}
505
506macro_rules! radix_tuple_impl {
507 ($(
508 $Tuple:ident {
509 $(($idx:tt) -> $T:ident)+
510 }
511 )+) => {
512 $(
513 impl<$($T:Radix),+> Radix for ($($T,)+) {
514 #[inline]
515 fn radix_similarity(&self, other: &($($T,)+)) -> u32 {
516 let similarity = 0;
517
518 $(
519 let s = self.$idx.radix_similarity(&other.$idx);
520 let similarity = similarity + s;
521 if s < <$T as Radix>::RADIX_BITS { return similarity }
522 )+
523
524 return similarity;
525 }
526 const RADIX_BITS: u32 = 0 $(+<$T as Radix>::RADIX_BITS)+;
527 }
528 )+
529 }
530}
531
532radix_tuple_impl! {
533 Tuple1 {
534 (0) -> A
535 }
536 Tuple2 {
537 (0) -> A
538 (1) -> B
539 }
540 Tuple3 {
541 (0) -> A
542 (1) -> B
543 (2) -> C
544 }
545 Tuple4 {
546 (0) -> A
547 (1) -> B
548 (2) -> C
549 (3) -> D
550 }
551 Tuple5 {
552 (0) -> A
553 (1) -> B
554 (2) -> C
555 (3) -> D
556 (4) -> E
557 }
558 Tuple6 {
559 (0) -> A
560 (1) -> B
561 (2) -> C
562 (3) -> D
563 (4) -> E
564 (5) -> F
565 }
566 Tuple7 {
567 (0) -> A
568 (1) -> B
569 (2) -> C
570 (3) -> D
571 (4) -> E
572 (5) -> F
573 (6) -> G
574 }
575 Tuple8 {
576 (0) -> A
577 (1) -> B
578 (2) -> C
579 (3) -> D
580 (4) -> E
581 (5) -> F
582 (6) -> G
583 (7) -> H
584 }
585 Tuple9 {
586 (0) -> A
587 (1) -> B
588 (2) -> C
589 (3) -> D
590 (4) -> E
591 (5) -> F
592 (6) -> G
593 (7) -> H
594 (8) -> I
595 }
596 Tuple10 {
597 (0) -> A
598 (1) -> B
599 (2) -> C
600 (3) -> D
601 (4) -> E
602 (5) -> F
603 (6) -> G
604 (7) -> H
605 (8) -> I
606 (9) -> J
607 }
608 Tuple11 {
609 (0) -> A
610 (1) -> B
611 (2) -> C
612 (3) -> D
613 (4) -> E
614 (5) -> F
615 (6) -> G
616 (7) -> H
617 (8) -> I
618 (9) -> J
619 (10) -> K
620 }
621 Tuple12 {
622 (0) -> A
623 (1) -> B
624 (2) -> C
625 (3) -> D
626 (4) -> E
627 (5) -> F
628 (6) -> G
629 (7) -> H
630 (8) -> I
631 (9) -> J
632 (10) -> K
633 (11) -> L
634 }
635}
636
637#[cfg(test)]
638mod tests {
639 extern crate quickcheck;
640
641 use self::quickcheck::{quickcheck, TestResult};
642 use super::Radix;
643 use super::RadixHeapMap;
644 use std::cmp::Reverse;
645
646 #[test]
647 fn radix_dist() {
648 assert!(4u32.radix_distance(&2) == 3);
649 assert!(3u32.radix_distance(&2) == 1);
650 assert!(2u32.radix_distance(&2) == 0);
651 assert!(1u32.radix_distance(&2) == 2);
652 assert!(0u32.radix_distance(&2) == 2);
653 }
654
655 #[test]
656 fn clear() {
657 let mut heap = RadixHeapMap::new();
658 heap.push(0u32, 'a');
659 heap.clear();
660 assert!(heap.pop().is_none());
661 }
662
663 #[test]
664 fn push_pop() {
665 let mut heap = RadixHeapMap::new();
666 heap.push(0u32, 'a');
667 heap.push(3, 'b');
668 heap.push(2, 'c');
669
670 assert!(heap.len() == 3);
671 assert!(!heap.is_empty());
672
673 assert!(heap.pop() == Some((3, 'b')));
674 assert!(heap.pop() == Some((2, 'c')));
675 assert!(heap.pop() == Some((0, 'a')));
676 assert!(heap.pop() == None);
677
678 assert!(heap.len() == 0);
679 assert!(heap.is_empty());
680 }
681
682 #[test]
683 fn rev_push_pop() {
684 let mut heap = RadixHeapMap::new();
685 heap.push(Reverse(0), 'a');
686 heap.push(Reverse(3), 'b');
687 heap.push(Reverse(2), 'c');
688
689 assert!(heap.len() == 3);
690 assert!(!heap.is_empty());
691
692 assert!(heap.pop() == Some((Reverse(0), 'a')));
693 assert!(heap.pop() == Some((Reverse(2), 'c')));
694 assert!(heap.pop() == Some((Reverse(3), 'b')));
695 assert!(heap.pop() == None);
696
697 assert!(heap.len() == 0);
698 assert!(heap.is_empty());
699 }
700
701 #[test]
702 #[should_panic]
703 fn push_pop_panic() {
704 let mut heap = RadixHeapMap::new();
705 heap.push(0u32, 'a');
706 heap.push(3, 'b');
707
708 assert!(heap.pop() == Some((3, 'b')));
709 heap.push(4, 'd');
710 }
711
712 #[test]
713 fn sort() {
714 fn prop<T: Ord + Radix + Copy>(mut xs: Vec<T>) -> bool {
715 let mut heap: RadixHeapMap<_, _> =
716 xs.iter().enumerate().map(|(i, &d)| (d, i)).collect();
717
718 xs.sort();
719
720 while xs.pop() == heap.pop().map(|(k, _)| k) {
721 if xs.is_empty() {
722 return true;
723 }
724 }
725
726 return false;
727 }
728
729 quickcheck(prop as fn(Vec<()>) -> bool);
730 quickcheck(prop as fn(Vec<u32>) -> bool);
731 quickcheck(prop as fn(Vec<i32>) -> bool);
732 quickcheck(prop as fn(Vec<(u32, i32)>) -> bool);
733 quickcheck(prop as fn(Vec<u8>) -> bool);
734 quickcheck(prop as fn(Vec<i16>) -> bool);
735 quickcheck(prop as fn(Vec<(i64, usize)>) -> bool);
736 quickcheck(prop as fn(Vec<i128>) -> bool);
737 quickcheck(prop as fn(Vec<u128>) -> bool);
738 }
739
740 #[cfg(feature = "ordered-float")]
741 #[test]
742 fn sort_float() {
743 fn prop(xs: Vec<f32>) -> TestResult {
744 if xs.iter().any(|x| x.is_nan()) {
745 return TestResult::discard();
746 }
747
748 let mut xs: Vec<_> = xs
749 .into_iter()
750 .map(|x| ordered_float::NotNan::new(x).unwrap())
751 .collect();
752 xs.sort();
753
754 let mut heap: RadixHeapMap<_, _> =
755 xs.iter().enumerate().map(|(i, &d)| (d, i)).collect();
756
757 while xs.pop() == heap.pop().map(|(k, _)| k) {
758 if xs.is_empty() {
759 return TestResult::passed();
760 }
761 }
762
763 return TestResult::failed();
764 }
765
766 quickcheck(prop as fn(Vec<f32>) -> TestResult);
767 }
768
769 #[test]
770 fn iter_yeilds_all_elements() {
771 fn prop<T: Ord + Radix + Copy>(mut xs: Vec<T>) -> TestResult {
772 let heap = xs.iter().map(|&d| (d, ())).collect::<RadixHeapMap<_, _>>();
773
774 for (k, ()) in heap.iter() {
776 for i in 0..xs.len() {
777 if xs[i] == *k {
778 xs.remove(i);
779 break;
780 }
781 }
782 }
783
784 if xs.is_empty() {
785 TestResult::passed()
786 } else {
787 TestResult::failed()
788 }
789 }
790
791 quickcheck(prop as fn(Vec<u32>) -> TestResult);
792 quickcheck(prop as fn(Vec<i32>) -> TestResult);
793 quickcheck(prop as fn(Vec<(u32, i32)>) -> TestResult);
794 quickcheck(prop as fn(Vec<u8>) -> TestResult);
795 quickcheck(prop as fn(Vec<i16>) -> TestResult);
796 quickcheck(prop as fn(Vec<(i64, usize)>) -> TestResult);
797 }
798
799 #[test]
800 fn into_iter_inital() {
801 let mut heap = RadixHeapMap::new();
802 heap.push(1, 2);
803 heap.push(5, 2);
804
805 let mut vec: Vec<_> = heap.into_iter().collect();
806 vec.sort();
807 assert_eq!(vec, vec![(1, 2), (5, 2)]);
808 }
809
810 #[test]
811 fn into_iter() {
812 let mut heap = RadixHeapMap::new();
813 heap.push(1, 2);
814 heap.push(5, 4);
815 heap.push(7, 1);
816
817 assert_eq!(Some((7, 1)), heap.pop());
818
819 let mut vec: Vec<_> = heap.into_iter().collect();
820 vec.sort();
821 assert_eq!(vec, vec![(1, 2), (5, 4)]);
822 }
823}