Select Iterators (with Channel iterators)

select.rs

pub struct SelectIterator<'i, I> {
    iters: Vec<&'i mut dyn Iterator<Item=I>>
}

pub fn select<'i, I>(iters: Vec<&'i mut dyn Iterator<Item=I>>)
    -> SelectIterator<'i, I>
{
    SelectIterator { iters }
}

impl<'i, I> Iterator for SelectIterator<'i, I> {
    type Item = I;
    
    fn next(&mut self) -> Option<Self::Item> {
        for iter in self.iters.iter_mut() {
            let n = iter.next();
            if n.is_some() {
                return n
            }
        }
        None
    }
}

#[cfg(test)]
mod test {
    use super::*;

    #[test]    
    fn basic_select() {
        let v1 = vec![1i32, 2];
        let mut iter1 = v1.iter();
        let v2 = vec![-1i32, -2, -3];
        let mut iter2 = v2.iter();
        
        let mut select_iter = select(vec![
            &mut iter1, &mut iter2]);
        
        assert_eq!(Some(&1), select_iter.next());
        assert_eq!(Some(&2), select_iter.next());
        assert_eq!(Some(&-1), select_iter.next());
        assert_eq!(Some(&-2), select_iter.next());
        assert_eq!(Some(&-3), select_iter.next());
        assert_eq!(None, select_iter.next());
    }
    
    #[test]
    fn with_mappings() {
        let left = vec!["a", "b", "c"];
        let mut left_merge = left.iter().map(|s| Merged::Left(s.to_string()));
        let right = vec![1i32, 2, 3];
        let mut right_merge = right.iter().map(|i| Merged::Right(*i));
        
        let mut select_iter: SelectIterator<Merged> = select(vec![
            &mut left_merge, &mut right_merge
        ]);
        
        assert_eq!(Some(Merged::Left("a".into())), select_iter.next());
        assert_eq!(Some(Merged::Left("b".into())), select_iter.next());
        assert_eq!(Some(Merged::Left("c".into())), select_iter.next());
        assert_eq!(Some(Merged::Right(1)), select_iter.next());
        assert_eq!(Some(Merged::Right(2)), select_iter.next());
        assert_eq!(Some(Merged::Right(3)), select_iter.next());
        assert_eq!(None, select_iter.next());
    }
    
    #[derive(Debug, PartialEq)]
    enum Merged {
        Left(String),
        Right(i32)
    }
    
    #[test]
    fn from_channels_try_iter() {
        use std::sync::mpsc::*;
        
        let (tx1, rx1) = channel();
        let (tx2, rx2) = channel();
        
        let mut iter1 = rx1.try_iter();
        let mut iter2 = rx2.try_iter();
        
        let mut select_iter = select(vec![
            &mut iter1, &mut iter2
        ]);
        
        tx2.send(1u32).unwrap();
        
        assert_eq!(Some(1u32), select_iter.next());
        
        tx1.send(10u32).unwrap();
        tx2.send(2u32).unwrap();
        tx1.send(20u32).unwrap();
        
        assert_eq!(Some(10u32), select_iter.next());
        assert_eq!(Some(20u32), select_iter.next());
        assert_eq!(Some(2u32), select_iter.next());
        assert_eq!(None, select_iter.next());
    }
    
    #[test]
    fn from_channels_blocking_iter() {
        use std::sync::mpsc::*;
        use std::thread;
        
        let (tx1, rx1) = channel();
        let (tx2, rx2) = channel();
        
        let (signal_tx, signal_rx) = channel();
        
        let _ = thread::spawn(move ||{
            let mut iter1 = rx1.iter();
            let mut iter2 = rx2.iter();
            
            let mut select_iter = select(vec![
                &mut iter1, &mut iter2
            ]);
        
            assert_eq!(Some(1u32), select_iter.next());
            assert_eq!(Some(10u32), select_iter.next());
            assert_eq!(Some(20u32), select_iter.next());
            assert_eq!(Some(2u32), select_iter.next());
            assert_eq!(None, select_iter.next());
            
            signal_tx.send(true).unwrap();
        });
        
        tx2.send(1u32).unwrap();
        tx1.send(10u32).unwrap();
        tx2.send(2u32).unwrap();
        tx1.send(20u32).unwrap();
        
        // should block until we get the signal
        assert!(signal_rx.recv().unwrap());
    }
}