math/iter/
concatenated_iter.rs

1pub struct ConcatenatedIter<I: Iterator> {
2    iters: Vec<I>,
3    current_iter_index: usize,
4}
5
6impl<I: Iterator> ConcatenatedIter<I> {
7    pub fn from_iters(iters: Vec<I>) -> Self {
8        ConcatenatedIter {
9            iters,
10            current_iter_index: 0,
11        }
12    }
13
14    pub fn concat_iter(self, other: I) -> Self {
15        let mut iters = self.iters;
16        iters.push(other);
17        ConcatenatedIter {
18            iters,
19            current_iter_index: self.current_iter_index,
20        }
21    }
22}
23
24impl<I: Iterator> From<Vec<I>> for ConcatenatedIter<I> {
25    fn from(iters: Vec<I>) -> Self {
26        Self::from_iters(iters)
27    }
28}
29
30impl<I: Iterator> Iterator for ConcatenatedIter<I> {
31    type Item = <I as Iterator>::Item;
32
33    fn next(&mut self) -> Option<Self::Item> {
34        match self.iters.get_mut(self.current_iter_index) {
35            None => None,
36            Some(iter) => match iter.next() {
37                Some(x) => Some(x),
38                None => {
39                    self.current_iter_index += 1;
40                    self.next()
41                }
42            },
43        }
44    }
45}
46
47pub trait IntoConcatIter: Iterator + Sized {
48    fn into_concat_iter(self, other: Self) -> ConcatenatedIter<Self> {
49        ConcatenatedIter {
50            iters: vec![self, other],
51            current_iter_index: 0,
52        }
53    }
54}
55
56impl<I: Iterator + Sized> IntoConcatIter for I {}
57
58#[cfg(test)]
59mod tests {
60    use crate::iter::{concatenated_iter::IntoConcatIter, ConcatenatedIter};
61    use std::collections::BTreeMap;
62
63    #[test]
64    fn test_from_iters() {
65        let arr1 = vec![0, 1, 2];
66        let arr2 = vec![3, 4];
67        let arr3 = vec![5, 6];
68        let mut concat_iter =
69            ConcatenatedIter::from(vec![arr1.iter(), arr2.iter(), arr3.iter()]);
70        assert_eq!(concat_iter.next(), Some(&0));
71        assert_eq!(concat_iter.next(), Some(&1));
72        assert_eq!(concat_iter.next(), Some(&2));
73        assert_eq!(concat_iter.next(), Some(&3));
74        assert_eq!(concat_iter.next(), Some(&4));
75        assert_eq!(concat_iter.next(), Some(&5));
76        assert_eq!(concat_iter.next(), Some(&6));
77        assert_eq!(concat_iter.next(), None);
78    }
79
80    #[test]
81    fn test_concat_vec_iter() {
82        let arr1 = vec![0, 1, 2];
83        let arr2 = vec![3, 4];
84        let arr3 = vec![6, 7, 8];
85        {
86            let mut concat_1_2 = arr1.iter().into_concat_iter(arr2.iter());
87            assert_eq!(concat_1_2.next(), Some(&0));
88            assert_eq!(concat_1_2.next(), Some(&1));
89            assert_eq!(concat_1_2.next(), Some(&2));
90            assert_eq!(concat_1_2.next(), Some(&3));
91            assert_eq!(concat_1_2.next(), Some(&4));
92            assert_eq!(concat_1_2.next(), None);
93        }
94        {
95            let mut concat_1_2 = arr1.iter().into_concat_iter(arr2.iter());
96            assert_eq!(concat_1_2.next(), Some(&0));
97            assert_eq!(concat_1_2.next(), Some(&1));
98            concat_1_2 = concat_1_2.concat_iter(arr3.iter());
99            assert_eq!(concat_1_2.next(), Some(&2));
100            assert_eq!(concat_1_2.next(), Some(&3));
101            assert_eq!(concat_1_2.next(), Some(&4));
102            assert_eq!(concat_1_2.next(), Some(&6));
103            assert_eq!(concat_1_2.next(), Some(&7));
104            concat_1_2 = concat_1_2.concat_iter(arr3.iter());
105            assert_eq!(concat_1_2.next(), Some(&8));
106            assert_eq!(concat_1_2.next(), Some(&6));
107            assert_eq!(concat_1_2.next(), Some(&7));
108            assert_eq!(concat_1_2.next(), Some(&8));
109            concat_1_2 = concat_1_2.concat_iter(arr1.iter());
110            assert_eq!(concat_1_2.next(), Some(&0));
111            assert_eq!(concat_1_2.next(), Some(&1));
112            assert_eq!(concat_1_2.next(), Some(&2));
113            assert_eq!(concat_1_2.next(), None);
114        }
115    }
116
117    #[test]
118    fn test_concat_btreemap_iter() {
119        let m1: BTreeMap<i32, i32> = vec![(1, 2), (3, 4)].into_iter().collect();
120        let m2: BTreeMap<i32, i32> =
121            vec![(11, 12), (13, 14), (15, 16)].into_iter().collect();
122        let mut iter = m1.iter().into_concat_iter(m2.iter());
123        assert_eq!(iter.next(), Some((&1, &2)));
124        assert_eq!(iter.next(), Some((&3, &4)));
125        assert_eq!(iter.next(), Some((&11, &12)));
126        assert_eq!(iter.next(), Some((&13, &14)));
127        assert_eq!(iter.next(), Some((&15, &16)));
128    }
129}
math/iter/concatenated_iter.rs

math/iter/
concatenated_iter.rs
