1use core::{cmp::Ordering, marker::PhantomData, num::NonZero};
2use alloc::vec::Vec;
3
4use crate::comparator::Comparator;
5use crate::cursor::CursorLendingIterator;
6use crate::lending_iterator_support::{LendItem, LentItem};
7use crate::seekable::{ItemToKey, Seekable};
8use crate::seekable_iterators::SeekableLendingIterator;
9
10
11#[derive(Debug, Clone, Copy)]
12enum Direction {
13 Forwards,
14 Backwards,
15}
16
17#[derive(Debug)]
79#[cfg_attr(docsrs, doc(cfg(feature = "alloc")))]
80pub struct MergingIter<Key: ?Sized, Cmp, Iter> {
81 iterators: Vec<Iter>,
82 cmp: Cmp,
83 _key: PhantomData<fn(&Key)>,
86 current_iter: Option<NonZero<usize>>,
94 direction: Direction,
100}
101
102impl<Key, Cmp, Iter> MergingIter<Key, Cmp, Iter>
103where
104 Key: ?Sized,
105 Cmp: Comparator<Key>,
106 Iter: SeekableLendingIterator<Key, Cmp> + ItemToKey<Key>,
107{
108 #[inline]
120 #[must_use]
121 pub fn new(iterators: Vec<Iter>, cmp: Cmp) -> Self {
122 assert_ne!(
123 iterators.len(),
124 usize::MAX,
125 "Cannot create a MergingIter over `usize::MAX`-many iterators",
126 );
127
128 Self {
129 iterators,
130 cmp,
131 _key: PhantomData,
132 current_iter: None,
133 direction: Direction::Forwards,
134 }
135 }
136}
137
138impl<Key, Cmp, Iter> MergingIter<Key, Cmp, Iter>
139where
140 Key: ?Sized,
141 Cmp: Comparator<Key>,
142 Iter: SeekableLendingIterator<Key, Cmp> + ItemToKey<Key>,
143{
144 #[must_use]
145 fn get_current_iter_ref(&self) -> Option<&Iter> {
146 let current_idx = self.current_iter?.get() - 1;
147
148 #[expect(
149 clippy::indexing_slicing,
150 reason = "`self.iterators` is never truncated, \
151 and `self.current_idx` is always a valid idx if `Some`",
152 )]
153 Some(&self.iterators[current_idx])
154 }
155
156 fn find_smallest_iter(&mut self) {
159 let mut smallest: Option<(usize, &Key)> = None;
160
161 for (idx, iter) in self.iterators.iter().enumerate() {
162 if let Some(curr_item) = iter.current() {
163 let curr_key = Iter::item_to_key(curr_item);
164 if let Some((_, smallest_key)) = smallest {
165 if self.cmp.cmp(curr_key, smallest_key) == Ordering::Less {
166 smallest = Some((idx, curr_key));
168 }
169 } else {
170 smallest = Some((idx, curr_key));
172 }
173 } else {
174 }
176 }
177
178 #[expect(clippy::unwrap_used, reason = "MergingIter cannot have `usize::MAX` iterators")]
179 {
180 self.current_iter = smallest.map(|(idx, _)| NonZero::new(idx + 1).unwrap());
181 }
182 }
183
184 fn find_largest_iter(&mut self) {
187 let mut largest: Option<(usize, &Key)> = None;
188
189 for (idx, iter) in self.iterators.iter().enumerate().rev() {
190 if let Some(curr_item) = iter.current() {
191 let curr_key = Iter::item_to_key(curr_item);
192 if let Some((_, largest_key)) = largest {
193 if self.cmp.cmp(curr_key, largest_key) == Ordering::Greater {
194 largest = Some((idx, curr_key));
196 }
197 } else {
198 largest = Some((idx, curr_key));
200 }
201 } else {
202 }
204 }
205
206 #[expect(clippy::unwrap_used, reason = "MergingIter cannot have `usize::MAX` iterators")]
207 {
208 self.current_iter = largest.map(|(idx, _)| NonZero::new(idx + 1).unwrap());
209 }
210 }
211
212 fn switch_to_forwards(&mut self, current_idx: NonZero<usize>) -> &mut Iter {
216 let current_idx = current_idx.get() - 1;
217
218 let (iters, current_and_later) = self.iterators.split_at_mut(current_idx);
220 let (current_iter, other_iters) = current_and_later.split_at_mut(1);
221 #[expect(clippy::indexing_slicing, reason = "`current_idx` is a valid index")]
222 let current_iter = &mut current_iter[0];
223 #[expect(
224 clippy::unwrap_used,
225 reason = "the current iterator is `valid()` as an invariant",
226 )]
227 let current_key = Iter::item_to_key(current_iter.current().unwrap());
228
229 for iter in iters {
230 iter.seek(current_key);
231
232 if iter.current().is_some_and(|item| {
234 self.cmp.cmp(current_key, Iter::item_to_key(item)) == Ordering::Equal
235 }) {
236 iter.next();
237 }
238 }
239
240 for iter in other_iters {
241 iter.seek(current_key);
242
243 if iter.current().is_some_and(|item| {
244 self.cmp.cmp(current_key, Iter::item_to_key(item)) == Ordering::Equal
245 }) {
246 iter.next();
247 }
248 }
249
250 self.direction = Direction::Forwards;
251
252 current_iter
253 }
254
255 fn switch_to_backwards(&mut self, current_idx: NonZero<usize>) -> &mut Iter {
259 let current_idx = current_idx.get() - 1;
260
261 let (iters, current_and_later) = self.iterators.split_at_mut(current_idx);
263 let (current_iter, other_iters) = current_and_later.split_at_mut(1);
264 #[expect(clippy::indexing_slicing, reason = "`current_idx` is a valid index")]
265 let current_iter = &mut current_iter[0];
266 #[expect(
267 clippy::unwrap_used,
268 reason = "the current iterator is `valid()` as an invariant",
269 )]
270 let current_key = Iter::item_to_key(current_iter.current().unwrap());
271
272 for iter in iters {
273 iter.seek_before(current_key);
274 }
275 for iter in other_iters {
276 iter.seek_before(current_key);
277 }
278
279 self.direction = Direction::Backwards;
280
281 current_iter
282 }
283}
284
285impl<'lend, Key, Cmp, Iter> LendItem<'lend> for MergingIter<Key, Cmp, Iter>
286where
287 Key: ?Sized,
288 Iter: LendItem<'lend>,
289{
290 type Item = Iter::Item;
291}
292
293impl<Key, Cmp, Iter> CursorLendingIterator for MergingIter<Key, Cmp, Iter>
294where
295 Key: ?Sized,
296 Cmp: Comparator<Key>,
297 Iter: SeekableLendingIterator<Key, Cmp> + ItemToKey<Key>,
298{
299 #[inline]
300 fn valid(&self) -> bool {
301 self.current_iter.is_some()
302 }
303
304 fn next(&mut self) -> Option<LentItem<'_, Self>> {
305 if let Some(current_idx) = self.current_iter {
306 let current_iter = if matches!(self.direction, Direction::Backwards) {
307 self.switch_to_forwards(current_idx)
308 } else {
309 #[expect(clippy::indexing_slicing, reason = "we know that it's a valid index")]
310 &mut self.iterators[current_idx.get() - 1]
311 };
312
313 current_iter.next();
316 self.find_smallest_iter();
318
319 } else {
320 for iter in &mut self.iterators {
324 iter.next();
325 }
326
327 self.find_smallest_iter();
328 self.direction = Direction::Forwards;
329 }
330
331 self.current()
332 }
333
334 #[inline]
335 fn current(&self) -> Option<LentItem<'_, Self>> {
336 self.get_current_iter_ref()?.current()
337 }
338
339 fn prev(&mut self) -> Option<LentItem<'_, Self>> {
347 if let Some(current_idx) = self.current_iter {
348 let current_iter = if matches!(self.direction, Direction::Forwards) {
349 self.switch_to_backwards(current_idx)
350 } else {
351 #[expect(clippy::indexing_slicing, reason = "we know that it's a valid index")]
352 &mut self.iterators[current_idx.get() - 1]
353 };
354
355 current_iter.prev();
357 self.find_largest_iter();
359
360 } else {
361 for iter in &mut self.iterators {
365 iter.prev();
366 }
367
368 self.find_largest_iter();
369 self.direction = Direction::Backwards;
370 }
371
372 self.current()
373 }
374}
375
376impl<Key, Cmp, Iter> ItemToKey<Key> for MergingIter<Key, Cmp, Iter>
377where
378 Key: ?Sized,
379 Iter: ItemToKey<Key>,
380{
381 #[inline]
382 fn item_to_key(item: LentItem<'_, Self>) -> &'_ Key {
383 Iter::item_to_key(item)
384 }
385}
386
387impl<Key, Cmp, Iter> Seekable<Key, Cmp> for MergingIter<Key, Cmp, Iter>
388where
389 Key: ?Sized,
390 Cmp: Comparator<Key>,
391 Iter: SeekableLendingIterator<Key, Cmp> + ItemToKey<Key>,
392{
393 fn reset(&mut self) {
394 for iter in &mut self.iterators {
395 iter.reset();
396 }
397 self.current_iter = None;
398 self.direction = Direction::Forwards;
399 }
400
401 fn seek(&mut self, min_bound: &Key) {
402 for iter in &mut self.iterators {
403 iter.seek(min_bound);
404 }
405
406 self.find_smallest_iter();
407 self.direction = Direction::Forwards;
408 }
409
410 fn seek_before(&mut self, strict_upper_bound: &Key) {
423 for iter in &mut self.iterators {
424 iter.seek_before(strict_upper_bound);
425 }
426
427 self.find_largest_iter();
428 self.direction = Direction::Backwards;
429 }
430
431 fn seek_to_first(&mut self) {
432 for iter in &mut self.iterators {
433 iter.seek_to_first();
434 }
435
436 self.find_smallest_iter();
437 self.direction = Direction::Forwards;
438 }
439
440 fn seek_to_last(&mut self) {
447 for iter in &mut self.iterators {
448 iter.seek_to_last();
449 }
450
451 self.find_largest_iter();
452 self.direction = Direction::Backwards;
453 }
454}
455
456
457#[cfg(test)]
458mod tests {
459 use alloc::vec;
460 use crate::{comparator::OrdComparator, test_iter::TestIter};
461 use super::*;
462
463 fn iteration_without_duplicates(iter: &mut MergingIter<u8, OrdComparator, TestIter<'_>>) {
465 assert_eq!(*iter.next().unwrap(), 0);
466
467 for i in 1..=9 {
468 assert!(iter.valid());
469 let next = iter.next().unwrap();
470 assert_eq!(*next, i);
471 }
472
473 assert!(iter.next().is_none());
474
475 for i in (0..=9).rev() {
476 let current = iter.current().copied();
477 let prev = *iter.prev().unwrap();
478 assert!(!current.is_some_and(|curr| curr == prev));
479
480 assert!(iter.valid());
481
482 let new_current = iter.current().unwrap();
483
484 assert_eq!(prev, i);
485 assert_eq!(*new_current, i);
486 }
487
488 iter.seek_before(&4);
489 for i in 4..=9 {
490 assert_eq!(*iter.next().unwrap(), i);
491 }
492 assert!(iter.next().is_none());
493
494 iter.seek(&5);
495 for i in (0..=4).rev() {
496 assert_eq!(*iter.prev().unwrap(), i);
497 }
498 assert!(iter.prev().is_none());
499 }
500
501 fn seek_tests(
506 merged_data: &[u8],
507 iter: &mut MergingIter<u8, OrdComparator, TestIter<'_>>,
508 ) {
509 assert!(merged_data.is_sorted());
510
511 iter.reset();
512 let mut data_iter = merged_data.iter();
513 while let Some(item) = iter.next() {
514 assert_eq!(item, data_iter.next().unwrap());
515 }
516 assert!(data_iter.next().is_none());
517
518
519 iter.seek_to_first();
520 assert_eq!(*iter.current().unwrap(), 0);
521
522 iter.seek(&1);
523 assert_eq!(*iter.current().unwrap(), 1);
524
525 iter.seek(&8);
526 assert_eq!(*iter.current().unwrap(), 8);
527
528 iter.seek(&10);
529 assert_eq!(*iter.current().unwrap(), 99);
530
531 iter.seek_before(&92);
532 assert_eq!(*iter.current().unwrap(), 9);
533
534 iter.seek(&9);
535 assert_eq!(*iter.current().unwrap(), 9);
536
537 iter.seek_before(&99);
538 assert_eq!(*iter.current().unwrap(), 9);
539
540 iter.seek_before(&100);
541 assert_eq!(*iter.current().unwrap(), 99);
542
543 iter.seek_before(&1);
544 assert_eq!(*iter.current().unwrap(), 0);
545
546 iter.seek_before(&0);
547 assert!(!iter.valid());
548
549 iter.seek(&100);
550 assert!(!iter.valid());
551
552 iter.seek(&99);
553 assert_eq!(*iter.current().unwrap(), 99);
554
555 iter.seek_to_last();
556 assert_eq!(*iter.current().unwrap(), 99);
557
558 iter.seek(&0);
559 assert_eq!(*iter.current().unwrap(), 0);
560
561 iter.seek_before(&4);
562 assert_eq!(*iter.current().unwrap(), 3);
563 }
564
565 #[test]
566 fn single_merged() {
567 let data: &[u8] = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9].as_slice();
568 let mut iter = MergingIter::new(vec![TestIter::new(data).unwrap()], OrdComparator);
569
570 iteration_without_duplicates(&mut iter);
571 }
572
573 #[test]
574 fn seek_single_merged() {
575 let data: &[u8] = [0, 1, 2, 3, 4, 4, 4, 4, 4, 4, 5, 6, 7, 8, 9, 99].as_slice();
576 let mut iter = MergingIter::new(vec![TestIter::new(data).unwrap()], OrdComparator);
577
578 seek_tests(data, &mut iter);
579 }
580
581 #[test]
582 fn three_merged_interleaved() {
583 let data_one: &[u8] = [0, 3, 6, 7].as_slice();
584 let data_two: &[u8] = [1, 5, 8].as_slice();
585 let data_three: &[u8] = [2, 4, 9].as_slice();
586 let mut iter = MergingIter::new(
587 vec![
588 TestIter::new(data_one).unwrap(),
589 TestIter::new(data_two).unwrap(),
590 TestIter::new(data_three).unwrap(),
591 ],
592 OrdComparator,
593 );
594
595 iteration_without_duplicates(&mut iter);
596 }
597
598 #[test]
599 fn seek_three_merged_interleaved() {
600 let data_one: &[u8] = [0, 3, 6, 7].as_slice();
601 let data_two: &[u8] = [1, 5, 8].as_slice();
602 let data_three: &[u8] = [2, 4, 9, 99].as_slice();
603 let merged_data: &[u8] = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 99].as_slice();
604 let mut iter = MergingIter::new(
605 vec![
606 TestIter::new(data_one).unwrap(),
607 TestIter::new(data_two).unwrap(),
608 TestIter::new(data_three).unwrap(),
609 ],
610 OrdComparator,
611 );
612
613 seek_tests(merged_data, &mut iter);
614 }
615
616 #[test]
617 fn three_merged_chained() {
618 let data_one: &[u8] = [0, 1, 2, 3].as_slice();
619 let data_two: &[u8] = [7, 8, 9].as_slice();
620 let data_three: &[u8] = [4, 5, 6].as_slice();
621 let mut iter = MergingIter::new(
622 vec![
623 TestIter::new(data_one).unwrap(),
624 TestIter::new(data_two).unwrap(),
625 TestIter::new(data_three).unwrap(),
626 ],
627 OrdComparator,
628 );
629
630 iteration_without_duplicates(&mut iter);
631 }
632
633 #[test]
634 fn seek_three_merged_chained() {
635 let data_one: &[u8] = [0, 1, 2, 3].as_slice();
636 let data_two: &[u8] = [7, 8, 9, 99].as_slice();
637 let data_three: &[u8] = [4, 5, 6].as_slice();
638 let merged_data: &[u8] = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 99].as_slice();
639 let mut iter = MergingIter::new(
640 vec![
641 TestIter::new(data_one).unwrap(),
642 TestIter::new(data_two).unwrap(),
643 TestIter::new(data_three).unwrap(),
644 ],
645 OrdComparator,
646 );
647
648 seek_tests(merged_data, &mut iter);
649 }
650
651 #[test]
653 fn single_duplicates_defined() {
654 let data = &[1, 2, 2, 2, 2, 3];
655 let mut iter = MergingIter::new(
656 vec![
657 TestIter::new(data).unwrap(),
658 ],
659 OrdComparator,
660 );
661
662 let mut data_iter = data.iter();
663 while let Some(item) = iter.next() {
664 assert_eq!(item, data_iter.next().unwrap());
665 }
666 assert!(data_iter.next().is_none());
667
668 let mut data_iter = data.iter().rev();
669 while let Some(item) = iter.prev() {
670 assert_eq!(item, data_iter.next().unwrap());
671 }
672 assert!(data_iter.next().is_none());
673
674 iter.seek(&1);
675 for _ in 0..4 {
676 assert_eq!(iter.next(), Some(&2));
677 }
678
679 iter.seek_before(&3);
680 for _ in 0..4 {
681 assert_eq!(iter.current(), Some(&2));
682 iter.prev();
683 }
684 }
685
686 #[test]
688 fn single_duplicates_unspecified() {
689 let data = &[1, 2, 2, 2, 2, 3];
690 let mut iter = MergingIter::new(
691 vec![
692 TestIter::new(data).unwrap(),
693 ],
694 OrdComparator,
695 );
696
697 for advance in 1..=4 {
698 iter.seek(&1);
699 for _ in 0..advance {
700 assert_eq!(iter.next(), Some(&2));
701 }
702 for _ in 0..advance {
703 assert_eq!(iter.current(), Some(&2));
704 iter.prev();
705 }
706 }
707 }
708
709 #[test]
711 fn two_duplicates_defined() {
712 let data_one = &[1, 2, 2, 3];
713 let data_two = &[0, 2, 2, 5];
714 let merged_data = &[0, 1, 2, 2, 2, 2, 3, 5];
715 let mut iter = MergingIter::new(
716 vec![
717 TestIter::new(data_one).unwrap(),
718 TestIter::new(data_two).unwrap(),
719 ],
720 OrdComparator,
721 );
722
723 let mut data_iter = merged_data.iter();
724 while let Some(item) = iter.next() {
725 assert_eq!(item, data_iter.next().unwrap());
726 }
727 assert!(data_iter.next().is_none());
728
729 let mut data_iter = merged_data.iter().rev();
730 while let Some(item) = iter.prev() {
731 assert_eq!(item, data_iter.next().unwrap());
732 }
733 assert!(data_iter.next().is_none());
734
735 iter.seek(&1);
736 for _ in 0..4 {
737 assert_eq!(iter.next(), Some(&2));
738 }
739
740 iter.seek_before(&3);
741 for _ in 0..4 {
742 assert_eq!(iter.current(), Some(&2));
743 iter.prev();
744 }
745 }
746
747 #[test]
749 fn seek_two_duplicates_defined() {
750 let data_one = &[1, 2, 2, 3];
751 let data_two = &[0, 2, 2, 5];
752 let mut iter = MergingIter::new(
753 vec![
754 TestIter::new(data_one).unwrap(),
755 TestIter::new(data_two).unwrap(),
756 ],
757 OrdComparator,
758 );
759
760 iter.seek_to_first();
761 assert_eq!(*iter.current().unwrap(), 0);
762
763 iter.seek(&1);
764 assert_eq!(*iter.current().unwrap(), 1);
765
766 iter.seek(&3);
767 assert_eq!(*iter.current().unwrap(), 3);
768
769 iter.seek(&4);
770 assert_eq!(*iter.current().unwrap(), 5);
771
772 iter.seek_before(&4);
773 assert_eq!(*iter.current().unwrap(), 3);
774
775 iter.seek(&5);
776 assert_eq!(*iter.current().unwrap(), 5);
777
778 iter.seek_before(&2);
779 assert_eq!(*iter.current().unwrap(), 1);
780
781 iter.seek_before(&5);
782 assert_eq!(*iter.current().unwrap(), 3);
783
784 iter.seek_before(&1);
785 assert_eq!(*iter.current().unwrap(), 0);
786 assert_eq!(*iter.next().unwrap(), 1);
787 assert_eq!(*iter.next().unwrap(), 2);
788
789 iter.seek_before(&0);
790 assert!(!iter.valid());
791
792 iter.seek(&100);
793 assert!(!iter.valid());
794
795 iter.seek(&2);
796 assert_eq!(*iter.current().unwrap(), 2);
797
798 iter.seek(&3);
799 assert_eq!(*iter.current().unwrap(), 3);
800 assert_eq!(*iter.prev().unwrap(), 2);
801
802 iter.seek_before(&2);
803 assert_eq!(*iter.current().unwrap(), 1);
804
805 iter.seek_to_last();
806 assert_eq!(*iter.current().unwrap(), 5);
807
808 iter.seek_before(&2);
809 assert_eq!(*iter.current().unwrap(), 1);
810 assert_eq!(*iter.next().unwrap(), 2);
811
812 iter.seek(&2);
813 assert_eq!(*iter.current().unwrap(), 2);
814 }
815
816 #[test]
818 fn two_duplicates_unspecified() {
819 let data_one = &[1, 2, 2, 3];
820 let data_two = &[0, 2, 2, 5];
821 let mut iter = MergingIter::new(
822 vec![
823 TestIter::new(data_one).unwrap(),
824 TestIter::new(data_two).unwrap(),
825 ],
826 OrdComparator,
827 );
828
829 assert_eq!(*iter.next().unwrap(), 0);
830 assert_eq!(*iter.next().unwrap(), 1);
831
832 assert_eq!(*iter.next().unwrap(), 2);
833 assert_eq!(*iter.next().unwrap(), 2);
834
835 assert_eq!(*iter.prev().unwrap(), 2);
837 assert_eq!(*iter.prev().unwrap(), 1);
838
839 assert_eq!(*iter.next().unwrap(), 2);
840 assert_eq!(*iter.next().unwrap(), 2);
841 assert_eq!(*iter.next().unwrap(), 2);
842
843 assert_eq!(*iter.prev().unwrap(), 1);
845
846 assert_eq!(*iter.next().unwrap(), 2);
847 assert_eq!(*iter.next().unwrap(), 2);
848 assert_eq!(*iter.next().unwrap(), 2);
849 assert_eq!(*iter.next().unwrap(), 2);
850
851 assert_eq!(*iter.next().unwrap(), 3);
852
853 assert_eq!(*iter.prev().unwrap(), 2);
855 assert_eq!(*iter.prev().unwrap(), 2);
856 assert_eq!(*iter.prev().unwrap(), 2);
857 assert_eq!(*iter.prev().unwrap(), 2);
858 assert_eq!(*iter.prev().unwrap(), 1);
859
860 assert_eq!(*iter.next().unwrap(), 2);
861 assert_eq!(*iter.next().unwrap(), 2);
862 assert_eq!(*iter.next().unwrap(), 2);
863 assert_eq!(*iter.next().unwrap(), 2);
864
865 assert_eq!(*iter.prev().unwrap(), 2);
867 assert_eq!(*iter.prev().unwrap(), 1);
868
869 iter.seek(&3);
870 assert_eq!(*iter.prev().unwrap(), 2);
872 assert_eq!(*iter.prev().unwrap(), 2);
873 assert_eq!(*iter.prev().unwrap(), 2);
874 assert_eq!(*iter.prev().unwrap(), 2);
875 assert_eq!(*iter.prev().unwrap(), 1);
876
877 iter.seek(&3);
878
879 assert_eq!(*iter.prev().unwrap(), 2);
881 assert_eq!(*iter.prev().unwrap(), 2);
882 assert_eq!(*iter.prev().unwrap(), 2);
883 assert_eq!(*iter.prev().unwrap(), 2);
884
885 assert_eq!(*iter.next().unwrap(), 2);
887 assert_eq!(*iter.next().unwrap(), 3);
888 }
889}