Specialize WithPosition::fold

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);

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<I: Iterator> Iterator for WithPosition<I> {
     fn size_hint(&self) -> (usize, Option<usize>) {
         self.peekable.size_hint()
     }
+
+    fn fold<B, F>(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<I> ExactSizeIterator for WithPosition<I> where I: ExactSizeIterator {}

tests/specializations.rs

@@ -74,6 +74,10 @@ where
 }
 
 quickcheck! {
+    fn with_position(v: Vec<u8>) -> () {
+        test_specializations(&v.iter().with_position());
+    }
+
     fn tuple_combinations(v: Vec<u8>) -> () {
         let mut v = v;
         v.truncate(10);