diff options
Diffstat (limited to 'library/core/src/iter/adapters/step_by.rs')
| -rw-r--r-- | library/core/src/iter/adapters/step_by.rs | 235 |
1 files changed, 235 insertions, 0 deletions
diff --git a/library/core/src/iter/adapters/step_by.rs b/library/core/src/iter/adapters/step_by.rs new file mode 100644 index 000000000..4252c34a0 --- /dev/null +++ b/library/core/src/iter/adapters/step_by.rs @@ -0,0 +1,235 @@ +use crate::{intrinsics, iter::from_fn, ops::Try}; + +/// An iterator for stepping iterators by a custom amount. +/// +/// This `struct` is created by the [`step_by`] method on [`Iterator`]. See +/// its documentation for more. +/// +/// [`step_by`]: Iterator::step_by +/// [`Iterator`]: trait.Iterator.html +#[must_use = "iterators are lazy and do nothing unless consumed"] +#[stable(feature = "iterator_step_by", since = "1.28.0")] +#[derive(Clone, Debug)] +pub struct StepBy<I> { + iter: I, + step: usize, + first_take: bool, +} + +impl<I> StepBy<I> { + pub(in crate::iter) fn new(iter: I, step: usize) -> StepBy<I> { + assert!(step != 0); + StepBy { iter, step: step - 1, first_take: true } + } +} + +#[stable(feature = "iterator_step_by", since = "1.28.0")] +impl<I> Iterator for StepBy<I> +where + I: Iterator, +{ + type Item = I::Item; + + #[inline] + fn next(&mut self) -> Option<Self::Item> { + if self.first_take { + self.first_take = false; + self.iter.next() + } else { + self.iter.nth(self.step) + } + } + + #[inline] + fn size_hint(&self) -> (usize, Option<usize>) { + #[inline] + fn first_size(step: usize) -> impl Fn(usize) -> usize { + move |n| if n == 0 { 0 } else { 1 + (n - 1) / (step + 1) } + } + + #[inline] + fn other_size(step: usize) -> impl Fn(usize) -> usize { + move |n| n / (step + 1) + } + + let (low, high) = self.iter.size_hint(); + + if self.first_take { + let f = first_size(self.step); + (f(low), high.map(f)) + } else { + let f = other_size(self.step); + (f(low), high.map(f)) + } + } + + #[inline] + fn nth(&mut self, mut n: usize) -> Option<Self::Item> { + if self.first_take { + self.first_take = false; + let first = self.iter.next(); + if n == 0 { + return first; + } + n -= 1; + } + // n and self.step are indices, we need to add 1 to get the amount of elements + // When calling `.nth`, we need to subtract 1 again to convert back to an index + // step + 1 can't overflow because `.step_by` sets `self.step` to `step - 1` + let mut step = self.step + 1; + // n + 1 could overflow + // thus, if n is usize::MAX, instead of adding one, we call .nth(step) + if n == usize::MAX { + self.iter.nth(step - 1); + } else { + n += 1; + } + + // overflow handling + loop { + let mul = n.checked_mul(step); + { + if intrinsics::likely(mul.is_some()) { + return self.iter.nth(mul.unwrap() - 1); + } + } + let div_n = usize::MAX / n; + let div_step = usize::MAX / step; + let nth_n = div_n * n; + let nth_step = div_step * step; + let nth = if nth_n > nth_step { + step -= div_n; + nth_n + } else { + n -= div_step; + nth_step + }; + self.iter.nth(nth - 1); + } + } + + fn 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>, + { + #[inline] + fn nth<I: Iterator>(iter: &mut I, step: usize) -> impl FnMut() -> Option<I::Item> + '_ { + move || iter.nth(step) + } + + if self.first_take { + self.first_take = false; + match self.iter.next() { + None => return try { acc }, + Some(x) => acc = f(acc, x)?, + } + } + from_fn(nth(&mut self.iter, self.step)).try_fold(acc, f) + } + + fn fold<Acc, F>(mut self, mut acc: Acc, mut f: F) -> Acc + where + F: FnMut(Acc, Self::Item) -> Acc, + { + #[inline] + fn nth<I: Iterator>(iter: &mut I, step: usize) -> impl FnMut() -> Option<I::Item> + '_ { + move || iter.nth(step) + } + + if self.first_take { + self.first_take = false; + match self.iter.next() { + None => return acc, + Some(x) => acc = f(acc, x), + } + } + from_fn(nth(&mut self.iter, self.step)).fold(acc, f) + } +} + +impl<I> StepBy<I> +where + I: ExactSizeIterator, +{ + // The zero-based index starting from the end of the iterator of the + // last element. Used in the `DoubleEndedIterator` implementation. + fn next_back_index(&self) -> usize { + let rem = self.iter.len() % (self.step + 1); + if self.first_take { + if rem == 0 { self.step } else { rem - 1 } + } else { + rem + } + } +} + +#[stable(feature = "double_ended_step_by_iterator", since = "1.38.0")] +impl<I> DoubleEndedIterator for StepBy<I> +where + I: DoubleEndedIterator + ExactSizeIterator, +{ + #[inline] + fn next_back(&mut self) -> Option<Self::Item> { + self.iter.nth_back(self.next_back_index()) + } + + #[inline] + fn nth_back(&mut self, n: usize) -> Option<Self::Item> { + // `self.iter.nth_back(usize::MAX)` does the right thing here when `n` + // is out of bounds because the length of `self.iter` does not exceed + // `usize::MAX` (because `I: ExactSizeIterator`) and `nth_back` is + // zero-indexed + let n = n.saturating_mul(self.step + 1).saturating_add(self.next_back_index()); + self.iter.nth_back(n) + } + + fn try_rfold<Acc, F, R>(&mut self, init: Acc, mut f: F) -> R + where + F: FnMut(Acc, Self::Item) -> R, + R: Try<Output = Acc>, + { + #[inline] + fn nth_back<I: DoubleEndedIterator>( + iter: &mut I, + step: usize, + ) -> impl FnMut() -> Option<I::Item> + '_ { + move || iter.nth_back(step) + } + + match self.next_back() { + None => try { init }, + Some(x) => { + let acc = f(init, x)?; + from_fn(nth_back(&mut self.iter, self.step)).try_fold(acc, f) + } + } + } + + #[inline] + fn rfold<Acc, F>(mut self, init: Acc, mut f: F) -> Acc + where + Self: Sized, + F: FnMut(Acc, Self::Item) -> Acc, + { + #[inline] + fn nth_back<I: DoubleEndedIterator>( + iter: &mut I, + step: usize, + ) -> impl FnMut() -> Option<I::Item> + '_ { + move || iter.nth_back(step) + } + + match self.next_back() { + None => init, + Some(x) => { + let acc = f(init, x); + from_fn(nth_back(&mut self.iter, self.step)).fold(acc, f) + } + } + } +} + +// StepBy can only make the iterator shorter, so the len will still fit. +#[stable(feature = "iterator_step_by", since = "1.28.0")] +impl<I> ExactSizeIterator for StepBy<I> where I: ExactSizeIterator {} |