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

Yang Kun via libcxx-commits libcxx-commits at lists.llvm.org
Fri Apr 25 21:53:33 PDT 2025


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

>From ba27b11fd4811580cb32b0585cd3a0365e694954 Mon Sep 17 00:00:00 2001
From: Yang Kun <ikspress at outlook.com>
Date: Sat, 4 Jan 2025 19:44:44 +0800
Subject: [PATCH 1/2] [libc++] Optimize heap operations

---
 libcxx/include/__algorithm/make_heap.h        |  9 ++-
 libcxx/include/__algorithm/partial_sort.h     |  2 +-
 .../include/__algorithm/partial_sort_copy.h   |  2 +-
 libcxx/include/__algorithm/pop_heap.h         |  3 +-
 libcxx/include/__algorithm/push_heap.h        | 48 ++++++------
 libcxx/include/__algorithm/sift_down.h        | 78 +++++++++++--------
 6 files changed, 75 insertions(+), 67 deletions(-)

diff --git a/libcxx/include/__algorithm/make_heap.h b/libcxx/include/__algorithm/make_heap.h
index e8f0cdb27333a..74d23c0f613ca 100644
--- a/libcxx/include/__algorithm/make_heap.h
+++ b/libcxx/include/__algorithm/make_heap.h
@@ -34,10 +34,11 @@ __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) {
-      std::__sift_down<_AlgPolicy>(__first, __comp_ref, __n, __first + __start);
-    }
+    difference_type __start = __n / 2;
+    do {
+      // start from the first parent, there is no need to consider children
+      std::__sift_down<_AlgPolicy>(__first, __comp_ref, __n, --__start);
+    } while (__start != 0);
   }
 }
 
diff --git a/libcxx/include/__algorithm/partial_sort.h b/libcxx/include/__algorithm/partial_sort.h
index 7f8d0c49147e3..46e48c33f798d 100644
--- a/libcxx/include/__algorithm/partial_sort.h
+++ b/libcxx/include/__algorithm/partial_sort.h
@@ -45,7 +45,7 @@ _LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX20 _RandomAccessIterator __part
   for (; __i != __last; ++__i) {
     if (__comp(*__i, *__first)) {
       _IterOps<_AlgPolicy>::iter_swap(__i, __first);
-      std::__sift_down<_AlgPolicy>(__first, __comp, __len, __first);
+      std::__sift_down<_AlgPolicy>(__first, __comp, __len, 0);
     }
   }
   std::__sort_heap<_AlgPolicy>(std::move(__first), std::move(__middle), __comp);
diff --git a/libcxx/include/__algorithm/partial_sort_copy.h b/libcxx/include/__algorithm/partial_sort_copy.h
index 172f53b290d54..a4437e4f932ce 100644
--- a/libcxx/include/__algorithm/partial_sort_copy.h
+++ b/libcxx/include/__algorithm/partial_sort_copy.h
@@ -60,7 +60,7 @@ _LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX20 pair<_InputIterator, _Random
     for (; __first != __last; ++__first)
       if (std::__invoke(__comp, std::__invoke(__proj1, *__first), std::__invoke(__proj2, *__result_first))) {
         *__result_first = *__first;
-        std::__sift_down<_AlgPolicy>(__result_first, __projected_comp, __len, __result_first);
+        std::__sift_down<_AlgPolicy>(__result_first, __projected_comp, __len, 0);
       }
     std::__sort_heap<_AlgPolicy>(__result_first, __r, __projected_comp);
   }
diff --git a/libcxx/include/__algorithm/pop_heap.h b/libcxx/include/__algorithm/pop_heap.h
index 6d23830097ff9..60481938c4bc5 100644
--- a/libcxx/include/__algorithm/pop_heap.h
+++ b/libcxx/include/__algorithm/pop_heap.h
@@ -51,9 +51,8 @@ __pop_heap(_RandomAccessIterator __first,
       *__hole = std::move(__top);
     } else {
       *__hole = _IterOps<_AlgPolicy>::__iter_move(__last);
-      ++__hole;
       *__last = std::move(__top);
-      std::__sift_up<_AlgPolicy>(__first, __hole, __comp_ref, __hole - __first);
+      std::__sift_up<_AlgPolicy>(__first, __hole, __comp_ref);
     }
   }
 }
diff --git a/libcxx/include/__algorithm/push_heap.h b/libcxx/include/__algorithm/push_heap.h
index ec0b445f2b70f..798cc3d28c961 100644
--- a/libcxx/include/__algorithm/push_heap.h
+++ b/libcxx/include/__algorithm/push_heap.h
@@ -29,37 +29,35 @@ _LIBCPP_BEGIN_NAMESPACE_STD
 
 template <class _AlgPolicy, class _Compare, class _RandomAccessIterator>
 _LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX14 void
-__sift_up(_RandomAccessIterator __first,
-          _RandomAccessIterator __last,
-          _Compare&& __comp,
-          typename iterator_traits<_RandomAccessIterator>::difference_type __len) {
-  using value_type = typename iterator_traits<_RandomAccessIterator>::value_type;
-
-  if (__len > 1) {
-    __len                       = (__len - 2) / 2;
-    _RandomAccessIterator __ptr = __first + __len;
-
-    if (__comp(*__ptr, *--__last)) {
-      value_type __t(_IterOps<_AlgPolicy>::__iter_move(__last));
-      do {
-        *__last = _IterOps<_AlgPolicy>::__iter_move(__ptr);
-        __last  = __ptr;
-        if (__len == 0)
-          break;
-        __len = (__len - 1) / 2;
-        __ptr = __first + __len;
-      } while (__comp(*__ptr, __t));
-
-      *__last = std::move(__t);
-    }
+__sift_up(_RandomAccessIterator __first, _RandomAccessIterator __last, _Compare&& __comp) {
+  using difference_type = typename iterator_traits<_RandomAccessIterator>::difference_type;
+  using value_type      = typename iterator_traits<_RandomAccessIterator>::value_type;
+
+  difference_type __parent = __last - __first;
+  _LIBCPP_ASSERT_INTERNAL(__parent > 0, "shouldn't be called unless __last - __first > 0");
+  __parent                         = (__parent - 1) / 2;
+  _RandomAccessIterator __parent_i = __first + __parent;
+
+  if (__comp(*__parent_i, *__last)) {
+    value_type __t(_IterOps<_AlgPolicy>::__iter_move(__last));
+    do {
+      *__last = _IterOps<_AlgPolicy>::__iter_move(__parent_i);
+      __last  = __parent_i;
+      if (__parent == 0)
+        break;
+      __parent   = (__parent - 1) / 2;
+      __parent_i = __first + __parent;
+    } while (__comp(*__parent_i, __t));
+
+    *__last = std::move(__t);
   }
 }
 
 template <class _AlgPolicy, class _RandomAccessIterator, class _Compare>
 inline _LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX14 void
 __push_heap(_RandomAccessIterator __first, _RandomAccessIterator __last, _Compare& __comp) {
-  typename iterator_traits<_RandomAccessIterator>::difference_type __len = __last - __first;
-  std::__sift_up<_AlgPolicy, __comp_ref_type<_Compare> >(std::move(__first), std::move(__last), __comp, __len);
+  if (__first != __last)
+    std::__sift_up<_AlgPolicy, __comp_ref_type<_Compare> >(std::move(__first), std::move(--__last), __comp);
 }
 
 template <class _RandomAccessIterator, class _Compare>
diff --git a/libcxx/include/__algorithm/sift_down.h b/libcxx/include/__algorithm/sift_down.h
index 42803e30631fb..f413a7f55c559 100644
--- a/libcxx/include/__algorithm/sift_down.h
+++ b/libcxx/include/__algorithm/sift_down.h
@@ -29,37 +29,39 @@ _LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX14 void
 __sift_down(_RandomAccessIterator __first,
             _Compare&& __comp,
             typename iterator_traits<_RandomAccessIterator>::difference_type __len,
-            _RandomAccessIterator __start) {
+            typename iterator_traits<_RandomAccessIterator>::difference_type __start) {
   using _Ops = _IterOps<_AlgPolicy>;
 
   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;
-  _RandomAccessIterator __child_i = __first + __child;
+  // left-child of __start is at 2 * __start + 1
+  // right-child of __start is at 2 * __start + 2
+  difference_type __child         = 2 * __start + 1;
+  _RandomAccessIterator __child_i = __first + __child, __start_i = __first + __start;
 
-  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
-  if (__comp(*__child_i, *__start))
+  if (__comp(*__child_i, *__start_i))
     // we are, __start is larger than its largest child
     return;
 
-  value_type __top(_Ops::__iter_move(__start));
+  value_type __top(_Ops::__iter_move(__start_i));
   do {
     // we are not in heap-order, swap the parent with its largest child
-    *__start = _Ops::__iter_move(__child_i);
-    __start  = __child_i;
+    *__start_i = _Ops::__iter_move(__child_i);
+    __start_i  = __child_i;
 
     if ((__len - 2) / 2 < __child)
       break;
@@ -68,15 +70,18 @@ __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
   } while (!__comp(*__child_i, __top));
-  *__start = std::move(__top);
+  *__start_i = std::move(__top);
 }
 
 template <class _AlgPolicy, class _Compare, class _RandomAccessIterator>
@@ -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)
       return __hole;
   }
 }

>From b46825829af7cd83e5e49cf1d5dc594d7a44becc Mon Sep 17 00:00:00 2001
From: Yang Kun <ikspress at outlook.com>
Date: Sat, 26 Apr 2025 12:42:03 +0800
Subject: [PATCH 2/2] Update

---
 libcxx/include/__algorithm/make_heap.h        | 11 +++---
 .../include/__algorithm/partial_sort_copy.h   | 12 +++----
 libcxx/include/__algorithm/pop_heap.h         | 34 +++++++++----------
 libcxx/include/__algorithm/sift_down.h        |  6 ++--
 libcxx/include/__algorithm/sort_heap.h        |  4 +--
 5 files changed, 31 insertions(+), 36 deletions(-)

diff --git a/libcxx/include/__algorithm/make_heap.h b/libcxx/include/__algorithm/make_heap.h
index 74d23c0f613ca..16291e45b49cf 100644
--- a/libcxx/include/__algorithm/make_heap.h
+++ b/libcxx/include/__algorithm/make_heap.h
@@ -32,13 +32,10 @@ __make_heap(_RandomAccessIterator __first, _RandomAccessIterator __last, _Compar
   __comp_ref_type<_Compare> __comp_ref = __comp;
 
   using difference_type = typename iterator_traits<_RandomAccessIterator>::difference_type;
-  difference_type __n   = __last - __first;
-  if (__n > 1) {
-    difference_type __start = __n / 2;
-    do {
-      // start from the first parent, there is no need to consider children
-      std::__sift_down<_AlgPolicy>(__first, __comp_ref, __n, --__start);
-    } while (__start != 0);
+  difference_type __n = __last - __first, __start = __n / 2;
+  while (__start != 0) {
+    // start from the first parent, there is no need to consider children
+    std::__sift_down<_AlgPolicy>(__first, __comp_ref, __n, --__start);
   }
 }
 
diff --git a/libcxx/include/__algorithm/partial_sort_copy.h b/libcxx/include/__algorithm/partial_sort_copy.h
index a4437e4f932ce..5a47f0c8a316b 100644
--- a/libcxx/include/__algorithm/partial_sort_copy.h
+++ b/libcxx/include/__algorithm/partial_sort_copy.h
@@ -11,6 +11,7 @@
 
 #include <__algorithm/comp.h>
 #include <__algorithm/comp_ref_type.h>
+#include <__algorithm/copy.h>
 #include <__algorithm/iterator_operations.h>
 #include <__algorithm/make_heap.h>
 #include <__algorithm/make_projected.h>
@@ -49,19 +50,18 @@ _LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX20 pair<_InputIterator, _Random
     _Compare&& __comp,
     _Proj1&& __proj1,
     _Proj2&& __proj2) {
-  _RandomAccessIterator __r = __result_first;
-  auto&& __projected_comp   = std::__make_projected(__comp, __proj2);
+  auto&& __projected_comp = std::__make_projected(__comp, __proj2);
 
-  if (__r != __result_last) {
-    for (; __first != __last && __r != __result_last; ++__first, (void)++__r)
-      *__r = *__first;
+  if (__result_first != __result_last) {
+    _RandomAccessIterator __r = std::copy(__first, __last, std::move(__result_first));
     std::__make_heap<_AlgPolicy>(__result_first, __r, __projected_comp);
     typename iterator_traits<_RandomAccessIterator>::difference_type __len = __r - __result_first;
-    for (; __first != __last; ++__first)
+    for (; __first != __last; ++__first) {
       if (std::__invoke(__comp, std::__invoke(__proj1, *__first), std::__invoke(__proj2, *__result_first))) {
         *__result_first = *__first;
         std::__sift_down<_AlgPolicy>(__result_first, __projected_comp, __len, 0);
       }
+    }
     std::__sort_heap<_AlgPolicy>(__result_first, __r, __projected_comp);
   }
 
diff --git a/libcxx/include/__algorithm/pop_heap.h b/libcxx/include/__algorithm/pop_heap.h
index 60481938c4bc5..35e2ba1593933 100644
--- a/libcxx/include/__algorithm/pop_heap.h
+++ b/libcxx/include/__algorithm/pop_heap.h
@@ -33,27 +33,20 @@ _LIBCPP_BEGIN_NAMESPACE_STD
 template <class _AlgPolicy, class _Compare, class _RandomAccessIterator>
 inline _LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX14 void
 __pop_heap(_RandomAccessIterator __first,
-           _RandomAccessIterator __last,
+           _RandomAccessIterator __bottom,
            _Compare& __comp,
            typename iterator_traits<_RandomAccessIterator>::difference_type __len) {
-  // Calling `pop_heap` on an empty range is undefined behavior, but in practice it will be a no-op.
-  _LIBCPP_ASSERT_PEDANTIC(__len > 0, "The heap given to pop_heap must be non-empty");
+  using value_type = typename iterator_traits<_RandomAccessIterator>::value_type;
 
-  __comp_ref_type<_Compare> __comp_ref = __comp;
+  value_type __top             = _IterOps<_AlgPolicy>::__iter_move(__first); // create a hole at __first
+  _RandomAccessIterator __hole = std::__floyd_sift_down<_AlgPolicy>(__first, __comp, __len);
 
-  using value_type = typename iterator_traits<_RandomAccessIterator>::value_type;
-  if (__len > 1) {
-    value_type __top             = _IterOps<_AlgPolicy>::__iter_move(__first); // create a hole at __first
-    _RandomAccessIterator __hole = std::__floyd_sift_down<_AlgPolicy>(__first, __comp_ref, __len);
-    --__last;
-
-    if (__hole == __last) {
-      *__hole = std::move(__top);
-    } else {
-      *__hole = _IterOps<_AlgPolicy>::__iter_move(__last);
-      *__last = std::move(__top);
-      std::__sift_up<_AlgPolicy>(__first, __hole, __comp_ref);
-    }
+  if (__hole == __bottom) {
+    *__hole = std::move(__top);
+  } else {
+    *__hole   = _IterOps<_AlgPolicy>::__iter_move(__bottom);
+    *__bottom = std::move(__top);
+    std::__sift_up<_AlgPolicy>(__first, __hole, __comp);
   }
 }
 
@@ -63,8 +56,13 @@ pop_heap(_RandomAccessIterator __first, _RandomAccessIterator __last, _Compare _
   static_assert(std::is_copy_constructible<_RandomAccessIterator>::value, "Iterators must be copy constructible.");
   static_assert(std::is_copy_assignable<_RandomAccessIterator>::value, "Iterators must be copy assignable.");
 
+  // Calling `pop_heap` on an empty range is undefined behavior, but in practice it will be a no-op.
+  _LIBCPP_ASSERT_PEDANTIC(__len > 0, "The heap given to pop_heap must be non-empty");
+
+  __comp_ref_type<_Compare> __comp_ref = __comp;
+
   typename iterator_traits<_RandomAccessIterator>::difference_type __len = __last - __first;
-  std::__pop_heap<_ClassicAlgPolicy>(std::move(__first), std::move(__last), __comp, __len);
+  std::__pop_heap<_ClassicAlgPolicy>(std::move(__first), std::move(--__last), __comp_ref, __len);
 }
 
 template <class _RandomAccessIterator>
diff --git a/libcxx/include/__algorithm/sift_down.h b/libcxx/include/__algorithm/sift_down.h
index f413a7f55c559..83194a0d4afeb 100644
--- a/libcxx/include/__algorithm/sift_down.h
+++ b/libcxx/include/__algorithm/sift_down.h
@@ -43,7 +43,7 @@ __sift_down(_RandomAccessIterator __first,
   difference_type __child         = 2 * __start + 1;
   _RandomAccessIterator __child_i = __first + __child, __start_i = __first + __start;
 
-  if ((__child + 1) < __len) {
+  if (__child < __len - 1) {
     _RandomAccessIterator __right_i = _Ops::next(__child_i);
     if (__comp(*__child_i, *__right_i)) {
       // right-child exists and is greater than left-child
@@ -63,14 +63,14 @@ __sift_down(_RandomAccessIterator __first,
     *__start_i = _Ops::__iter_move(__child_i);
     __start_i  = __child_i;
 
-    if ((__len - 2) / 2 < __child)
+    if (__len / 2 - 1 < __child)
       break;
 
     // recompute the child based off of the updated parent
     __child   = 2 * __child + 1;
     __child_i = __first + __child;
 
-    if ((__child + 1) < __len) {
+    if (__child < __len - 1) {
       _RandomAccessIterator __right_i = _Ops::next(__child_i);
       if (__comp(*__child_i, *__right_i)) {
         // right-child exists and is greater than left-child
diff --git a/libcxx/include/__algorithm/sort_heap.h b/libcxx/include/__algorithm/sort_heap.h
index f20b110c7fd12..f2d73ab793650 100644
--- a/libcxx/include/__algorithm/sort_heap.h
+++ b/libcxx/include/__algorithm/sort_heap.h
@@ -36,8 +36,8 @@ __sort_heap(_RandomAccessIterator __first, _RandomAccessIterator __last, _Compar
   __comp_ref_type<_Compare> __comp_ref = __comp;
 
   using difference_type = typename iterator_traits<_RandomAccessIterator>::difference_type;
-  for (difference_type __n = __last - __first; __n > 1; --__last, (void)--__n)
-    std::__pop_heap<_AlgPolicy>(__first, __last, __comp_ref, __n);
+  for (difference_type __n = __last - __first; __n > 1; --__n)
+    std::__pop_heap<_AlgPolicy>(__first, --__last, __comp_ref, __n);
   std::__check_strict_weak_ordering_sorted(__first, __saved_last, __comp_ref);
 }
 



More information about the libcxx-commits mailing list