Skip to content

Commit 4088981

Browse files
authored
Rollup merge of #78857 - SkiFire13:bheap-opt, r=KodrAus
Improve BinaryHeap performance By changing the condition in the loops from `child < end` to `child < end - 1` we're guaranteed that `right = child + 1 < end` and since finding the index of the biggest sibling can be done with an arithmetic operation we can remove a branch from the loop body. The case where there's no right child, i.e. `child == end - 1` is instead handled outside the loop, after it ends; note that if the loops ends early we can use `return` instead of `break` since the check `child == end - 1` will surely fail. I've also removed a call to `<[T]>::swap` that was hiding a bound check that [wasn't being optimized by LLVM](https://godbolt.org/z/zrhdGM). A quick benchmarks on my pc shows that the gains are pretty significant: |name |before ns/iter |after ns/iter |diff ns/iter |diff % |speedup | |---------------------|----------------|---------------|--------------|----------|--------| |find_smallest_1000 | 352,565 | 260,098 | -92,467 | -26.23% | x 1.36 | |from_vec | 676,795 | 473,934 | -202,861 | -29.97% | x 1.43 | |into_sorted_vec | 469,511 | 304,275 | -165,236 | -35.19% | x 1.54 | |pop | 483,198 | 373,778 | -109,420 | -22.64% | x 1.29 | The other 2 benchmarks for `BinaryHeap` (`peek_mut_deref_mut` and `push`) weren't impacted and as such didn't show any significant change.
2 parents 755dd14 + 387568c commit 4088981

File tree

1 file changed

+19
-13
lines changed

1 file changed

+19
-13
lines changed

library/alloc/src/collections/binary_heap.rs

+19-13
Original file line numberDiff line numberDiff line change
@@ -495,7 +495,14 @@ impl<T: Ord> BinaryHeap<T> {
495495
let mut end = self.len();
496496
while end > 1 {
497497
end -= 1;
498-
self.data.swap(0, end);
498+
// SAFETY: `end` goes from `self.len() - 1` to 1 (both included),
499+
// so it's always a valid index to access.
500+
// It is safe to access index 0 (i.e. `ptr`), because
501+
// 1 <= end < self.len(), which means self.len() >= 2.
502+
unsafe {
503+
let ptr = self.data.as_mut_ptr();
504+
ptr::swap(ptr, ptr.add(end));
505+
}
499506
self.sift_down_range(0, end);
500507
}
501508
self.into_vec()
@@ -531,19 +538,19 @@ impl<T: Ord> BinaryHeap<T> {
531538
unsafe {
532539
let mut hole = Hole::new(&mut self.data, pos);
533540
let mut child = 2 * pos + 1;
534-
while child < end {
535-
let right = child + 1;
541+
while child < end - 1 {
536542
// compare with the greater of the two children
537-
if right < end && hole.get(child) <= hole.get(right) {
538-
child = right;
539-
}
543+
child += (hole.get(child) <= hole.get(child + 1)) as usize;
540544
// if we are already in order, stop.
541545
if hole.element() >= hole.get(child) {
542-
break;
546+
return;
543547
}
544548
hole.move_to(child);
545549
child = 2 * hole.pos() + 1;
546550
}
551+
if child == end - 1 && hole.element() < hole.get(child) {
552+
hole.move_to(child);
553+
}
547554
}
548555
}
549556

@@ -563,15 +570,14 @@ impl<T: Ord> BinaryHeap<T> {
563570
unsafe {
564571
let mut hole = Hole::new(&mut self.data, pos);
565572
let mut child = 2 * pos + 1;
566-
while child < end {
567-
let right = child + 1;
568-
// compare with the greater of the two children
569-
if right < end && hole.get(child) <= hole.get(right) {
570-
child = right;
571-
}
573+
while child < end - 1 {
574+
child += (hole.get(child) <= hole.get(child + 1)) as usize;
572575
hole.move_to(child);
573576
child = 2 * hole.pos() + 1;
574577
}
578+
if child == end - 1 {
579+
hole.move_to(child);
580+
}
575581
pos = hole.pos;
576582
}
577583
self.sift_up(start, pos);

0 commit comments

Comments
 (0)