[libcxx-commits] [libcxx] [libc++] Optimize make_heap() and sift_down() (PR #121480)

[Yang Kun via libcxx-commits]

https://github.com/omikrun updated https://github.com/llvm/llvm-project/pull/121480

>From 497215e68f3863876396d1d88f1b80fca0d78d45 Mon Sep 17 00:00:00 2001
From: Yang Kun <193369907+omikrun at users.noreply.github.com>
Date: Sat, 4 Jan 2025 19:44:44 +0800
Subject: [PATCH] [libc++] Optimize make_heap() and sift_down()

---
 libcxx/include/__algorithm/make_heap.h |  8 ++++---
 libcxx/include/__algorithm/sift_down.h | 33 +++++++++++++-------------
 2 files changed, 22 insertions(+), 19 deletions(-)

diff --git a/libcxx/include/__algorithm/make_heap.h b/libcxx/include/__algorithm/make_heap.h
index e8f0cdb27333a4..746c2621fa70e6 100644
--- a/libcxx/include/__algorithm/make_heap.h
+++ b/libcxx/include/__algorithm/make_heap.h
@@ -34,10 +34,12 @@ __make_heap(_RandomAccessIterator __first, _RandomAccessIterator __last, _Compar
   using difference_type = typename iterator_traits<_RandomAccessIterator>::difference_type;
   difference_type __n   = __last - __first;
   if (__n > 1) {
-    // start from the first parent, there is no need to consider children
-    for (difference_type __start = (__n - 2) / 2; __start >= 0; --__start) {
+    difference_type __start = __n / 2;
+    do {
+      // start from the first parent, there is no need to consider children
+      --__start;
       std::__sift_down<_AlgPolicy>(__first, __comp_ref, __n, __first + __start);
-    }
+    } while (__start != 0);
   }
 }
 
diff --git a/libcxx/include/__algorithm/sift_down.h b/libcxx/include/__algorithm/sift_down.h
index 42803e30631fb1..0c5558f3c31698 100644
--- a/libcxx/include/__algorithm/sift_down.h
+++ b/libcxx/include/__algorithm/sift_down.h
@@ -36,18 +36,16 @@ __sift_down(_RandomAccessIterator __first,
   typedef typename iterator_traits<_RandomAccessIterator>::value_type value_type;
   // left-child of __start is at 2 * __start + 1
   // right-child of __start is at 2 * __start + 2
-  difference_type __child = __start - __first;
-
-  if (__len < 2 || (__len - 2) / 2 < __child)
-    return;
-
-  __child                         = 2 * __child + 1;
+  difference_type __child         = 2 * (__start - __first) + 1;
   _RandomAccessIterator __child_i = __first + __child;
 
-  if ((__child + 1) < __len && __comp(*__child_i, *(__child_i + difference_type(1)))) {
-    // right-child exists and is greater than left-child
-    ++__child_i;
-    ++__child;
+  if ((__child + 1) < __len) {
+    _RandomAccessIterator __right_child_i = _Ops::next(__child_i);
+    if (__comp(*__child_i, *__right_child_i)) {
+      // right-child exists and is greater than left-child
+      __child_i = __right_child_i;
+      ++__child;
+    }
   }
 
   // check if we are in heap-order
@@ -61,17 +59,20 @@ __sift_down(_RandomAccessIterator __first,
     *__start = _Ops::__iter_move(__child_i);
     __start  = __child_i;
 
-    if ((__len - 2) / 2 < __child)
+    __child = 2 * __child + 1;
+    if (!(__child < __len))
       break;
 
     // recompute the child based off of the updated parent
-    __child   = 2 * __child + 1;
     __child_i = __first + __child;
 
-    if ((__child + 1) < __len && __comp(*__child_i, *(__child_i + difference_type(1)))) {
-      // right-child exists and is greater than left-child
-      ++__child_i;
-      ++__child;
+    if ((__child + 1) < __len) {
+      _RandomAccessIterator __right_child_i = _Ops::next(__child_i);
+      if (__comp(*__child_i, *__right_child_i)) {
+        // right-child exists and is greater than left-child
+        __child_i = __right_child_i;
+        ++__child;
+      }
     }
 
     // check if we are in heap-order