[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 53ef8e0d319c0bce7c77ab3b72e70ac2da3ac069 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 | 64 +++++++++++++++-----------
 2 files changed, 42 insertions(+), 30 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..07d8087afaa838 100644
--- a/libcxx/include/__algorithm/sift_down.h
+++ b/libcxx/include/__algorithm/sift_down.h
@@ -34,20 +34,22 @@ __sift_down(_RandomAccessIterator __first,
 
   typedef typename iterator_traits<_RandomAccessIterator>::difference_type difference_type;
   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)
+  if (__len < 2)
     return;
 
-  __child                         = 2 * __child + 1;
+  // left-child of __start is at 2 * __start + 1
+  // right-child of __start is at 2 * __start + 2
+  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_i = _Ops::next(__child_i);
+    if (__comp(*__child_i, *__right_i)) {
+      // right-child exists and is greater than left-child
+      __child_i = __right_i;
+      ++__child;
+    }
   }
 
   // check if we are in heap-order
@@ -68,10 +70,13 @@ __sift_down(_RandomAccessIterator __first,
     __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_i = _Ops::next(__child_i);
+      if (__comp(*__child_i, *__right_i)) {
+        // right-child exists and is greater than left-child
+        __child_i = __right_i;
+        ++__child;
+      }
     }
 
     // check if we are in heap-order
@@ -84,29 +89,34 @@ _LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX14 _RandomAccessIterator __floy
     _RandomAccessIterator __first,
     _Compare&& __comp,
     typename iterator_traits<_RandomAccessIterator>::difference_type __len) {
-  using difference_type = typename iterator_traits<_RandomAccessIterator>::difference_type;
-  _LIBCPP_ASSERT_INTERNAL(__len >= 2, "shouldn't be called unless __len >= 2");
+  _LIBCPP_ASSERT_INTERNAL(__len > 1, "shouldn't be called unless __len > 1");
 
-  _RandomAccessIterator __hole    = __first;
-  _RandomAccessIterator __child_i = __first;
-  difference_type __child         = 0;
+  using _Ops = _IterOps<_AlgPolicy>;
 
-  while (true) {
-    __child_i += difference_type(__child + 1);
-    __child = 2 * __child + 1;
+  typedef typename iterator_traits<_RandomAccessIterator>::difference_type difference_type;
 
-    if ((__child + 1) < __len && __comp(*__child_i, *(__child_i + difference_type(1)))) {
-      // right-child exists and is greater than left-child
-      ++__child_i;
-      ++__child;
+  difference_type __child      = 1;
+  _RandomAccessIterator __hole = __first, __child_i = __first;
+
+  while (true) {
+    __child_i += __child;
+    __child *= 2;
+
+    if (__child < __len) {
+      _RandomAccessIterator __right_i = _Ops::next(__child_i);
+      if (__comp(*__child_i, *__right_i)) {
+        // right-child exists and is greater than left-child
+        __child_i = __right_i;
+        ++__child;
+      }
     }
 
     // swap __hole with its largest child
-    *__hole = _IterOps<_AlgPolicy>::__iter_move(__child_i);
+    *__hole = _Ops::__iter_move(__child_i);
     __hole  = __child_i;
 
     // if __hole is now a leaf, we're done
-    if (__child > (__len - 2) / 2)
+    if (__child > __len / 2 - 2)
       return __hole;
   }
 }