flat_zip/
flat_zip.rs

1use crate::{group::Group, groups::Groups};
2use std::iter::Fuse;
3
4pub struct FlatZip<I, K, G>
5where
6    I: Iterator<Item = (K, G)>,
7    K: Clone,
8    G: IntoIterator,
9{
10    // see FlattenCompat from std
11    groups: Fuse<Groups<I>>,
12    front: Option<Group<K, G::IntoIter>>,
13    back: Option<Group<K, G::IntoIter>>,
14}
15
16impl<I, K, G> FlatZip<I, K, G>
17where
18    I: Iterator<Item = (K, G)>,
19    K: Clone,
20    G: IntoIterator,
21{
22    pub fn new<T: IntoIterator<IntoIter = I>>(iter: T) -> Self {
23        Self {
24            groups: Groups::new(iter).fuse(),
25            front: None,
26            back: None,
27        }
28    }
29
30    fn fold_groups<B, F>(self, mut acc: B, mut f: F) -> B
31    where
32        F: FnMut(B, Group<K, G::IntoIter>) -> B,
33    {
34        if let Some(group) = self.front {
35            acc = f(acc, group);
36        }
37
38        acc = self.groups.fold(acc, &mut f);
39
40        if let Some(group) = self.back {
41            acc = f(acc, group);
42        }
43
44        acc
45    }
46
47    fn rfold_groups<B, F>(self, mut acc: B, mut f: F) -> B
48    where
49        F: FnMut(B, Group<K, G::IntoIter>) -> B,
50        I: DoubleEndedIterator,
51        G::IntoIter: DoubleEndedIterator,
52    {
53        if let Some(group) = self.back {
54            acc = f(acc, group);
55        }
56
57        acc = self.groups.rfold(acc, &mut f);
58
59        if let Some(group) = self.front {
60            acc = f(acc, group);
61        }
62
63        acc
64    }
65}
66
67impl<I, K, G, V> Iterator for FlatZip<I, K, G>
68where
69    I: Iterator<Item = (K, G)>,
70    K: Clone,
71    G: IntoIterator<Item = V>,
72{
73    type Item = (K, V);
74
75    fn next(&mut self) -> Option<Self::Item> {
76        loop {
77            if let Some(group) = &mut self.front {
78                if let Some(item) = group.next() {
79                    break Some(item);
80                }
81                self.front = None;
82            }
83
84            if let Some(group) = self.groups.next() {
85                self.front = Some(group);
86                continue;
87            }
88
89            if let Some(group) = &mut self.back {
90                if let Some(item) = group.next() {
91                    break Some(item);
92                }
93                self.back = None;
94            }
95
96            break None;
97        }
98    }
99
100    fn fold<B, F>(self, init: B, mut f: F) -> B
101    where
102        F: FnMut(B, Self::Item) -> B,
103    {
104        self.fold_groups(init, |acc, group| group.fold(acc, &mut f))
105    }
106
107    fn count(self) -> usize {
108        self.fold_groups(0, |n, group| n + group.count())
109    }
110
111    fn last(self) -> Option<Self::Item> {
112        self.fold_groups(None, |acc, group| group.last().or(acc))
113    }
114}
115
116impl<I, K, G, V> DoubleEndedIterator for FlatZip<I, K, G>
117where
118    I: DoubleEndedIterator<Item = (K, G)>,
119    K: Clone,
120    G: IntoIterator<Item = V>,
121    G::IntoIter: DoubleEndedIterator,
122{
123    fn next_back(&mut self) -> Option<Self::Item> {
124        loop {
125            if let Some(group) = &mut self.back {
126                if let Some(item) = group.next_back() {
127                    break Some(item);
128                }
129                self.back = None;
130            }
131
132            if let Some(group) = self.groups.next_back() {
133                self.back = Some(group);
134                continue;
135            }
136
137            if let Some(group) = &mut self.front {
138                if let Some(item) = group.next_back() {
139                    break Some(item);
140                }
141                self.front = None;
142            }
143
144            break None;
145        }
146    }
147
148    fn rfold<B, F>(self, init: B, mut f: F) -> B
149    where
150        F: FnMut(B, Self::Item) -> B,
151    {
152        self.rfold_groups(init, |acc, group| group.rfold(acc, &mut f))
153    }
154}
flat_zip/flat_zip.rs

flat_zip/
flat_zip.rs
