A Rust data structure to merge sorted iterators

merger.rs

use std::cmp::{Ordering, Reverse};
use std::collections::BinaryHeap;

struct Head<I: Iterator>(I, I::Item);

impl<I: Iterator> Eq for Head<I> where I::Item: Eq {}

impl<I: Iterator> PartialEq for Head<I>
where
    I::Item: Eq,
{
    fn eq(&self, other: &Self) -> bool {
        self.1.eq(&other.1)
    }
}

impl<I: Iterator> PartialOrd for Head<I>
where
    I::Item: Eq + PartialOrd,
{
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        self.1.partial_cmp(&other.1)
    }
}

impl<I: Iterator> Ord for Head<I>
where
    I::Item: Ord,
{
    fn cmp(&self, other: &Self) -> Ordering {
        self.1.cmp(&other.1)
    }
}

pub struct Merger<I: IntoIterator> {
    iters: BinaryHeap<Reverse<Head<I::IntoIter>>>,
}

impl<I: IntoIterator> Merger<I>
where
    I::Item: Ord,
{
    fn maybe_push(&mut self, mut iter: I::IntoIter) {
        if let Some(next) = iter.next() {
            self.iters.push(Reverse(Head(iter, next)));
        }
    }
}

impl<I: IntoIterator> Default for Merger<I>
where
    I::Item: Ord,
{
    fn default() -> Self {
        Self {
            iters: Default::default(),
        }
    }
}

impl<I: IntoIterator> Extend<I> for Merger<I>
where
    I::Item: Ord,
{
    fn extend<Is: IntoIterator<Item = I>>(&mut self, things: Is) {
        self.iters.extend(things.into_iter().flat_map(|into_iter| {
            let mut iter = into_iter.into_iter();
            iter.next().map(|next| Reverse(Head(iter, next)))
        }))
    }
}

impl<I: IntoIterator> Merger<I>
where
    I::Item: Ord,
{
    pub fn push(&mut self, into_iter: I) {
        self.maybe_push(into_iter.into_iter());
    }

    pub fn peek(&mut self) -> Option<&I::Item> {
        self.iters.peek().map(|head| &(head.0).1)
    }
}

impl<I: IntoIterator> Iterator for Merger<I>
where
    I::Item: Ord,
{
    type Item = I::Item;

    fn next(&mut self) -> Option<Self::Item> {
        if let Some(Reverse(Head(iter, item))) = self.iters.pop() {
            self.maybe_push(iter);
            Some(item)
        } else {
            None
        }
    }
}

impl<I: IntoIterator> From<Vec<I>> for Merger<I>
where
    I::Item: Ord,
{
    fn from(elements: Vec<I>) -> Self {
        let mut merger = Merger::default();
        merger.extend(elements.into_iter());
        merger
    }
}

impl<I: IntoIterator> FromIterator<I> for Merger<I>
where
    I::Item: Ord,
{
    #[inline]
    fn from_iter<II: IntoIterator<Item = I>>(iter: II) -> Self {
        let mut merger = Merger::default();
        merger.extend(iter.into_iter());
        merger
    }
}

#[cfg(test)]
mod test {
    use crate::Merger;

    #[test]
    fn one_element() {
        let mut merger: Merger<_> = Some(vec![1]).into_iter().collect();
        assert_eq!(merger.next(), Some(1));
        assert_eq!(merger.next(), None);
    }
    #[test]
    fn two_elements() {
        let mut merger: Merger<_> = [vec![0], vec![1]].into_iter().collect();
        assert_eq!(merger.next(), Some(0));
        assert_eq!(merger.next(), Some(1));
        assert_eq!(merger.next(), None);
        let mut merger: Merger<_> = [vec![1], vec![0]].into_iter().collect();
        assert_eq!(merger.next(), Some(0));
        assert_eq!(merger.next(), Some(1));
        assert_eq!(merger.next(), None);
    }
}