Skip to content

Commit

Permalink
Override StepBy::{try_fold, try_rfold}
Browse files Browse the repository at this point in the history
  • Loading branch information
timvermeulen committed Sep 4, 2019
1 parent d760df5 commit 78908f2
Show file tree
Hide file tree
Showing 2 changed files with 78 additions and 1 deletion.
44 changes: 43 additions & 1 deletion src/libcore/iter/adapters/mod.rs
Expand Up @@ -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;
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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.
Expand Down
35 changes: 35 additions & 0 deletions src/libcore/tests/iter.rs
Expand Up @@ -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);
Expand All @@ -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() {
Expand Down

0 comments on commit 78908f2

Please sign in to comment.