[libcxx-commits] [libcxx] Libcxx nth element small ranges (PR #215698)

[Adib Pathan via libcxx-commits]

================
@@ -28,6 +28,43 @@ _LIBCPP_PUSH_MACROS
 
 _LIBCPP_BEGIN_NAMESPACE_STD
 
+// Handles ranges of at most 5 elements by sorting them entirely, which trivially satisfies the
+// nth_element postconditions. Sorting networks are the most efficient way to handle such small
+// ranges, and this function boils down to one of them when the size of the range is known at
+// compile time. Returns false if the range was too large to be handled.
+template <class _AlgPolicy, class _Compare, class _RandomAccessIterator>
+_LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX14 bool
+__nth_element_small(_RandomAccessIterator __first, _RandomAccessIterator __last, _Compare __comp) {
+  using _Ops = _IterOps<_AlgPolicy>;
+  typedef typename iterator_traits<_RandomAccessIterator>::difference_type difference_type;
+  switch (__last - __first) {
+  case 0:
+  case 1:
+    return true;
+  case 2:
+    if (__comp(*--__last, *__first))
+      _Ops::iter_swap(__first, __last);
+    return true;
+  case 3:
+    std::__sort3<_AlgPolicy, _Compare>(__first, __first + difference_type(1), --__last, __comp);
+    return true;
+  case 4:
+    std::__sort4<_AlgPolicy, _Compare>(
+        __first, __first + difference_type(1), __first + difference_type(2), --__last, __comp);
+    return true;
+  case 5:
+    std::__sort5<_AlgPolicy, _Compare>(
----------------
adxb-pxthxn wrote:

You're right, a full sort does more work than nth_element needs. I switched the small-range path to selection instead of __sort3/__sort4/__sort5.

5 elements, median (nth == first + 2): 6 comparisons (Knuth, TAOCP 5.3.3). Pair, order the pairs by their larger element (that larger value cannot be the median), insert the fifth, order again, then take the larger of the two remaining candidates. That also partitions, so the nth_element postconditions hold. Other nth positions still use __sort5.

4 elements: min/max in 3 comparisons, 2nd smallest/largest in 4 (V₂(4)). A 4-sort is 5.

3 elements: min/max in 2. The middle element is a median-of-3, which is a 3-sort, so that case still calls __sort3.

For cheap arithmetic + less/greater I kept branchless networks (7 comparators for median-of-5, __sort4/__sort3 otherwise). A data-oblivious network can't hit the 6-comparison bound, but no control flow is faster there, same split as __sort5.

I checked the comparison counts with a counting comparator over all permutations of 3/4/5.

https://github.com/llvm/llvm-project/pull/215698