Override StepBy::{try_fold, try_rfold}
This commit is contained in:
parent
d760df5aea
commit
78908f2e09
@ -5,7 +5,7 @@ use crate::usize;
|
||||
use crate::intrinsics;
|
||||
|
||||
use super::{Iterator, DoubleEndedIterator, ExactSizeIterator, FusedIterator, TrustedLen};
|
||||
use super::LoopState;
|
||||
use super::{LoopState, from_fn};
|
||||
|
||||
mod chain;
|
||||
mod flatten;
|
||||
@ -541,6 +541,26 @@ impl<I> Iterator for StepBy<I> where I: Iterator {
|
||||
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<Ok = 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::from_ok(acc),
|
||||
Some(x) => acc = f(acc, x)?,
|
||||
}
|
||||
}
|
||||
from_fn(nth(&mut self.iter, self.step)).try_fold(acc, f)
|
||||
}
|
||||
}
|
||||
|
||||
impl<I> StepBy<I> where I: ExactSizeIterator {
|
||||
@ -574,6 +594,28 @@ impl<I> DoubleEndedIterator for StepBy<I> where I: DoubleEndedIterator + ExactSi
|
||||
.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<Ok = 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::from_ok(init),
|
||||
Some(x) => {
|
||||
let acc = f(init, x)?;
|
||||
from_fn(nth_back(&mut self.iter, self.step)).try_fold(acc, f)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// StepBy can only make the iterator shorter, so the len will still fit.
|
||||
|
@ -329,6 +329,23 @@ fn test_iterator_step_by_nth_overflow() {
|
||||
assert_eq!(it.0, (usize::MAX as Bigger) * 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_iterator_step_by_nth_try_fold() {
|
||||
let mut it = (0..).step_by(10);
|
||||
assert_eq!(it.try_fold(0, i8::checked_add), None);
|
||||
assert_eq!(it.next(), Some(60));
|
||||
assert_eq!(it.try_fold(0, i8::checked_add), None);
|
||||
assert_eq!(it.next(), Some(90));
|
||||
|
||||
let mut it = (100..).step_by(10);
|
||||
assert_eq!(it.try_fold(50, i8::checked_add), None);
|
||||
assert_eq!(it.next(), Some(110));
|
||||
|
||||
let mut it = (100..=100).step_by(10);
|
||||
assert_eq!(it.next(), Some(100));
|
||||
assert_eq!(it.try_fold(0, i8::checked_add), Some(0));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_iterator_step_by_nth_back() {
|
||||
let mut it = (0..16).step_by(5);
|
||||
@ -354,6 +371,24 @@ fn test_iterator_step_by_nth_back() {
|
||||
assert_eq!(it().nth_back(42), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_iterator_step_by_nth_try_rfold() {
|
||||
let mut it = (0..100).step_by(10);
|
||||
assert_eq!(it.try_rfold(0, i8::checked_add), None);
|
||||
assert_eq!(it.next_back(), Some(70));
|
||||
assert_eq!(it.next(), Some(0));
|
||||
assert_eq!(it.try_rfold(0, i8::checked_add), None);
|
||||
assert_eq!(it.next_back(), Some(30));
|
||||
|
||||
let mut it = (0..100).step_by(10);
|
||||
assert_eq!(it.try_rfold(50, i8::checked_add), None);
|
||||
assert_eq!(it.next_back(), Some(80));
|
||||
|
||||
let mut it = (100..=100).step_by(10);
|
||||
assert_eq!(it.next_back(), Some(100));
|
||||
assert_eq!(it.try_fold(0, i8::checked_add), Some(0));
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[should_panic]
|
||||
fn test_iterator_step_by_zero() {
|
||||
|
Loading…
x
Reference in New Issue
Block a user