Specialize KMergeBy fold for the final input iterator

benches/fold_specialization.rs

@@ -1,6 +1,6 @@
 #![allow(unstable_name_collisions)]
 
-use criterion::{criterion_group, criterion_main, Criterion};
+use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion};
 use itertools::Itertools;
 
 struct Unspecialized<I>(I);
@@ -66,10 +66,63 @@ mod specialization {
     }
 }
 
+fn kmerge_fold(c: &mut Criterion) {
+    let cases = [
+        ("single", vec![(0_u64..4096).collect_vec()]),
+        (
+            "skewed",
+            vec![
+                (0_u64..32).collect_vec(),
+                (32_u64..64).collect_vec(),
+                (64_u64..4096).collect_vec(),
+            ],
+        ),
+        (
+            "balanced",
+            (0_u64..4)
+                .map(|offset| (offset..4096).step_by(4).collect_vec())
+                .collect_vec(),
+        ),
+    ];
+    let mut group = c.benchmark_group("kmerge_fold");
+    for (name, sources) in cases {
+        group.bench_with_input(
+            BenchmarkId::new("specialized", name),
+            &sources,
+            |b, sources| {
+                b.iter(|| {
+                    black_box(sources)
+                        .iter()
+                        .map(|items| items.iter().copied())
+                        .kmerge()
+                        .fold(0_u64, u64::wrapping_add)
+                });
+            },
+        );
+        group.bench_with_input(
+            BenchmarkId::new("unspecialized", name),
+            &sources,
+            |b, sources| {
+                b.iter(|| {
+                    Unspecialized(
+                        black_box(sources)
+                            .iter()
+                            .map(|items| items.iter().copied())
+                            .kmerge(),
+                    )
+                    .fold(0_u64, u64::wrapping_add)
+                });
+            },
+        );
+    }
+    group.finish();
+}
+
 criterion_group!(
     benches,
     specialization::intersperse::external,
     specialization::intersperse::internal_specialized,
     specialization::intersperse::internal_unspecialized,
+    kmerge_fold,
 );
 criterion_main!(benches);

src/kmerge_impl.rs

@@ -222,6 +222,25 @@ where
         Some(result)
     }
 
+    fn fold<B, G>(mut self, init: B, mut f: G) -> B
+    where
+        Self: Sized,
+        G: FnMut(B, Self::Item) -> B,
+    {
+        let mut acc = init;
+        while self.heap.len() > 1 {
+            acc = f(acc, self.next().unwrap());
+        }
+        match self.heap.pop() {
+            None => acc,
+            Some(HeadTail { head, tail }) => {
+                // As in `next`, advance the tail before consuming its previous head.
+                let (acc, last) = tail.fold((acc, head), |(acc, head), next| (f(acc, head), next));
+                f(acc, last)
+            }
+        }
+    }
+
     fn size_hint(&self) -> (usize, Option<usize>) {
         self.heap
             .iter()

tests/specializations.rs

@@ -495,6 +495,95 @@ quickcheck! {
     }
 }
 
+#[test]
+fn kmerge_fold_preserves_iterator_events() {
+    use std::cell::RefCell;
+    use std::rc::Rc;
+
+    #[derive(Clone, Debug, PartialEq)]
+    enum Event {
+        Next(usize, Option<i32>),
+        Compare(i32, i32),
+        Consume(i32),
+        Drop(usize),
+    }
+
+    struct Tracked {
+        id: usize,
+        items: std::vec::IntoIter<Option<i32>>,
+        events: Rc<RefCell<Vec<Event>>>,
+    }
+
+    impl Iterator for Tracked {
+        type Item = i32;
+
+        fn next(&mut self) -> Option<i32> {
+            let item = self.items.next().flatten();
+            self.events.borrow_mut().push(Event::Next(self.id, item));
+            item
+        }
+    }
+
+    impl Drop for Tracked {
+        fn drop(&mut self) {
+            self.events.borrow_mut().push(Event::Drop(self.id));
+        }
+    }
+
+    fn run(
+        sources: &[Vec<Option<i32>>],
+        advanced: usize,
+        specialized: bool,
+    ) -> (Vec<i32>, Vec<Event>) {
+        let events = Rc::new(RefCell::new(Vec::new()));
+        let mut it = sources
+            .iter()
+            .enumerate()
+            .map(|(id, items)| Tracked {
+                id,
+                items: items.clone().into_iter(),
+                events: events.clone(),
+            })
+            .kmerge_by(|&a, &b| {
+                events.borrow_mut().push(Event::Compare(a, b));
+                a < b
+            });
+        for _ in 0..advanced {
+            it.next();
+        }
+        let consume = |mut acc: Vec<i32>, item| {
+            events.borrow_mut().push(Event::Consume(item));
+            acc.push(item);
+            acc
+        };
+        let items = if specialized {
+            it.fold(Vec::new(), consume)
+        } else {
+            Unspecialized(it).fold(Vec::new(), consume)
+        };
+        let trace = events.borrow().clone();
+        (items, trace)
+    }
+
+    let cases = [
+        vec![],
+        vec![vec![]],
+        vec![vec![Some(1)]],
+        vec![vec![Some(1), Some(2), None, Some(99)]],
+        vec![
+            vec![Some(0), Some(3), None, Some(98)],
+            vec![Some(1)],
+            vec![Some(2), Some(4), Some(5)],
+        ],
+        vec![vec![None, Some(99)], vec![Some(1), Some(2)]],
+    ];
+    for sources in &cases {
+        for advanced in 0..=3 {
+            assert_eq!(run(sources, advanced, true), run(sources, advanced, false));
+        }
+    }
+}
+
 quickcheck! {
     // TODO Replace this function by a normal call to test_specializations
     fn process_results(v: Vec<Result<u8, u8>>) -> () {