use core::{cmp::Ordering, marker::PhantomData, num::NonZero};
use alloc::vec::Vec;
use crate::comparator::Comparator;
use crate::cursor::CursorLendingIterator;
use crate::lending_iterator_support::{LendItem, LentItem};
use crate::seekable::{ItemToKey, Seekable};
use crate::seekable_iterators::SeekableLendingIterator;
#[derive(Debug, Clone, Copy)]
enum Direction {
Forwards,
Backwards,
}
#[derive(Debug)]
#[cfg_attr(docsrs, doc(cfg(feature = "alloc")))]
pub struct MergingIter<Key: ?Sized, Cmp, Iter> {
iterators: Vec<Iter>,
cmp: Cmp,
_key: PhantomData<fn(&Key)>,
current_iter: Option<NonZero<usize>>,
direction: Direction,
}
impl<Key, Cmp, Iter> MergingIter<Key, Cmp, Iter>
where
Key: ?Sized,
Cmp: Comparator<Key>,
Iter: SeekableLendingIterator<Key, Cmp> + ItemToKey<Key>,
{
#[inline]
#[must_use]
pub fn new(iterators: Vec<Iter>, cmp: Cmp) -> Self {
assert_ne!(
iterators.len(),
usize::MAX,
"Cannot create a MergingIter over `usize::MAX`-many iterators",
);
Self {
iterators,
cmp,
_key: PhantomData,
current_iter: None,
direction: Direction::Forwards,
}
}
}
impl<Key, Cmp, Iter> MergingIter<Key, Cmp, Iter>
where
Key: ?Sized,
Cmp: Comparator<Key>,
Iter: SeekableLendingIterator<Key, Cmp> + ItemToKey<Key>,
{
#[must_use]
fn get_current_iter_ref(&self) -> Option<&Iter> {
let current_idx = self.current_iter?.get() - 1;
#[expect(
clippy::indexing_slicing,
reason = "`self.iterators` is never truncated, \
and `self.current_idx` is always a valid idx if `Some`",
)]
Some(&self.iterators[current_idx])
}
fn find_smallest_iter(&mut self) {
let mut smallest: Option<(usize, &Key)> = None;
for (idx, iter) in self.iterators.iter().enumerate() {
if let Some(curr_item) = iter.current() {
let curr_key = Iter::item_to_key(curr_item);
if let Some((_, smallest_key)) = smallest {
if self.cmp.cmp(curr_key, smallest_key) == Ordering::Less {
smallest = Some((idx, curr_key));
}
} else {
smallest = Some((idx, curr_key));
}
} else {
}
}
#[expect(clippy::unwrap_used, reason = "MergingIter cannot have `usize::MAX` iterators")]
{
self.current_iter = smallest.map(|(idx, _)| NonZero::new(idx + 1).unwrap());
}
}
fn find_largest_iter(&mut self) {
let mut largest: Option<(usize, &Key)> = None;
for (idx, iter) in self.iterators.iter().enumerate().rev() {
if let Some(curr_item) = iter.current() {
let curr_key = Iter::item_to_key(curr_item);
if let Some((_, largest_key)) = largest {
if self.cmp.cmp(curr_key, largest_key) == Ordering::Greater {
largest = Some((idx, curr_key));
}
} else {
largest = Some((idx, curr_key));
}
} else {
}
}
#[expect(clippy::unwrap_used, reason = "MergingIter cannot have `usize::MAX` iterators")]
{
self.current_iter = largest.map(|(idx, _)| NonZero::new(idx + 1).unwrap());
}
}
fn switch_to_forwards(&mut self, current_idx: NonZero<usize>) -> &mut Iter {
let current_idx = current_idx.get() - 1;
let (iters, current_and_later) = self.iterators.split_at_mut(current_idx);
let (current_iter, other_iters) = current_and_later.split_at_mut(1);
#[expect(clippy::indexing_slicing, reason = "`current_idx` is a valid index")]
let current_iter = &mut current_iter[0];
#[expect(
clippy::unwrap_used,
reason = "the current iterator is `valid()` as an invariant",
)]
let current_key = Iter::item_to_key(current_iter.current().unwrap());
for iter in iters {
iter.seek(current_key);
if iter.current().is_some_and(|item| {
self.cmp.cmp(current_key, Iter::item_to_key(item)) == Ordering::Equal
}) {
iter.next();
}
}
for iter in other_iters {
iter.seek(current_key);
if iter.current().is_some_and(|item| {
self.cmp.cmp(current_key, Iter::item_to_key(item)) == Ordering::Equal
}) {
iter.next();
}
}
self.direction = Direction::Forwards;
current_iter
}
fn switch_to_backwards(&mut self, current_idx: NonZero<usize>) -> &mut Iter {
let current_idx = current_idx.get() - 1;
let (iters, current_and_later) = self.iterators.split_at_mut(current_idx);
let (current_iter, other_iters) = current_and_later.split_at_mut(1);
#[expect(clippy::indexing_slicing, reason = "`current_idx` is a valid index")]
let current_iter = &mut current_iter[0];
#[expect(
clippy::unwrap_used,
reason = "the current iterator is `valid()` as an invariant",
)]
let current_key = Iter::item_to_key(current_iter.current().unwrap());
for iter in iters {
iter.seek_before(current_key);
}
for iter in other_iters {
iter.seek_before(current_key);
}
self.direction = Direction::Backwards;
current_iter
}
}
impl<'lend, Key, Cmp, Iter> LendItem<'lend> for MergingIter<Key, Cmp, Iter>
where
Key: ?Sized,
Iter: LendItem<'lend>,
{
type Item = Iter::Item;
}
impl<Key, Cmp, Iter> CursorLendingIterator for MergingIter<Key, Cmp, Iter>
where
Key: ?Sized,
Cmp: Comparator<Key>,
Iter: SeekableLendingIterator<Key, Cmp> + ItemToKey<Key>,
{
#[inline]
fn valid(&self) -> bool {
self.current_iter.is_some()
}
fn next(&mut self) -> Option<LentItem<'_, Self>> {
if let Some(current_idx) = self.current_iter {
let current_iter = if matches!(self.direction, Direction::Backwards) {
self.switch_to_forwards(current_idx)
} else {
#[expect(clippy::indexing_slicing, reason = "we know that it's a valid index")]
&mut self.iterators[current_idx.get() - 1]
};
current_iter.next();
self.find_smallest_iter();
} else {
for iter in &mut self.iterators {
iter.next();
}
self.find_smallest_iter();
self.direction = Direction::Forwards;
}
self.current()
}
#[inline]
fn current(&self) -> Option<LentItem<'_, Self>> {
self.get_current_iter_ref()?.current()
}
fn prev(&mut self) -> Option<LentItem<'_, Self>> {
if let Some(current_idx) = self.current_iter {
let current_iter = if matches!(self.direction, Direction::Forwards) {
self.switch_to_backwards(current_idx)
} else {
#[expect(clippy::indexing_slicing, reason = "we know that it's a valid index")]
&mut self.iterators[current_idx.get() - 1]
};
current_iter.prev();
self.find_largest_iter();
} else {
for iter in &mut self.iterators {
iter.prev();
}
self.find_largest_iter();
self.direction = Direction::Backwards;
}
self.current()
}
}
impl<Key, Cmp, Iter> ItemToKey<Key> for MergingIter<Key, Cmp, Iter>
where
Key: ?Sized,
Iter: ItemToKey<Key>,
{
#[inline]
fn item_to_key(item: LentItem<'_, Self>) -> &'_ Key {
Iter::item_to_key(item)
}
}
impl<Key, Cmp, Iter> Seekable<Key, Cmp> for MergingIter<Key, Cmp, Iter>
where
Key: ?Sized,
Cmp: Comparator<Key>,
Iter: SeekableLendingIterator<Key, Cmp> + ItemToKey<Key>,
{
fn reset(&mut self) {
for iter in &mut self.iterators {
iter.reset();
}
self.current_iter = None;
self.direction = Direction::Forwards;
}
fn seek(&mut self, min_bound: &Key) {
for iter in &mut self.iterators {
iter.seek(min_bound);
}
self.find_smallest_iter();
self.direction = Direction::Forwards;
}
fn seek_before(&mut self, strict_upper_bound: &Key) {
for iter in &mut self.iterators {
iter.seek_before(strict_upper_bound);
}
self.find_largest_iter();
self.direction = Direction::Backwards;
}
fn seek_to_first(&mut self) {
for iter in &mut self.iterators {
iter.seek_to_first();
}
self.find_smallest_iter();
self.direction = Direction::Forwards;
}
fn seek_to_last(&mut self) {
for iter in &mut self.iterators {
iter.seek_to_last();
}
self.find_largest_iter();
self.direction = Direction::Backwards;
}
}
#[cfg(test)]
mod tests {
use alloc::vec;
use crate::{comparator::OrdComparator, test_iter::TestIter};
use super::*;
fn iteration_without_duplicates(iter: &mut MergingIter<u8, OrdComparator, TestIter<'_>>) {
assert_eq!(*iter.next().unwrap(), 0);
for i in 1..=9 {
assert!(iter.valid());
let next = iter.next().unwrap();
assert_eq!(*next, i);
}
assert!(iter.next().is_none());
for i in (0..=9).rev() {
let current = iter.current().copied();
let prev = *iter.prev().unwrap();
assert!(!current.is_some_and(|curr| curr == prev));
assert!(iter.valid());
let new_current = iter.current().unwrap();
assert_eq!(prev, i);
assert_eq!(*new_current, i);
}
iter.seek_before(&4);
for i in 4..=9 {
assert_eq!(*iter.next().unwrap(), i);
}
assert!(iter.next().is_none());
iter.seek(&5);
for i in (0..=4).rev() {
assert_eq!(*iter.prev().unwrap(), i);
}
assert!(iter.prev().is_none());
}
fn seek_tests(
merged_data: &[u8],
iter: &mut MergingIter<u8, OrdComparator, TestIter<'_>>,
) {
assert!(merged_data.is_sorted());
iter.reset();
let mut data_iter = merged_data.iter();
while let Some(item) = iter.next() {
assert_eq!(item, data_iter.next().unwrap());
}
assert!(data_iter.next().is_none());
iter.seek_to_first();
assert_eq!(*iter.current().unwrap(), 0);
iter.seek(&1);
assert_eq!(*iter.current().unwrap(), 1);
iter.seek(&8);
assert_eq!(*iter.current().unwrap(), 8);
iter.seek(&10);
assert_eq!(*iter.current().unwrap(), 99);
iter.seek_before(&92);
assert_eq!(*iter.current().unwrap(), 9);
iter.seek(&9);
assert_eq!(*iter.current().unwrap(), 9);
iter.seek_before(&99);
assert_eq!(*iter.current().unwrap(), 9);
iter.seek_before(&100);
assert_eq!(*iter.current().unwrap(), 99);
iter.seek_before(&1);
assert_eq!(*iter.current().unwrap(), 0);
iter.seek_before(&0);
assert!(!iter.valid());
iter.seek(&100);
assert!(!iter.valid());
iter.seek(&99);
assert_eq!(*iter.current().unwrap(), 99);
iter.seek_to_last();
assert_eq!(*iter.current().unwrap(), 99);
iter.seek(&0);
assert_eq!(*iter.current().unwrap(), 0);
iter.seek_before(&4);
assert_eq!(*iter.current().unwrap(), 3);
}
#[test]
fn single_merged() {
let data: &[u8] = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9].as_slice();
let mut iter = MergingIter::new(vec![TestIter::new(data).unwrap()], OrdComparator);
iteration_without_duplicates(&mut iter);
}
#[test]
fn seek_single_merged() {
let data: &[u8] = [0, 1, 2, 3, 4, 4, 4, 4, 4, 4, 5, 6, 7, 8, 9, 99].as_slice();
let mut iter = MergingIter::new(vec![TestIter::new(data).unwrap()], OrdComparator);
seek_tests(data, &mut iter);
}
#[test]
fn three_merged_interleaved() {
let data_one: &[u8] = [0, 3, 6, 7].as_slice();
let data_two: &[u8] = [1, 5, 8].as_slice();
let data_three: &[u8] = [2, 4, 9].as_slice();
let mut iter = MergingIter::new(
vec![
TestIter::new(data_one).unwrap(),
TestIter::new(data_two).unwrap(),
TestIter::new(data_three).unwrap(),
],
OrdComparator,
);
iteration_without_duplicates(&mut iter);
}
#[test]
fn seek_three_merged_interleaved() {
let data_one: &[u8] = [0, 3, 6, 7].as_slice();
let data_two: &[u8] = [1, 5, 8].as_slice();
let data_three: &[u8] = [2, 4, 9, 99].as_slice();
let merged_data: &[u8] = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 99].as_slice();
let mut iter = MergingIter::new(
vec![
TestIter::new(data_one).unwrap(),
TestIter::new(data_two).unwrap(),
TestIter::new(data_three).unwrap(),
],
OrdComparator,
);
seek_tests(merged_data, &mut iter);
}
#[test]
fn three_merged_chained() {
let data_one: &[u8] = [0, 1, 2, 3].as_slice();
let data_two: &[u8] = [7, 8, 9].as_slice();
let data_three: &[u8] = [4, 5, 6].as_slice();
let mut iter = MergingIter::new(
vec![
TestIter::new(data_one).unwrap(),
TestIter::new(data_two).unwrap(),
TestIter::new(data_three).unwrap(),
],
OrdComparator,
);
iteration_without_duplicates(&mut iter);
}
#[test]
fn seek_three_merged_chained() {
let data_one: &[u8] = [0, 1, 2, 3].as_slice();
let data_two: &[u8] = [7, 8, 9, 99].as_slice();
let data_three: &[u8] = [4, 5, 6].as_slice();
let merged_data: &[u8] = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 99].as_slice();
let mut iter = MergingIter::new(
vec![
TestIter::new(data_one).unwrap(),
TestIter::new(data_two).unwrap(),
TestIter::new(data_three).unwrap(),
],
OrdComparator,
);
seek_tests(merged_data, &mut iter);
}
#[test]
fn single_duplicates_defined() {
let data = &[1, 2, 2, 2, 2, 3];
let mut iter = MergingIter::new(
vec![
TestIter::new(data).unwrap(),
],
OrdComparator,
);
let mut data_iter = data.iter();
while let Some(item) = iter.next() {
assert_eq!(item, data_iter.next().unwrap());
}
assert!(data_iter.next().is_none());
let mut data_iter = data.iter().rev();
while let Some(item) = iter.prev() {
assert_eq!(item, data_iter.next().unwrap());
}
assert!(data_iter.next().is_none());
iter.seek(&1);
for _ in 0..4 {
assert_eq!(iter.next(), Some(&2));
}
iter.seek_before(&3);
for _ in 0..4 {
assert_eq!(iter.current(), Some(&2));
iter.prev();
}
}
#[test]
fn single_duplicates_unspecified() {
let data = &[1, 2, 2, 2, 2, 3];
let mut iter = MergingIter::new(
vec![
TestIter::new(data).unwrap(),
],
OrdComparator,
);
for advance in 1..=4 {
iter.seek(&1);
for _ in 0..advance {
assert_eq!(iter.next(), Some(&2));
}
for _ in 0..advance {
assert_eq!(iter.current(), Some(&2));
iter.prev();
}
}
}
#[test]
fn two_duplicates_defined() {
let data_one = &[1, 2, 2, 3];
let data_two = &[0, 2, 2, 5];
let merged_data = &[0, 1, 2, 2, 2, 2, 3, 5];
let mut iter = MergingIter::new(
vec![
TestIter::new(data_one).unwrap(),
TestIter::new(data_two).unwrap(),
],
OrdComparator,
);
let mut data_iter = merged_data.iter();
while let Some(item) = iter.next() {
assert_eq!(item, data_iter.next().unwrap());
}
assert!(data_iter.next().is_none());
let mut data_iter = merged_data.iter().rev();
while let Some(item) = iter.prev() {
assert_eq!(item, data_iter.next().unwrap());
}
assert!(data_iter.next().is_none());
iter.seek(&1);
for _ in 0..4 {
assert_eq!(iter.next(), Some(&2));
}
iter.seek_before(&3);
for _ in 0..4 {
assert_eq!(iter.current(), Some(&2));
iter.prev();
}
}
#[test]
fn seek_two_duplicates_defined() {
let data_one = &[1, 2, 2, 3];
let data_two = &[0, 2, 2, 5];
let mut iter = MergingIter::new(
vec![
TestIter::new(data_one).unwrap(),
TestIter::new(data_two).unwrap(),
],
OrdComparator,
);
iter.seek_to_first();
assert_eq!(*iter.current().unwrap(), 0);
iter.seek(&1);
assert_eq!(*iter.current().unwrap(), 1);
iter.seek(&3);
assert_eq!(*iter.current().unwrap(), 3);
iter.seek(&4);
assert_eq!(*iter.current().unwrap(), 5);
iter.seek_before(&4);
assert_eq!(*iter.current().unwrap(), 3);
iter.seek(&5);
assert_eq!(*iter.current().unwrap(), 5);
iter.seek_before(&2);
assert_eq!(*iter.current().unwrap(), 1);
iter.seek_before(&5);
assert_eq!(*iter.current().unwrap(), 3);
iter.seek_before(&1);
assert_eq!(*iter.current().unwrap(), 0);
assert_eq!(*iter.next().unwrap(), 1);
assert_eq!(*iter.next().unwrap(), 2);
iter.seek_before(&0);
assert!(!iter.valid());
iter.seek(&100);
assert!(!iter.valid());
iter.seek(&2);
assert_eq!(*iter.current().unwrap(), 2);
iter.seek(&3);
assert_eq!(*iter.current().unwrap(), 3);
assert_eq!(*iter.prev().unwrap(), 2);
iter.seek_before(&2);
assert_eq!(*iter.current().unwrap(), 1);
iter.seek_to_last();
assert_eq!(*iter.current().unwrap(), 5);
iter.seek_before(&2);
assert_eq!(*iter.current().unwrap(), 1);
assert_eq!(*iter.next().unwrap(), 2);
iter.seek(&2);
assert_eq!(*iter.current().unwrap(), 2);
}
#[test]
fn two_duplicates_unspecified() {
let data_one = &[1, 2, 2, 3];
let data_two = &[0, 2, 2, 5];
let mut iter = MergingIter::new(
vec![
TestIter::new(data_one).unwrap(),
TestIter::new(data_two).unwrap(),
],
OrdComparator,
);
assert_eq!(*iter.next().unwrap(), 0);
assert_eq!(*iter.next().unwrap(), 1);
assert_eq!(*iter.next().unwrap(), 2);
assert_eq!(*iter.next().unwrap(), 2);
assert_eq!(*iter.prev().unwrap(), 2);
assert_eq!(*iter.prev().unwrap(), 1);
assert_eq!(*iter.next().unwrap(), 2);
assert_eq!(*iter.next().unwrap(), 2);
assert_eq!(*iter.next().unwrap(), 2);
assert_eq!(*iter.prev().unwrap(), 1);
assert_eq!(*iter.next().unwrap(), 2);
assert_eq!(*iter.next().unwrap(), 2);
assert_eq!(*iter.next().unwrap(), 2);
assert_eq!(*iter.next().unwrap(), 2);
assert_eq!(*iter.next().unwrap(), 3);
assert_eq!(*iter.prev().unwrap(), 2);
assert_eq!(*iter.prev().unwrap(), 2);
assert_eq!(*iter.prev().unwrap(), 2);
assert_eq!(*iter.prev().unwrap(), 2);
assert_eq!(*iter.prev().unwrap(), 1);
assert_eq!(*iter.next().unwrap(), 2);
assert_eq!(*iter.next().unwrap(), 2);
assert_eq!(*iter.next().unwrap(), 2);
assert_eq!(*iter.next().unwrap(), 2);
assert_eq!(*iter.prev().unwrap(), 2);
assert_eq!(*iter.prev().unwrap(), 1);
iter.seek(&3);
assert_eq!(*iter.prev().unwrap(), 2);
assert_eq!(*iter.prev().unwrap(), 2);
assert_eq!(*iter.prev().unwrap(), 2);
assert_eq!(*iter.prev().unwrap(), 2);
assert_eq!(*iter.prev().unwrap(), 1);
iter.seek(&3);
assert_eq!(*iter.prev().unwrap(), 2);
assert_eq!(*iter.prev().unwrap(), 2);
assert_eq!(*iter.prev().unwrap(), 2);
assert_eq!(*iter.prev().unwrap(), 2);
assert_eq!(*iter.next().unwrap(), 2);
assert_eq!(*iter.next().unwrap(), 3);
}
}