Merge #518

518: Remove unpredicted branch from kmerge::sift_down r=jswrenn a=SkiFire13

This is pretty much a port from rust-lang/rust#78857

Compared with the previous implementation, this adds more branches for bound checks which aren't present on the stdlib version, however they should be predicted almost always. The speedup should come from the removal of an unpredictable branch from the loop body, in favor of boolean arithmetic.

The benchmarks seem to agree:

```
before:

test kmerge default ... bench:        6812 ns/iter (+/- 18)
test kmerge tenway ... bench:      223673 ns/iter (+/- 769)

after:

test kmerge default ... bench:        6212 ns/iter (+/- 43)
test kmerge tenway ... bench:      190700 ns/iter (+/- 419)
```

Co-authored-by: Giacomo Stevanato <[email protected]>

Author: bors[bot]

src/kmerge_impl.rs

@@ -74,14 +74,13 @@ fn sift_down<T, S>(heap: &mut [T], index: usize, mut less_than: S)
     debug_assert!(index <= heap.len());
     let mut pos = index;
     let mut child = 2 * pos + 1;
-    // the `pos` conditional is to avoid a bounds check
-    while pos < heap.len() && child < heap.len() {
-        let right = child + 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
-        if right < heap.len() && less_than(&heap[right], &heap[child]) {
-            child = right;
-        }
+        // use aritmethic to avoid an unpredictable branch
+        child += less_than(&heap[child+1], &heap[child]) as usize;
 
         // sift down is done if we are already in order
         if !less_than(&heap[child], &heap[pos]) {
@@ -91,6 +90,11 @@ fn sift_down<T, S>(heap: &mut [T], index: usize, mut less_than: S)
         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 abitrary number of base iterators in ascending order.