From 76be90719aabba98e4027ea62d77821d7e3634cc Mon Sep 17 00:00:00 2001
From: Philippe-Cholet <44676486+Philippe-Cholet@users.noreply.github.com>
Date: Thu, 14 Mar 2024 14:08:36 +0100
Subject: [PATCH 1/2] Fuse the iterator in `k_smallest`

---
 src/lib.rs | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index 0506a3d84..925cd9c6e 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -2949,7 +2949,7 @@ pub trait Itertools: Iterator {
     /// itertools::assert_equal(five_smallest, 0..5);
     /// ```
     #[cfg(feature = "use_alloc")]
-    fn k_smallest(mut self, k: usize) -> VecIntoIter<Self::Item>
+    fn k_smallest(self, k: usize) -> VecIntoIter<Self::Item>
     where
         Self: Sized,
         Self::Item: Ord,
@@ -2963,9 +2963,10 @@ pub trait Itertools: Iterator {
             return Vec::new().into_iter();
         }
 
-        let mut heap = self.by_ref().take(k).collect::<BinaryHeap<_>>();
+        let mut iter = self.fuse();
+        let mut heap: BinaryHeap<_> = iter.by_ref().take(k).collect();
 
-        self.for_each(|i| {
+        iter.for_each(|i| {
             debug_assert_eq!(heap.len(), k);
             // Equivalent to heap.push(min(i, heap.pop())) but more efficient.
             // This should be done with a single `.peek_mut().unwrap()` but

From 0936c8dd76ee18d3b3409e19590bd31c111f739d Mon Sep 17 00:00:00 2001
From: Philippe-Cholet <44676486+Philippe-Cholet@users.noreply.github.com>
Date: Thu, 14 Mar 2024 14:09:10 +0100
Subject: [PATCH 2/2] Fuse the iterator in `k_smallest_general`

---
 src/k_smallest.rs | 25 +++++++++++--------------
 1 file changed, 11 insertions(+), 14 deletions(-)

diff --git a/src/k_smallest.rs b/src/k_smallest.rs
index 4de1b25e0..fe699fbd4 100644
--- a/src/k_smallest.rs
+++ b/src/k_smallest.rs
@@ -2,7 +2,7 @@ use alloc::vec::Vec;
 use core::cmp::Ordering;
 
 /// Consumes a given iterator, returning the minimum elements in **ascending** order.
-pub(crate) fn k_smallest_general<I, F>(mut iter: I, k: usize, mut comparator: F) -> Vec<I::Item>
+pub(crate) fn k_smallest_general<I, F>(iter: I, k: usize, mut comparator: F) -> Vec<I::Item>
 where
     I: Iterator,
     F: FnMut(&I::Item, &I::Item) -> Ordering,
@@ -44,6 +44,7 @@ where
     if k == 0 {
         return Vec::new();
     }
+    let mut iter = iter.fuse();
     let mut storage: Vec<I::Item> = iter.by_ref().take(k).collect();
 
     let mut is_less_than = move |a: &_, b: &_| comparator(a, b) == Ordering::Less;
@@ -55,19 +56,15 @@ where
         sift_down(&mut storage, &mut is_less_than, i);
     }
 
-    if k == storage.len() {
-        // If we fill the storage, there may still be iterator elements left so feed them into the heap.
-        // Also avoids unexpected behaviour with restartable iterators.
-        iter.for_each(|val| {
-            if is_less_than(&val, &storage[0]) {
-                // Treating this as an push-and-pop saves having to write a sift-up implementation.
-                // https://en.wikipedia.org/wiki/Binary_heap#Insert_then_extract
-                storage[0] = val;
-                // We retain the smallest items we've seen so far, but ordered largest first so we can drop the largest efficiently.
-                sift_down(&mut storage, &mut is_less_than, 0);
-            }
-        });
-    }
+    iter.for_each(|val| {
+        if is_less_than(&val, &storage[0]) {
+            // Treating this as an push-and-pop saves having to write a sift-up implementation.
+            // https://en.wikipedia.org/wiki/Binary_heap#Insert_then_extract
+            storage[0] = val;
+            // We retain the smallest items we've seen so far, but ordered largest first so we can drop the largest efficiently.
+            sift_down(&mut storage, &mut is_less_than, 0);
+        }
+    });
 
     // Ultimately the items need to be in least-first, strict order, but the heap is currently largest-first.
     // To achieve this, repeatedly,