Override StepBy::{try_fold, try_rfold}

This commit is contained in:
Tim Vermeulen 2019-09-03 11:56:08 +02:00
parent d760df5aea
commit 78908f2e09
2 changed files with 78 additions and 1 deletions

View File

@ -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.

View File

@ -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() {