Extract heap helpers into a shared module

src/heap.rs

@@ -0,0 +1,44 @@
+//! Binary-heap helpers on slices, parameterized by a "less than" predicate.
+//!
+//! The heap is a min-heap w.r.t. `less_than`: the root is the "least" element.
+
+/// Make `data` a heap (min-heap w.r.t the sorting).
+pub(crate) fn heapify<T, S>(data: &mut [T], mut less_than: S)
+where
+    S: FnMut(&T, &T) -> bool,
+{
+    for i in (0..data.len() / 2).rev() {
+        sift_down(data, i, &mut less_than);
+    }
+}
+
+/// Sift down element at `index` (`heap` is a min-heap wrt the ordering)
+pub(crate) fn sift_down<T, S>(heap: &mut [T], index: usize, mut less_than: S)
+where
+    S: FnMut(&T, &T) -> bool,
+{
+    debug_assert!(index <= heap.len());
+    let mut pos = index;
+    let mut child = 2 * pos + 1;
+    // Require the right child to be present
+    // This allows to find the index of the smallest child without a branch
+    // that wouldn't be predicted if present
+    while child + 1 < heap.len() {
+        // pick the smaller of the two children
+        // use arithmetic to avoid an unpredictable branch
+        child += usize::from(less_than(&heap[child + 1], &heap[child]));
+
+        // sift down is done if we are already in order
+        if !less_than(&heap[child], &heap[pos]) {
+            return;
+        }
+        heap.swap(pos, child);
+        pos = child;
+        child = 2 * pos + 1;
+    }
+    // Check if the last (left) child was an only child
+    // if it is then it has to be compared with the parent
+    if child + 1 == heap.len() && less_than(&heap[child], &heap[pos]) {
+        heap.swap(pos, child);
+    }
+}

src/k_smallest.rs

@@ -1,3 +1,4 @@
+use crate::heap::{heapify, sift_down};
 use alloc::vec::Vec;
 use core::cmp::Ordering;
 
@@ -7,40 +8,6 @@ where
     I: Iterator,
     F: FnMut(&I::Item, &I::Item) -> Ordering,
 {
-    /// Sift the element currently at `origin` away from the root until it is properly ordered.
-    ///
-    /// This will leave **larger** elements closer to the root of the heap.
-    fn sift_down<T, F>(heap: &mut [T], is_less_than: &mut F, mut origin: usize)
-    where
-        F: FnMut(&T, &T) -> bool,
-    {
-        #[inline]
-        fn children_of(n: usize) -> (usize, usize) {
-            (2 * n + 1, 2 * n + 2)
-        }
-
-        while origin < heap.len() {
-            let (left_idx, right_idx) = children_of(origin);
-            if left_idx >= heap.len() {
-                return;
-            }
-
-            let replacement_idx =
-                if right_idx < heap.len() && is_less_than(&heap[left_idx], &heap[right_idx]) {
-                    right_idx
-                } else {
-                    left_idx
-                };
-
-            if is_less_than(&heap[origin], &heap[replacement_idx]) {
-                heap.swap(origin, replacement_idx);
-                origin = replacement_idx;
-            } else {
-                return;
-            }
-        }
-    }
-
     if k == 0 {
         iter.last();
         return Vec::new();
@@ -51,23 +18,24 @@ where
     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;
+    // The shared heap puts the *least* element (w.r.t. its predicate) on top.
+    // We want the **largest** item on top, so we give it the reversed comparison.
+    let mut is_greater_than = move |a: &_, b: &_| comparator(a, b) == Ordering::Greater;
 
     // Rearrange the storage into a valid heap by reordering from the second-bottom-most layer up to the root.
     // Slightly faster than ordering on each insert, but only by a factor of lg(k).
     // The resulting heap has the **largest** item on top.
-    for i in (0..=(storage.len() / 2)).rev() {
-        sift_down(&mut storage, &mut is_less_than, i);
-    }
+    heapify(&mut storage, &mut is_greater_than);
 
     iter.for_each(|val| {
         debug_assert_eq!(storage.len(), k);
-        if is_less_than(&val, &storage[0]) {
+        // `val < storage[0]`
+        if is_greater_than(&storage[0], &val) {
             // 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);
+            sift_down(&mut storage, 0, &mut is_greater_than);
         }
     });
 
@@ -82,7 +50,7 @@ where
         heap.swap(0, last_idx);
         // Sifting over a truncated slice means that the sifting will not disturb already popped elements.
         heap = &mut heap[..last_idx];
-        sift_down(heap, &mut is_less_than, 0);
+        sift_down(heap, 0, &mut is_greater_than);
     }
 
     storage

src/kmerge_impl.rs

@@ -1,3 +1,4 @@
+use crate::heap::{heapify, sift_down};
 use crate::size_hint;
 
 use alloc::vec::Vec;
@@ -56,47 +57,6 @@ where
     clone_fields!(head, tail);
 }
 
-/// Make `data` a heap (min-heap w.r.t the sorting).
-fn heapify<T, S>(data: &mut [T], mut less_than: S)
-where
-    S: FnMut(&T, &T) -> bool,
-{
-    for i in (0..data.len() / 2).rev() {
-        sift_down(data, i, &mut less_than);
-    }
-}
-
-/// Sift down element at `index` (`heap` is a min-heap wrt the ordering)
-fn sift_down<T, S>(heap: &mut [T], index: usize, mut less_than: S)
-where
-    S: FnMut(&T, &T) -> bool,
-{
-    debug_assert!(index <= heap.len());
-    let mut pos = index;
-    let mut child = 2 * pos + 1;
-    // Require the right child to be present
-    // This allows to find the index of the smallest child without a branch
-    // that wouldn't be predicted if present
-    while child + 1 < heap.len() {
-        // pick the smaller of the two children
-        // use arithmetic to avoid an unpredictable branch
-        child += usize::from(less_than(&heap[child + 1], &heap[child]));
-
-        // sift down is done if we are already in order
-        if !less_than(&heap[child], &heap[pos]) {
-            return;
-        }
-        heap.swap(pos, child);
-        pos = child;
-        child = 2 * pos + 1;
-    }
-    // Check if the last (left) child was an only child
-    // if it is then it has to be compared with the parent
-    if child + 1 == heap.len() && less_than(&heap[child], &heap[pos]) {
-        heap.swap(pos, child);
-    }
-}
-
 /// An iterator adaptor that merges an arbitrary number of base iterators in ascending order.
 /// If all base iterators are sorted (ascending), the result is sorted.
 ///

src/lib.rs

@@ -209,6 +209,8 @@ mod group_map;
 mod groupbylazy;
 #[cfg(feature = "use_std")]
 mod grouping_map;
+#[cfg(feature = "use_alloc")]
+mod heap;
 mod intersperse;
 mod iter_index;
 #[cfg(feature = "use_alloc")]