[libc++] Add assertions for potential OOB reads in std::sort

We introduced an optimization to std::sort in 4eddbf9f10. However,
that optimization led to issues where users that were passing invalid
comparators to std::sort could start seeing OOB reads. This led to
the revert of the std::sort optimization from the LLVM 16 release
(see https://llvm.org/D146421).

This patch introduces _LIBCPP_ASSERTs to the places in the algorithm
where we make an assumption that the comparator will be consistent and
hence avoid a bounds check based on that. If the comparator happens not
to be consistent with itself, these are the places where we would
incorrectly go out of bounds. This allows users that enable libc++
assertions to catch such misuse at the cost of basically a bounds
check. For users that do not enable libc++ assertions (which is 99.9%
of users since assertions are off by default), this is basically a
no-op, and in fact the assertion will turn into a __builtin_assume,
making it explicit to the compiler that it can rely on the fact that
we're not going out of bounds.

I think this patch strikes the right balance. Folks that want absolute
performance will get what they want, since it is a precondition for the
comparator to be consistent, so the bounds checks are technically not
mandatory. Folks who want more safety *already* need to be enabling
libc++ assertions to catch other types of bugs (like operator[] OOB),
so this solution should also work for them.

I do think we have a lot of work towards popularizing the use of libc++
assertions and integrating it better so that users don't have to know
about the obscure _LIBCPP_ENABLE_ASSERTIONS macro to enable them, but
that's a separate concern.

rdar://106897934

Differential Revision: https://reviews.llvm.org/D147089

libcxx/docs/ReleaseNotes.rst

@@ -53,6 +53,9 @@ Improvements and New Features
 - The performance of ``dynamic_cast`` on its hot paths is greatly improved and is as efficient as the
   ``libsupc++`` implementation. Note that the performance improvements are shipped in ``libcxxabi``.
 
+- `D122780 <https://reviews.llvm.org/D122780>`_ Improved the performance of ``std::sort`` and ``std::ranges::sort``
+  by up to 50% for arithmetic types and by approximately 10% for other types.
+
 Deprecations and Removals
 -------------------------
 

libcxx/include/__algorithm/sort.h

@@ -279,12 +279,13 @@ void __insertion_sort(_BidirectionalIterator __first, _BidirectionalIterator __l
 // element in the input range is greater or equal to the element at __first - 1.
 template <class _AlgPolicy, class _Compare, class _RandomAccessIterator>
 _LIBCPP_HIDE_FROM_ABI void
-__insertion_sort_unguarded(_RandomAccessIterator __first, _RandomAccessIterator __last, _Compare __comp) {
+__insertion_sort_unguarded(_RandomAccessIterator const __first, _RandomAccessIterator __last, _Compare __comp) {
   using _Ops = _IterOps<_AlgPolicy>;
   typedef typename iterator_traits<_RandomAccessIterator>::difference_type difference_type;
   typedef typename iterator_traits<_RandomAccessIterator>::value_type value_type;
   if (__first == __last)
     return;
+  const _RandomAccessIterator __leftmost = __first - difference_type(1); (void)__leftmost; // can be unused when assertions are disabled
   for (_RandomAccessIterator __i = __first + difference_type(1); __i != __last; ++__i) {
     _RandomAccessIterator __j = __i - difference_type(1);
     if (__comp(*__i, *__j)) {
@@ -294,6 +295,7 @@ __insertion_sort_unguarded(_RandomAccessIterator __first, _RandomAccessIterator
       do {
         *__j = _Ops::__iter_move(__k);
         __j = __k;
+        _LIBCPP_ASSERT(__k != __leftmost, "Would read out of bounds, is your comparator a valid strict-weak ordering?");
       } while (__comp(__t, *--__k)); // No need for bounds check due to the assumption stated above.
       *__j = std::move(__t);
     }
@@ -496,14 +498,17 @@ __bitset_partition(_RandomAccessIterator __first, _RandomAccessIterator __last,
   typedef typename std::iterator_traits<_RandomAccessIterator>::value_type value_type;
   typedef typename std::iterator_traits<_RandomAccessIterator>::difference_type difference_type;
   _LIBCPP_ASSERT(__last - __first >= difference_type(3), "");
+  const _RandomAccessIterator __begin = __first;            // used for bounds checking, those are not moved around
+  const _RandomAccessIterator __end = __last; (void)__end;  //
 
-  _RandomAccessIterator __begin = __first;
   value_type __pivot(_Ops::__iter_move(__first));
   // Find the first element greater than the pivot.
   if (__comp(__pivot, *(__last - difference_type(1)))) {
     // Not guarded since we know the last element is greater than the pivot.
-    while (!__comp(__pivot, *++__first)) {
-    }
+    do {
+      ++__first;
+      _LIBCPP_ASSERT(__first != __end, "Would read out of bounds, is your comparator a valid strict-weak ordering?");
+    } while (!__comp(__pivot, *__first));
   } else {
     while (++__first < __last && !__comp(__pivot, *__first)) {
     }
@@ -512,8 +517,10 @@ __bitset_partition(_RandomAccessIterator __first, _RandomAccessIterator __last,
   if (__first < __last) {
     // It will be always guarded because __introsort will do the median-of-three
     // before calling this.
-    while (__comp(__pivot, *--__last)) {
-    }
+    do {
+      _LIBCPP_ASSERT(__last != __begin, "Would read out of bounds, is your comparator a valid strict-weak ordering?");
+      --__last;
+    } while (__comp(__pivot, *__last));
   }
   // If the first element greater than the pivot is at or after the
   // last element less than or equal to the pivot, then we have covered the
@@ -578,13 +585,16 @@ __partition_with_equals_on_right(_RandomAccessIterator __first, _RandomAccessIte
   typedef typename iterator_traits<_RandomAccessIterator>::difference_type difference_type;
   typedef typename std::iterator_traits<_RandomAccessIterator>::value_type value_type;
   _LIBCPP_ASSERT(__last - __first >= difference_type(3), "");
-  _RandomAccessIterator __begin = __first;
+  const _RandomAccessIterator __begin = __first;            // used for bounds checking, those are not moved around
+  const _RandomAccessIterator __end = __last; (void)__end;  //
   value_type __pivot(_Ops::__iter_move(__first));
   // Find the first element greater or equal to the pivot.  It will be always
   // guarded because __introsort will do the median-of-three before calling
   // this.
-  while (__comp(*++__first, __pivot))
-    ;
+  do {
+    ++__first;
+    _LIBCPP_ASSERT(__first != __end, "Would read out of bounds, is your comparator a valid strict-weak ordering?");
+  } while (__comp(*__first, __pivot));
 
   // Find the last element less than the pivot.
   if (__begin == __first - difference_type(1)) {
@@ -592,8 +602,10 @@ __partition_with_equals_on_right(_RandomAccessIterator __first, _RandomAccessIte
       ;
   } else {
     // Guarded.
-    while (!__comp(*--__last, __pivot))
-      ;
+    do {
+      _LIBCPP_ASSERT(__last != __begin, "Would read out of bounds, is your comparator a valid strict-weak ordering?");
+      --__last;
+    } while (!__comp(*__last, __pivot));
   }
 
   // If the first element greater than or equal to the pivot is at or after the
@@ -605,10 +617,14 @@ __partition_with_equals_on_right(_RandomAccessIterator __first, _RandomAccessIte
   // correct side of the pivot.
   while (__first < __last) {
     _Ops::iter_swap(__first, __last);
-    while (__comp(*++__first, __pivot))
-      ;
-    while (!__comp(*--__last, __pivot))
-      ;
+    do {
+      ++__first;
+      _LIBCPP_ASSERT(__first != __end, "Would read out of bounds, is your comparator a valid strict-weak ordering?");
+    } while (__comp(*__first, __pivot));
+    do {
+      _LIBCPP_ASSERT(__last != __begin, "Would read out of bounds, is your comparator a valid strict-weak ordering?");
+      --__last;
+    } while (!__comp(*__last, __pivot));
   }
   // Move the pivot to its correct position.
   _RandomAccessIterator __pivot_pos = __first - difference_type(1);
@@ -627,12 +643,15 @@ __partition_with_equals_on_left(_RandomAccessIterator __first, _RandomAccessIter
   using _Ops = _IterOps<_AlgPolicy>;
   typedef typename iterator_traits<_RandomAccessIterator>::difference_type difference_type;
   typedef typename std::iterator_traits<_RandomAccessIterator>::value_type value_type;
-  _RandomAccessIterator __begin = __first;
+  const _RandomAccessIterator __begin = __first;            // used for bounds checking, those are not moved around
+  const _RandomAccessIterator __end = __last; (void)__end;  //
   value_type __pivot(_Ops::__iter_move(__first));
   if (__comp(__pivot, *(__last - difference_type(1)))) {
     // Guarded.
-    while (!__comp(__pivot, *++__first)) {
-    }
+    do {
+      ++__first;
+      _LIBCPP_ASSERT(__first != __end, "Would read out of bounds, is your comparator a valid strict-weak ordering?");
+    } while (!__comp(__pivot, *__first));
   } else {
     while (++__first < __last && !__comp(__pivot, *__first)) {
     }
@@ -641,15 +660,21 @@ __partition_with_equals_on_left(_RandomAccessIterator __first, _RandomAccessIter
   if (__first < __last) {
     // It will be always guarded because __introsort will do the
     // median-of-three before calling this.
-    while (__comp(__pivot, *--__last)) {
-    }
+    do {
+      _LIBCPP_ASSERT(__last != __begin, "Would read out of bounds, is your comparator a valid strict-weak ordering?");
+      --__last;
+    } while (__comp(__pivot, *__last));
   }
   while (__first < __last) {
     _Ops::iter_swap(__first, __last);
-    while (!__comp(__pivot, *++__first))
-      ;
-    while (__comp(__pivot, *--__last))
-      ;
+    do {
+      ++__first;
+      _LIBCPP_ASSERT(__first != __end, "Would read out of bounds, is your comparator a valid strict-weak ordering?");
+    } while (!__comp(__pivot, *__first));
+    do {
+      _LIBCPP_ASSERT(__last != __begin, "Would read out of bounds, is your comparator a valid strict-weak ordering?");
+      --__last;
+    } while (__comp(__pivot, *__last));
   }
   _RandomAccessIterator __pivot_pos = __first - difference_type(1);
   if (__begin != __pivot_pos) {

libcxx/test/libcxx/algorithms/alg.sorting/assert.sort.invalid_comparator.pass.cpp

@@ -0,0 +1,152 @@
+//===----------------------------------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+// REQUIRES: stdlib=libc++
+
+// REQUIRES: has-unix-headers
+// UNSUPPORTED: c++03, c++11, c++14, c++17
+// XFAIL: availability-verbose_abort-missing
+// ADDITIONAL_COMPILE_FLAGS: -D_LIBCPP_ENABLE_ASSERTIONS=1
+
+// This test uses a specific combination of an invalid comparator and sequence of values to
+// ensure that our sorting functions do not go out-of-bounds in that case. Instead, we should
+// fail loud with an assertion. The specific issue we're looking for here is when the comparator
+// does not satisfy the following property:
+//
+//    comp(a, b) implies that !comp(b, a)
+//
+// In other words,
+//
+//    a < b implies that !(b < a)
+//
+// If this is not satisfied, we have seen issues in the past where the std::sort implementation
+// would proceed to do OOB reads.
+
+// When the debug mode is enabled, this test fails because we actually catch that the comparator
+// is not a strict-weak ordering before we catch that we'd dereference out-of-bounds inside std::sort,
+// which leads to different errors than the ones tested below.
+// XFAIL: libcpp-has-debug-mode
+
+#include <algorithm>
+#include <cassert>
+#include <cstddef>
+#include <map>
+#include <memory>
+#include <ranges>
+#include <set>
+#include <string>
+#include <vector>
+
+#include "check_assertion.h"
+
+std::string DATA =
+#   include "bad_comparator_values.dat"
+;
+
+int main(int, char**) {
+    std::map<std::size_t, std::map<std::size_t, bool>> comparison_results; // terrible for performance, but really convenient
+    for (auto line : std::views::split(DATA, '\n') | std::views::filter([](auto const& line) { return !line.empty(); })) {
+        auto values = std::views::split(line, ' ');
+        auto it = values.begin();
+        std::size_t left = std::stol(std::string((*it).data(), (*it).size()));
+        it = std::next(it);
+        std::size_t right = std::stol(std::string((*it).data(), (*it).size()));
+        it = std::next(it);
+        bool result = static_cast<bool>(std::stol(std::string((*it).data(), (*it).size())));
+        comparison_results[left][right] = result;
+    }
+    auto predicate = [&](std::size_t* left, std::size_t* right) {
+        assert(left != nullptr && right != nullptr && "something is wrong with the test");
+        assert(comparison_results.contains(*left) && comparison_results[*left].contains(*right) && "malformed input data?");
+        return comparison_results[*left][*right];
+    };
+
+    std::vector<std::unique_ptr<std::size_t>> elements;
+    std::set<std::size_t*> valid_ptrs;
+    for (std::size_t i = 0; i != comparison_results.size(); ++i) {
+        elements.push_back(std::make_unique<std::size_t>(i));
+        valid_ptrs.insert(elements.back().get());
+    }
+
+    auto checked_predicate = [&](size_t* left, size_t* right) {
+        // If the pointers passed to the comparator are not in the set of pointers we
+        // set up above, then we're being passed garbage values from the algorithm
+        // because we're reading OOB.
+        assert(valid_ptrs.contains(left));
+        assert(valid_ptrs.contains(right));
+        return predicate(left, right);
+    };
+
+    // Check the classic sorting algorithms
+    {
+        std::vector<std::size_t*> copy;
+        for (auto const& e : elements)
+            copy.push_back(e.get());
+        TEST_LIBCPP_ASSERT_FAILURE(std::sort(copy.begin(), copy.end(), checked_predicate), "Would read out of bounds");
+    }
+    {
+        std::vector<std::size_t*> copy;
+        for (auto const& e : elements)
+            copy.push_back(e.get());
+        std::stable_sort(copy.begin(), copy.end(), checked_predicate); // doesn't go OOB even with invalid comparator
+    }
+    {
+        std::vector<std::size_t*> copy;
+        for (auto const& e : elements)
+            copy.push_back(e.get());
+        std::partial_sort(copy.begin(), copy.begin(), copy.end(), checked_predicate); // doesn't go OOB even with invalid comparator
+    }
+    {
+        std::vector<std::size_t*> copy;
+        for (auto const& e : elements)
+            copy.push_back(e.get());
+        std::vector<std::size_t*> results(copy.size(), nullptr);
+        std::partial_sort_copy(copy.begin(), copy.end(), results.begin(), results.end(), checked_predicate); // doesn't go OOB even with invalid comparator
+    }
+    {
+        std::vector<std::size_t*> copy;
+        for (auto const& e : elements)
+            copy.push_back(e.get());
+        std::nth_element(copy.begin(), copy.end(), copy.end(), checked_predicate); // doesn't go OOB even with invalid comparator
+    }
+
+    // Check the Ranges sorting algorithms
+    {
+        std::vector<std::size_t*> copy;
+        for (auto const& e : elements)
+            copy.push_back(e.get());
+        TEST_LIBCPP_ASSERT_FAILURE(std::ranges::sort(copy, checked_predicate), "Would read out of bounds");
+    }
+    {
+        std::vector<std::size_t*> copy;
+        for (auto const& e : elements)
+            copy.push_back(e.get());
+        std::ranges::stable_sort(copy, checked_predicate); // doesn't go OOB even with invalid comparator
+    }
+    {
+        std::vector<std::size_t*> copy;
+        for (auto const& e : elements)
+            copy.push_back(e.get());
+        std::ranges::partial_sort(copy, copy.begin(), checked_predicate); // doesn't go OOB even with invalid comparator
+    }
+    {
+        std::vector<std::size_t*> copy;
+        for (auto const& e : elements)
+            copy.push_back(e.get());
+        std::vector<std::size_t*> results(copy.size(), nullptr);
+        std::ranges::partial_sort_copy(copy, results, checked_predicate); // doesn't go OOB even with invalid comparator
+    }
+    {
+        std::vector<std::size_t*> copy;
+        for (auto const& e : elements)
+            copy.push_back(e.get());
+        std::ranges::nth_element(copy, copy.end(), checked_predicate); // doesn't go OOB even with invalid comparator
+    }
+
+    return 0;
+}

libcxx/test/support/check_assertion.h

@@ -237,8 +237,6 @@ private:
 };
 
 void std::__libcpp_verbose_abort(char const* format, ...) {
-  assert(!GlobalMatcher().empty());
-
   // Extract information from the error message. This has to stay synchronized with
   // how we format assertions in the library.
   va_list list;