1 use crate::{intrinsics, iter::from_fn, ops::Try};
3 /// An iterator for stepping iterators by a custom amount.
5 /// This `struct` is created by the [`step_by`] method on [`Iterator`]. See
6 /// its documentation for more.
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> {
20 pub(in crate::iter) fn new(iter: I, step: usize) -> StepBy<I> {
22 StepBy { iter, step: step - 1, first_take: true }
26 #[stable(feature = "iterator_step_by", since = "1.28.0")]
27 impl<I> Iterator for StepBy<I>
34 fn next(&mut self) -> Option<Self::Item> {
36 self.first_take = false;
39 self.iter.nth(self.step)
44 fn size_hint(&self) -> (usize, Option<usize>) {
46 fn first_size(step: usize) -> impl Fn(usize) -> usize {
47 move |n| if n == 0 { 0 } else { 1 + (n - 1) / (step + 1) }
51 fn other_size(step: usize) -> impl Fn(usize) -> usize {
52 move |n| n / (step + 1)
55 let (low, high) = self.iter.size_hint();
58 let f = first_size(self.step);
61 let f = other_size(self.step);
67 fn nth(&mut self, mut n: usize) -> Option<Self::Item> {
69 self.first_take = false;
70 let first = self.iter.next();
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)
83 self.iter.nth(step - 1);
90 let mul = n.checked_mul(step);
92 if intrinsics::likely(mul.is_some()) {
93 return self.iter.nth(mul.unwrap() - 1);
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 {
107 self.iter.nth(nth - 1);
111 fn try_fold<Acc, F, R>(&mut self, mut acc: Acc, mut f: F) -> R
113 F: FnMut(Acc, Self::Item) -> R,
114 R: Try<Output = Acc>,
117 fn nth<I: Iterator>(iter: &mut I, step: usize) -> impl FnMut() -> Option<I::Item> + '_ {
118 move || iter.nth(step)
122 self.first_take = false;
123 match self.iter.next() {
124 None => return try { acc },
125 Some(x) => acc = f(acc, x)?,
128 from_fn(nth(&mut self.iter, self.step)).try_fold(acc, f)
131 fn fold<Acc, F>(mut self, mut acc: Acc, mut f: F) -> Acc
133 F: FnMut(Acc, Self::Item) -> Acc,
136 fn nth<I: Iterator>(iter: &mut I, step: usize) -> impl FnMut() -> Option<I::Item> + '_ {
137 move || iter.nth(step)
141 self.first_take = false;
142 match self.iter.next() {
144 Some(x) => acc = f(acc, x),
147 from_fn(nth(&mut self.iter, self.step)).fold(acc, f)
153 I: ExactSizeIterator,
155 // The zero-based index starting from the end of the iterator of the
156 // last element. Used in the `DoubleEndedIterator` implementation.
157 fn next_back_index(&self) -> usize {
158 let rem = self.iter.len() % (self.step + 1);
160 if rem == 0 { self.step } else { rem - 1 }
167 #[stable(feature = "double_ended_step_by_iterator", since = "1.38.0")]
168 impl<I> DoubleEndedIterator for StepBy<I>
170 I: DoubleEndedIterator + ExactSizeIterator,
173 fn next_back(&mut self) -> Option<Self::Item> {
174 self.iter.nth_back(self.next_back_index())
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
183 let n = n.saturating_mul(self.step + 1).saturating_add(self.next_back_index());
184 self.iter.nth_back(n)
187 fn try_rfold<Acc, F, R>(&mut self, init: Acc, mut f: F) -> R
189 F: FnMut(Acc, Self::Item) -> R,
190 R: Try<Output = Acc>,
193 fn nth_back<I: DoubleEndedIterator>(
196 ) -> impl FnMut() -> Option<I::Item> + '_ {
197 move || iter.nth_back(step)
200 match self.next_back() {
201 None => try { init },
203 let acc = f(init, x)?;
204 from_fn(nth_back(&mut self.iter, self.step)).try_fold(acc, f)
210 fn rfold<Acc, F>(mut self, init: Acc, mut f: F) -> Acc
213 F: FnMut(Acc, Self::Item) -> Acc,
216 fn nth_back<I: DoubleEndedIterator>(
219 ) -> impl FnMut() -> Option<I::Item> + '_ {
220 move || iter.nth_back(step)
223 match self.next_back() {
226 let acc = f(init, x);
227 from_fn(nth_back(&mut self.iter, self.step)).fold(acc, f)
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 {}