1 files changed, 48 insertions, 23 deletions
diff --git a/libcxx/include/__algorithm/partial_sort.h b/libcxx/include/__algorithm/partial_sort.h
index e008c0c99679..24016e5cf5a5 100644
--- a/libcxx/include/__algorithm/partial_sort.h
+++ b/libcxx/include/__algorithm/partial_sort.h
@@ -11,6 +11,7 @@
 
 #include <__algorithm/comp.h>
 #include <__algorithm/comp_ref_type.h>
+#include <__algorithm/iterator_operations.h>
 #include <__algorithm/make_heap.h>
 #include <__algorithm/sift_down.h>
 #include <__algorithm/sort_heap.h>
@@ -18,7 +19,8 @@
 #include <__debug>
 #include <__debug_utils/randomize_range.h>
 #include <__iterator/iterator_traits.h>
-#include <__utility/swap.h>
+#include <__utility/move.h>
+#include <type_traits>
 
 #if !defined(_LIBCPP_HAS_NO_PRAGMA_SYSTEM_HEADER)
 #  pragma GCC system_header
@@ -26,24 +28,47 @@
 
 _LIBCPP_BEGIN_NAMESPACE_STD
 
-template <class _Compare, class _RandomAccessIterator>
-_LIBCPP_CONSTEXPR_AFTER_CXX17 void
-__partial_sort(_RandomAccessIterator __first, _RandomAccessIterator __middle, _RandomAccessIterator __last,
-               _Compare __comp)
-{
-    if (__first == __middle)
-        return;
-    _VSTD::__make_heap<_Compare>(__first, __middle, __comp);
-    typename iterator_traits<_RandomAccessIterator>::difference_type __len = __middle - __first;
-    for (_RandomAccessIterator __i = __middle; __i != __last; ++__i)
-    {
-        if (__comp(*__i, *__first))
-        {
-            swap(*__i, *__first);
-            _VSTD::__sift_down<_Compare>(__first, __comp, __len, __first);
-        }
-    }
-    _VSTD::__sort_heap<_Compare>(__first, __middle, __comp);
+template <class _AlgPolicy, class _Compare, class _RandomAccessIterator, class _Sentinel>
+_LIBCPP_CONSTEXPR_AFTER_CXX17
+_RandomAccessIterator __partial_sort_impl(
+    _RandomAccessIterator __first, _RandomAccessIterator __middle, _Sentinel __last, _Compare __comp) {
+  if (__first == __middle) {
+    return _IterOps<_AlgPolicy>::next(__middle, __last);
+  }
+
+  std::__make_heap<_AlgPolicy, _Compare>(__first, __middle, __comp);
+
+  typename iterator_traits<_RandomAccessIterator>::difference_type __len = __middle - __first;
+  _RandomAccessIterator __i = __middle;
+  for (; __i != __last; ++__i)
+  {
+      if (__comp(*__i, *__first))
+      {
+          _IterOps<_AlgPolicy>::iter_swap(__i, __first);
+          std::__sift_down<_AlgPolicy, _Compare>(__first, __comp, __len, __first);
+      }
+
+  }
+  std::__sort_heap<_AlgPolicy, _Compare>(std::move(__first), std::move(__middle), __comp);
+
+  return __i;
+}
+
+template <class _AlgPolicy, class _Compare, class _RandomAccessIterator, class _Sentinel>
+_LIBCPP_CONSTEXPR_AFTER_CXX17
+_RandomAccessIterator __partial_sort(_RandomAccessIterator __first, _RandomAccessIterator __middle, _Sentinel __last,
+                                     _Compare& __comp) {
+  if (__first == __middle)
+      return _IterOps<_AlgPolicy>::next(__middle, __last);
+
+  std::__debug_randomize_range<_AlgPolicy>(__first, __last);
+
+  using _Comp_ref = typename __comp_ref_type<_Compare>::type;
+  auto __last_iter = std::__partial_sort_impl<_AlgPolicy, _Comp_ref>(__first, __middle, __last, __comp);
+
+  std::__debug_randomize_range<_AlgPolicy>(__middle, __last);
+
+  return __last_iter;
 }
 
 template <class _RandomAccessIterator, class _Compare>
@@ -52,10 +77,10 @@ void
 partial_sort(_RandomAccessIterator __first, _RandomAccessIterator __middle, _RandomAccessIterator __last,
              _Compare __comp)
 {
-  std::__debug_randomize_range(__first, __last);
-  typedef typename __comp_ref_type<_Compare>::type _Comp_ref;
-  _VSTD::__partial_sort<_Comp_ref>(__first, __middle, __last, __comp);
-  std::__debug_randomize_range(__middle, __last);
+  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.");
+
+  (void)std::__partial_sort<_ClassicAlgPolicy>(std::move(__first), std::move(__middle), std::move(__last), __comp);
 }
 
 template <class _RandomAccessIterator>