1 use crate::{intrinsics, iter::from_fn, ops::Try};
2 
3 /// An iterator for stepping iterators by a custom amount.
4 ///
5 /// This `struct` is created by the [`step_by`] method on [`Iterator`]. See
6 /// its documentation for more.
7 ///
8 /// [`step_by`]: Iterator::step_by
9 /// [`Iterator`]: trait.Iterator.html
10 #[must_use = "iterators are lazy and do nothing unless consumed"]
11 #[stable(feature = "iterator_step_by", since = "1.28.0")]
12 #[derive(Clone, Debug)]
13 pub struct StepBy<I> {
14     iter: I,
15     step: usize,
16     first_take: bool,
17 }
18 
19 impl<I> StepBy<I> {
new(iter: I, step: usize) -> StepBy<I>20     pub(in crate::iter) fn new(iter: I, step: usize) -> StepBy<I> {
21         assert!(step != 0);
22         StepBy { iter, step: step - 1, first_take: true }
23     }
24 }
25 
26 #[stable(feature = "iterator_step_by", since = "1.28.0")]
27 impl<I> Iterator for StepBy<I>
28 where
29     I: Iterator,
30 {
31     type Item = I::Item;
32 
33     #[inline]
next(&mut self) -> Option<Self::Item>34     fn next(&mut self) -> Option<Self::Item> {
35         if self.first_take {
36             self.first_take = false;
37             self.iter.next()
38         } else {
39             self.iter.nth(self.step)
40         }
41     }
42 
43     #[inline]
size_hint(&self) -> (usize, Option<usize>)44     fn size_hint(&self) -> (usize, Option<usize>) {
45         #[inline]
46         fn first_size(step: usize) -> impl Fn(usize) -> usize {
47             move |n| if n == 0 { 0 } else { 1 + (n - 1) / (step + 1) }
48         }
49 
50         #[inline]
51         fn other_size(step: usize) -> impl Fn(usize) -> usize {
52             move |n| n / (step + 1)
53         }
54 
55         let (low, high) = self.iter.size_hint();
56 
57         if self.first_take {
58             let f = first_size(self.step);
59             (f(low), high.map(f))
60         } else {
61             let f = other_size(self.step);
62             (f(low), high.map(f))
63         }
64     }
65 
66     #[inline]
nth(&mut self, mut n: usize) -> Option<Self::Item>67     fn nth(&mut self, mut n: usize) -> Option<Self::Item> {
68         if self.first_take {
69             self.first_take = false;
70             let first = self.iter.next();
71             if n == 0 {
72                 return first;
73             }
74             n -= 1;
75         }
76         // n and self.step are indices, we need to add 1 to get the amount of elements
77         // When calling `.nth`, we need to subtract 1 again to convert back to an index
78         // step + 1 can't overflow because `.step_by` sets `self.step` to `step - 1`
79         let mut step = self.step + 1;
80         // n + 1 could overflow
81         // thus, if n is usize::MAX, instead of adding one, we call .nth(step)
82         if n == usize::MAX {
83             self.iter.nth(step - 1);
84         } else {
85             n += 1;
86         }
87 
88         // overflow handling
89         loop {
90             let mul = n.checked_mul(step);
91             {
92                 if intrinsics::likely(mul.is_some()) {
93                     return self.iter.nth(mul.unwrap() - 1);
94                 }
95             }
96             let div_n = usize::MAX / n;
97             let div_step = usize::MAX / step;
98             let nth_n = div_n * n;
99             let nth_step = div_step * step;
100             let nth = if nth_n > nth_step {
101                 step -= div_n;
102                 nth_n
103             } else {
104                 n -= div_step;
105                 nth_step
106             };
107             self.iter.nth(nth - 1);
108         }
109     }
110 
try_fold<Acc, F, R>(&mut self, mut acc: Acc, mut f: F) -> R where F: FnMut(Acc, Self::Item) -> R, R: Try<Output = Acc>,111     fn try_fold<Acc, F, R>(&mut self, mut acc: Acc, mut f: F) -> R
112     where
113         F: FnMut(Acc, Self::Item) -> R,
114         R: Try<Output = Acc>,
115     {
116         #[inline]
117         fn nth<I: Iterator>(iter: &mut I, step: usize) -> impl FnMut() -> Option<I::Item> + '_ {
118             move || iter.nth(step)
119         }
120 
121         if self.first_take {
122             self.first_take = false;
123             match self.iter.next() {
124                 None => return try { acc },
125                 Some(x) => acc = f(acc, x)?,
126             }
127         }
128         from_fn(nth(&mut self.iter, self.step)).try_fold(acc, f)
129     }
130 
fold<Acc, F>(mut self, mut acc: Acc, mut f: F) -> Acc where F: FnMut(Acc, Self::Item) -> Acc,131     fn fold<Acc, F>(mut self, mut acc: Acc, mut f: F) -> Acc
132     where
133         F: FnMut(Acc, Self::Item) -> Acc,
134     {
135         #[inline]
136         fn nth<I: Iterator>(iter: &mut I, step: usize) -> impl FnMut() -> Option<I::Item> + '_ {
137             move || iter.nth(step)
138         }
139 
140         if self.first_take {
141             self.first_take = false;
142             match self.iter.next() {
143                 None => return acc,
144                 Some(x) => acc = f(acc, x),
145             }
146         }
147         from_fn(nth(&mut self.iter, self.step)).fold(acc, f)
148     }
149 }
150 
151 impl<I> StepBy<I>
152 where
153     I: ExactSizeIterator,
154 {
155     // The zero-based index starting from the end of the iterator of the
156     // last element. Used in the `DoubleEndedIterator` implementation.
next_back_index(&self) -> usize157     fn next_back_index(&self) -> usize {
158         let rem = self.iter.len() % (self.step + 1);
159         if self.first_take {
160             if rem == 0 { self.step } else { rem - 1 }
161         } else {
162             rem
163         }
164     }
165 }
166 
167 #[stable(feature = "double_ended_step_by_iterator", since = "1.38.0")]
168 impl<I> DoubleEndedIterator for StepBy<I>
169 where
170     I: DoubleEndedIterator + ExactSizeIterator,
171 {
172     #[inline]
next_back(&mut self) -> Option<Self::Item>173     fn next_back(&mut self) -> Option<Self::Item> {
174         self.iter.nth_back(self.next_back_index())
175     }
176 
177     #[inline]
nth_back(&mut self, n: usize) -> Option<Self::Item>178     fn nth_back(&mut self, n: usize) -> Option<Self::Item> {
179         // `self.iter.nth_back(usize::MAX)` does the right thing here when `n`
180         // is out of bounds because the length of `self.iter` does not exceed
181         // `usize::MAX` (because `I: ExactSizeIterator`) and `nth_back` is
182         // zero-indexed
183         let n = n.saturating_mul(self.step + 1).saturating_add(self.next_back_index());
184         self.iter.nth_back(n)
185     }
186 
try_rfold<Acc, F, R>(&mut self, init: Acc, mut f: F) -> R where F: FnMut(Acc, Self::Item) -> R, R: Try<Output = Acc>,187     fn try_rfold<Acc, F, R>(&mut self, init: Acc, mut f: F) -> R
188     where
189         F: FnMut(Acc, Self::Item) -> R,
190         R: Try<Output = Acc>,
191     {
192         #[inline]
193         fn nth_back<I: DoubleEndedIterator>(
194             iter: &mut I,
195             step: usize,
196         ) -> impl FnMut() -> Option<I::Item> + '_ {
197             move || iter.nth_back(step)
198         }
199 
200         match self.next_back() {
201             None => try { init },
202             Some(x) => {
203                 let acc = f(init, x)?;
204                 from_fn(nth_back(&mut self.iter, self.step)).try_fold(acc, f)
205             }
206         }
207     }
208 
209     #[inline]
rfold<Acc, F>(mut self, init: Acc, mut f: F) -> Acc where Self: Sized, F: FnMut(Acc, Self::Item) -> Acc,210     fn rfold<Acc, F>(mut self, init: Acc, mut f: F) -> Acc
211     where
212         Self: Sized,
213         F: FnMut(Acc, Self::Item) -> Acc,
214     {
215         #[inline]
216         fn nth_back<I: DoubleEndedIterator>(
217             iter: &mut I,
218             step: usize,
219         ) -> impl FnMut() -> Option<I::Item> + '_ {
220             move || iter.nth_back(step)
221         }
222 
223         match self.next_back() {
224             None => init,
225             Some(x) => {
226                 let acc = f(init, x);
227                 from_fn(nth_back(&mut self.iter, self.step)).fold(acc, f)
228             }
229         }
230     }
231 }
232 
233 // StepBy can only make the iterator shorter, so the len will still fit.
234 #[stable(feature = "iterator_step_by", since = "1.28.0")]
235 impl<I> ExactSizeIterator for StepBy<I> where I: ExactSizeIterator {}
236