Repository: rust-itertools/itertools
Pull request #772 (2023-10-02 10:45:00)

Specialize `WithPosition::fold`
Each call to `f` is done with a `Position` variant known at compile time so I think it might be optimized away in some cases?

    Without specialization: [5.3231 µs 5.3487 µs 5.3806 µs]
    With specialization:    [2.6974 µs 2.7042 µs 2.7132 µs]

We should mention #755 in PRs that specialize `fold` (such as this one) to help keep track of the advancement of the TODO list I'm gonna make there.

Philippe-Cholet: @phimuemue 
I do not update `handled_first` field anymore, as it's consumed.
I avoid to `peek` entirely, the bench become even faster.
And rebased on master.

    [5.4031 µs 5.4195 µs 5.4353 µs] // without specialization
    [563.64 ns 565.92 ns 568.61 ns] // with new specialization
Philippe-Cholet: Just adding that I initially tried to avoid `std::mem::swap` with `(init, head) = self.peekable.fold((init, head), ...);` which worked for me but not for the msrv CI test. It did not like assigning to a tuple. The performance was not affected.
phimuemue: Nicely done, thanks.

As an aside: Does the mail address from your commits work?

Patch:
diff --git a/benches/bench1.rs b/benches/bench1.rs
index 8a384ee..fbcaa05 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 9cdc831..89cddeb 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<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 {}


Test patch:
diff --git a/tests/specializations.rs b/tests/specializations.rs
index f4463a1..4fa6c33 100644
--- a/tests/specializations.rs
+++ b/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);