diff --git a/benches/bench1.rs b/benches/bench1.rs index 8a384ee69..fbcaa05a7 100644 --- a/benches/bench1.rs +++ b/benches/bench1.rs @@ -1,6 +1,7 @@ use criterion::{black_box, criterion_group, criterion_main, Criterion}; use itertools::free::cloned; use itertools::Itertools; +use itertools::Position; use itertools::{iproduct, EitherOrBoth}; use std::cmp; @@ -819,6 +820,22 @@ fn permutations_slice(c: &mut Criterion) { }); } +fn with_position_fold(c: &mut Criterion) { + let v = black_box((0..10240).collect_vec()); + c.bench_function("with_position fold", move |b| { + b.iter(|| { + v.iter() + .with_position() + .fold(0, |acc, (pos, &item)| match pos { + Position::Middle => acc + item, + Position::First => acc - 2 * item, + Position::Last => acc + 2 * item, + Position::Only => acc + 5 * item, + }) + }) + }); +} + criterion_group!( benches, slice_iter, @@ -866,5 +883,6 @@ criterion_group!( permutations_iter, permutations_range, permutations_slice, + with_position_fold, ); criterion_main!(benches); diff --git a/src/with_position.rs b/src/with_position.rs index 9cdc8319e..89cddeb8a 100644 --- a/src/with_position.rs +++ b/src/with_position.rs @@ -37,7 +37,7 @@ where /// Indicates the position of this element in the iterator results. /// /// See [`.with_position()`](crate::Itertools::with_position) for more information. -#[derive(Copy, Clone, Debug, PartialEq)] +#[derive(Copy, Clone, Debug, PartialEq, Eq)] pub enum Position { /// This is the first element. First, @@ -81,6 +81,33 @@ impl Iterator for WithPosition { fn size_hint(&self) -> (usize, Option) { self.peekable.size_hint() } + + fn fold(mut self, mut init: B, mut f: F) -> B + where + F: FnMut(B, Self::Item) -> B, + { + if let Some(mut head) = self.peekable.next() { + if !self.handled_first { + // The current head is `First` or `Only`, + // it depends if there is another item or not. + match self.peekable.next() { + Some(second) => { + let first = std::mem::replace(&mut head, second); + init = f(init, (Position::First, first)); + } + None => return f(init, (Position::Only, head)), + } + } + // Have seen the first item, and there's something left. + init = self.peekable.fold(init, |acc, mut item| { + std::mem::swap(&mut head, &mut item); + f(acc, (Position::Middle, item)) + }); + // The "head" is now the last item. + init = f(init, (Position::Last, head)); + } + init + } } impl ExactSizeIterator for WithPosition where I: ExactSizeIterator {} diff --git a/tests/specializations.rs b/tests/specializations.rs index f4463a1b2..4fa6c33df 100644 --- a/tests/specializations.rs +++ b/tests/specializations.rs @@ -74,6 +74,10 @@ where } quickcheck! { + fn with_position(v: Vec) -> () { + test_specializations(&v.iter().with_position()); + } + fn tuple_combinations(v: Vec) -> () { let mut v = v; v.truncate(10);