Unbundling heap type and adjusting sift_down

The type was not providing much benefit after inlining the other functions

Author: ejmount

src/k_smallest.rs

@@ -5,49 +5,44 @@ use core::cmp::{Ord, Ordering, Reverse};
 fn k_smallest_general<T, I: Iterator<Item = T>>(
     mut iter: I,
     k: usize,
-    comparator: impl Fn(&T, &T) -> Ordering,
+    mut comparator: impl FnMut(&T, &T) -> Ordering,
 ) -> IntoIter<T> {
     if k == 0 {
         return Vec::new().into_iter();
     }
     let mut storage: Vec<T> = iter.by_ref().take(k).collect();
-
-    let mut heap = MaxHeap {
-        len: storage.len(),
-        storage: &mut storage[..],
-        comparator,
-    };
+    let mut heap = &mut storage[..];
 
     // Rearrange the into a valid heap by reordering from the second-bottom-most layer up
     // Slightly faster than ordering on each insert
     // (But only by a factor of lg(k) and I'd love to hear of a use case where that matters)
-    for i in (0..(heap.len / 2 + 1)).rev() {
-        heap.sift_down(i);
+    for i in (0..(heap.len() / 2 + 1)).rev() {
+        sift_down(heap, &mut comparator, i);
     }
 
-    if k == heap.storage.len() {
+    if k == heap.len() {
         // Nothing else needs done if we didn't fill the storage in the first place
         // Also avoids unexpected behaviour with restartable iterators
         for val in iter {
-            if (heap.comparator)(&heap.storage[0], &val) == Ordering::Greater {
-                heap.storage[0] = val;
-                heap.sift_down(0);
+            if comparator(&heap[0], &val) == Ordering::Greater {
+                heap[0] = val;
+                sift_down(heap, &mut comparator, 0);
             }
         }
     }
 
-    while heap.len > 1 {
-        heap.storage.swap(0, heap.len - 1);
-        // Leaves the length shorter than the number of initialized elements
+    while heap.len() > 1 {
+        let last_idx = heap.len() - 1;
+        heap.swap(0, last_idx);
+        // Leaves the length shorter than the number of elements
         // so that sifting does not disturb already popped elements
-        heap.len -= 1;
-        heap.sift_down(0);
+        heap = &mut heap[..last_idx];
+        sift_down(heap, &mut comparator, 0);
     }
 
     storage.into_iter()
 }
 
-
 pub(crate) fn reverse_cmp<T, F>(cmp: F) -> impl Fn(&T, &T) -> Ordering
 where
     F: Fn(&T, &T) -> Ordering,
@@ -92,59 +87,36 @@ where
     results.into_iter()
 }
 
-/// An efficient heapsort requires that the heap ordering is the inverse of the desired sort order
-/// So the basic case of retrieving minimum elements requires a max heap
-///
-/// This type does not attempt to reproduce all the functionality of [std::collections::BinaryHeap] and instead only implements what is needed for iter operations,
-/// e.g. we do not need to insert single elements.
-/// Additionally, some minor optimizations used in the std BinaryHeap are not used here, e.g. elements are actually swapped rather than managing a "hole"
-struct MaxHeap<'a, T, C> {
-    // It may be not be possible to shrink the storage for smaller sequencess
-    // so manually manage the initialization
-    // This is **assumed not to be empty**
-    storage: &'a mut [T],
-    comparator: C,
-    // this is always less or equal to the count of actual elements
-    // allowing it to be less means the heap property can cover only a subset of the vec
-    // while reusing the storage
-    len: usize,
-}
+/// Sift the element currently at `origin` **away** from the root until it is properly ordered
+fn sift_down<T>(
+    heap: &mut [T],
+    comparator: &mut impl FnMut(&T, &T) -> Ordering,
+    mut origin: usize,
+) {
+    fn children_of(n: usize) -> (usize, usize) {
+        (2 * n + 1, 2 * n + 2)
+    }
 
-impl<'a, T, C> MaxHeap<'a, T, C>
-where
-    C: FnMut(&T, &T) -> Ordering,
-{
-    /// Sift the element currently at `origin` **away** from the root until it is properly ordered
-    fn sift_down(&mut self, mut origin: usize) {
-        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;
         }
 
-        while origin < self.len {
-            let (left_idx, right_idx) = children_of(origin);
-
-            if left_idx >= self.len {
-                // the left is the earlier child, so if it doesn't exist there's nothing to swap with
-                return;
-            }
-
-            let items = &self.storage;
-
-            let replacement_idx = if right_idx < self.len
-                && Ordering::Less == (self.comparator)(&items[left_idx], &items[right_idx])
-            {
-                right_idx
-            } else {
-                left_idx
-            };
-
-            let cmp = (self.comparator)(&items[origin], &items[replacement_idx]);
-            if Ordering::Less == cmp {
-                self.storage.swap(origin, replacement_idx);
-                origin = replacement_idx;
-            } else {
-                return;
-            }
+        let replacement_idx = if right_idx < heap.len()
+            && Ordering::Less == comparator(&heap[left_idx], &heap[right_idx])
+        {
+            right_idx
+        } else {
+            left_idx
+        };
+
+        let cmp = comparator(&heap[origin], &heap[replacement_idx]);
+        if Ordering::Less == cmp {
+            heap.swap(origin, replacement_idx);
+            origin = replacement_idx;
+        } else {
+            return;
         }
     }
 }