math/iter/
flat_zip.rs

1pub trait IntoFlatZipIter<I> {
2    fn flat_zip(self, other: I) -> FlatZipIter<I>;
3}
4
5pub struct FlatZipIter<I> {
6    iters: Vec<I>,
7}
8
9impl<I: Iterator> IntoFlatZipIter<I> for I {
10    fn flat_zip(self, other: I) -> FlatZipIter<I> {
11        FlatZipIter {
12            iters: vec![self, other],
13        }
14    }
15}
16
17impl<I: Iterator> IntoFlatZipIter<I> for FlatZipIter<I> {
18    fn flat_zip(mut self, other: I) -> FlatZipIter<I> {
19        self.iters.push(other);
20        FlatZipIter {
21            iters: self.iters,
22        }
23    }
24}
25
26impl<I: Iterator> Iterator for FlatZipIter<I> {
27    type Item = Vec<<I as Iterator>::Item>;
28
29    fn next(&mut self) -> Option<Self::Item> {
30        match self
31            .iters
32            .iter_mut()
33            .map(|i| {
34                i.next().map_or_else(|| Err("None encountered"), |x| Ok(x))
35            })
36            .collect::<Result<Self::Item, &str>>()
37        {
38            Err(_) => None,
39            Ok(v) => Some(v),
40        }
41    }
42}
43
44#[cfg(test)]
45mod tests {
46    use crate::iter::flat_zip::IntoFlatZipIter;
47
48    #[test]
49    fn test_flat_zip() {
50        let arr1 = vec![1, 2, 3];
51        let arr2 = vec![4, 5, 6];
52        let arr3 = vec![7, 8, 9];
53        let expected_1 =
54            vec![vec![&1, &4, &7], vec![&2, &5, &8], vec![&3, &6, &9]];
55        for (i1, i2) in arr1
56            .iter()
57            .flat_zip(arr2.iter())
58            .flat_zip(arr3.iter())
59            .zip(expected_1.into_iter())
60        {
61            assert_eq!(i1, i2);
62        }
63
64        let expected_2 =
65            vec![vec![&1, &4, &7, &10], vec![&2, &5, &8, &11], vec![
66                &3, &6, &9, &12,
67            ]];
68        let arr4 = vec![10, 11, 12, 13];
69        for (i1, i2) in arr1
70            .iter()
71            .flat_zip(arr2.iter())
72            .flat_zip(arr3.iter())
73            .flat_zip(arr4.iter())
74            .zip(expected_2.into_iter())
75        {
76            assert_eq!(i1, i2);
77        }
78    }
79}
math/iter/flat_zip.rs

math/iter/
flat_zip.rs
