toodee/
flattenexact.rs

1#![forbid(unsafe_code)]
2
3#![allow(missing_debug_implementations)]
4
5use crate::iter::TooDeeIterator;
6
7/// An iterator that behaves like `core::iter::adapters::Flatten` but has the added advantage of implementing
8/// `ExactSizeIterator` (we know how many cells there are per row in a `TooDee` array).
9pub struct FlattenExact<I>
10where
11    I : ExactSizeIterator + DoubleEndedIterator + TooDeeIterator,
12    I::Item : IntoIterator,
13    <I::Item as IntoIterator>::IntoIter : DoubleEndedIterator + ExactSizeIterator,
14{
15    iter: I,
16    frontiter: Option<<I::Item as IntoIterator>::IntoIter>,
17    backiter: Option<<I::Item as IntoIterator>::IntoIter>,
18}
19
20impl<I> FlattenExact<I>
21where
22    I : ExactSizeIterator + DoubleEndedIterator + TooDeeIterator,
23    I::Item : IntoIterator,
24    <I::Item as IntoIterator>::IntoIter : DoubleEndedIterator + ExactSizeIterator,
25{
26    pub(super) fn new(iter: I) -> FlattenExact<I> {
27        FlattenExact { iter, frontiter: None, backiter: None }
28    }
29}
30
31impl<I> Iterator for FlattenExact<I>
32where
33    I : ExactSizeIterator + DoubleEndedIterator + TooDeeIterator,
34    I::Item : IntoIterator,
35    <I::Item as IntoIterator>::IntoIter : DoubleEndedIterator + ExactSizeIterator,
36{
37    type Item = <I::Item as IntoIterator>::Item;
38
39    #[inline]
40    fn next(&mut self) -> Option<<I::Item as IntoIterator>::Item> {
41        loop {
42            if let Some(ref mut inner) = self.frontiter {
43                if let elt @ Some(_) = inner.next() {
44                    return elt;
45                }
46            }
47            match self.iter.next() {
48                None => return self.backiter.as_mut()?.next(),
49                Some(inner) => self.frontiter = Some(inner.into_iter()),
50            }
51        }
52    }
53
54    #[inline]
55    fn size_hint(&self) -> (usize, Option<usize>) {
56        let mut len = self.num_cols() * self.iter.len();
57        len += self.frontiter.as_ref().map_or(0, |i| i.len());
58        len += self.backiter.as_ref().map_or(0, |i| i.len());
59        (len, Some(len))
60    }
61
62    #[inline]
63    fn last(mut self) -> Option<Self::Item> {
64        self.next_back()
65    }
66
67    #[inline]
68    fn nth(&mut self, mut n: usize) -> Option<<I::Item as IntoIterator>::Item> {
69
70        let num_cols = self.num_cols();
71
72        if num_cols == 0 {
73            return None;
74        }
75
76        if let Some(ref mut inner) = self.frontiter {
77            if n < inner.len() {
78                return inner.nth(n);
79            } else {
80                n -= inner.len();
81                self.frontiter = None;
82            }
83        }
84
85        let iter_skip = self.iter.len().min(n / num_cols);
86        if let Some(inner) = self.iter.nth(iter_skip) {
87            let mut tmp = inner.into_iter();
88            n -= iter_skip * num_cols;
89            debug_assert!(n < tmp.len());
90            let ret_val = tmp.nth(n);
91            self.frontiter = Some(tmp);
92            ret_val
93        } else {
94            n -= iter_skip * num_cols;
95            self.backiter.as_mut()?.nth(n)
96        }
97
98    }
99    
100    #[inline]
101    #[allow(clippy::toplevel_ref_arg)]
102    fn fold<Acc, Fold>(self, init: Acc, ref mut fold: Fold) -> Acc
103    where
104        Fold: FnMut(Acc, Self::Item) -> Acc,
105    {
106        #[inline]
107        fn flatten<U: Iterator, Acc>(
108            fold: &mut impl FnMut(Acc, U::Item) -> Acc,
109        ) -> impl FnMut(Acc, U) -> Acc + '_ {
110            move |acc, iter| iter.fold(acc, &mut *fold)
111        }
112
113        self.frontiter
114            .into_iter()
115            .chain(self.iter.map(IntoIterator::into_iter))
116            .chain(self.backiter)
117            .fold(init, flatten(fold))
118    }
119    
120}
121
122impl<I> DoubleEndedIterator for FlattenExact<I>
123where
124    I : ExactSizeIterator + DoubleEndedIterator + TooDeeIterator,
125    I::Item : IntoIterator,
126    <I::Item as IntoIterator>::IntoIter : DoubleEndedIterator + ExactSizeIterator,
127{
128    #[inline]
129    fn next_back(&mut self) -> Option<<I::Item as IntoIterator>::Item> {
130        loop {
131            if let Some(ref mut inner) = self.backiter {
132                if let elt @ Some(_) = inner.next_back() {
133                    return elt;
134                }
135            }
136            match self.iter.next_back() {
137                None => return self.frontiter.as_mut()?.next_back(),
138                Some(next) => self.backiter = Some(next.into_iter()),
139            }
140        }
141    }
142
143    #[inline]
144    fn nth_back(&mut self, mut n: usize) -> Option<<I::Item as IntoIterator>::Item> {
145        
146        let num_cols = self.num_cols();
147        
148        if num_cols == 0 {
149            return None;
150        }
151        
152        if let Some(ref mut inner) = self.backiter {
153            if n < inner.len() {
154                return inner.nth_back(n);
155            } else {
156                n -= inner.len();
157                self.backiter = None;
158            }
159        }
160        
161        let iter_skip = self.iter.len().min(n / num_cols);
162        if let Some(inner) = self.iter.nth_back(iter_skip) {
163            let mut tmp = inner.into_iter();
164            n -= iter_skip * num_cols;
165            debug_assert!(n < tmp.len());
166            let ret_val = tmp.nth_back(n);
167            self.backiter = Some(tmp);
168            ret_val
169        } else {
170            n -= iter_skip * num_cols;
171            self.frontiter.as_mut()?.nth_back(n)
172        }
173        
174    }
175    
176    #[inline]
177    #[allow(clippy::toplevel_ref_arg)]
178    fn rfold<Acc, Fold>(self, init: Acc, ref mut fold: Fold) -> Acc
179    where
180        Fold: FnMut(Acc, Self::Item) -> Acc,
181    {
182        #[inline]
183        fn flatten<U: DoubleEndedIterator, Acc>(
184            fold: &mut impl FnMut(Acc, U::Item) -> Acc,
185        ) -> impl FnMut(Acc, U) -> Acc + '_ {
186            move |acc, iter| iter.rfold(acc, &mut *fold)
187        }
188
189        self.frontiter
190            .into_iter()
191            .chain(self.iter.map(IntoIterator::into_iter))
192            .chain(self.backiter)
193            .rfold(init, flatten(fold))
194    }
195    
196}
197
198impl<I> ExactSizeIterator for FlattenExact<I>
199where
200    I : ExactSizeIterator + DoubleEndedIterator + TooDeeIterator,
201    I::Item : IntoIterator,
202    <I::Item as IntoIterator>::IntoIter : DoubleEndedIterator + ExactSizeIterator,
203{}
204
205impl<I> TooDeeIterator for FlattenExact<I>
206where
207    I : ExactSizeIterator + DoubleEndedIterator + TooDeeIterator,
208    I::Item : IntoIterator,
209    <I::Item as IntoIterator>::IntoIter : DoubleEndedIterator + ExactSizeIterator,
210{
211    fn num_cols(&self) -> usize {
212        self.iter.num_cols()
213    }
214}