more_itertools/combining/
interleave.rs

1use std::collections::VecDeque;
2
3use crate::error::Error;
4
5
6pub struct Interleave<T> {
7    buf: VecDeque<T>,
8    buf2: VecDeque<Option<T>>,
9    iter_vec: Vec<Box<dyn Iterator<Item = Result<T,Error>>>>,
10    iter_finished: bool
11}
12
13impl<T> Iterator for Interleave<T> {
14    type Item = Result<T,Error>;
15
16    fn next(&mut self) -> Option<Self::Item> {
17        loop {
18            if self.iter_finished {
19                return None;
20            }
21
22            if self.buf.len() > 0 {
23                let ret = self.buf.pop_front().unwrap();
24                return Some(Ok(ret));
25            }
26
27            assert_eq!(0, self.buf2.len());
28
29            for i in self.iter_vec.iter_mut() {
30                match i.next() {
31                    None => {
32                        self.buf2.push_back(None);
33                    },
34                    Some(v) => {
35                        match v {
36                            Ok(ok_v) => {
37                                self.buf2.push_back(Some(ok_v));
38                            },
39                            Err(err_v) => { // upstream error
40                                self.iter_finished = true;
41                                return Some(Err(err_v));
42                            }
43                        }
44                    }
45                }
46            }
47
48            if self.buf2.iter().any(|x| match x {
49                None => { return true; },
50                Some(_) => { return false; }
51            }) {
52                self.iter_finished = true;
53            } else {
54                while self.buf2.len() > 0 {
55                    self.buf.push_back(self.buf2.pop_front().unwrap().unwrap());
56                }
57            }
58        }   
59    }
60}
61
62pub fn interleave<T>(iter_vec: Vec<Box<dyn Iterator<Item = Result<T,Error>>>>) -> Box<dyn Iterator<Item = Result<T,Error>>> 
63where T: 'static
64{
65    Box::new(Interleave {
66        buf: VecDeque::new(),
67        buf2: VecDeque::new(),
68        iter_vec,
69        iter_finished: false
70    })
71}
72
73#[cfg(test)]
74mod tests {
75
76    use crate::{error, utils::{extract_value_from_result_vec, generate_okok_iterator, generate_okokerr_iterator}};
77
78    use super::*;
79
80    #[test]
81    fn test1() {
82        let mut v = Vec::new();
83        v.push(generate_okok_iterator(vec![1,2,3]));
84        v.push(generate_okok_iterator(vec![4,5]));
85        v.push(generate_okok_iterator(vec![6,7,8]));
86
87        let ret = interleave(v).collect::<Vec<_>>();
88        assert_eq!(vec![1, 4, 6, 2, 5, 7], extract_value_from_result_vec(ret).0);
89
90
91        let mut v = Vec::new();
92        v.push(generate_okok_iterator(vec![1,2,3]));
93        v.push(generate_okokerr_iterator(vec![4], error::overflow_error("[test]".to_string())));
94        v.push(generate_okok_iterator(vec![6,7,8]));
95
96        let ret = interleave(v).collect::<Vec<_>>();
97        let ret2 = extract_value_from_result_vec(ret);
98        assert_eq!(vec![1, 4, 6], ret2.0);
99        assert_eq!(error::Kind::OverflowError, ret2.1.unwrap().kind());
100    }
101}