Merge 2 sorted iterators

merge_sorted.rs

use std::iter::{Iterator, Peekable};

pub struct MergeSorted<A: Iterator, B: Iterator> {
    a: Peekable<A>,
    b: Peekable<B>,
}

impl<A: Iterator, B: Iterator> MergeSorted<A, B> {
    fn new(a: A, b: B) -> MergeSorted<A, B> {
        MergeSorted {
            a: a.peekable(),
            b: b.peekable(),
        }
    }
}

impl<T: Ord, A: Iterator<Item = T>, B: Iterator<Item = T>> Iterator for MergeSorted<A, B> {
    type Item = T;

    #[inline]
    fn next(&mut self) -> Option<T> {
        match (self.a.peek(), self.b.peek()) {
            (None, None) => None,
            (None, Some(_)) => self.b.next(),
            (Some(_), None) => self.a.next(),
            (Some(a), Some(b)) => {
                if a < b {
                    self.a.next()
                } else {
                    self.b.next()
                }
            }
        }
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        let (a_min, a_max) = self.a.size_hint();
        let (b_min, b_max) = self.b.size_hint();
        (
            a_min + b_min,
            a_max.and_then(|a_max| b_max.map(|b_max| a_max + b_max)),
        )
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::iter::Iterator;

    #[test]
    fn smoke_test() {
        let a: Vec<usize> = vec![1, 3, 5, 7];
        let b: Vec<usize> = vec![2, 4, 6, 8];
        let r: Vec<usize> = MergeSorted::new(a.iter(), b.iter()).cloned().collect();
        assert_eq!(r, vec![1, 2, 3, 4, 5, 6, 7, 8])
    }
}