1//===----------------------------------------------------------------------===//
2//
3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4// See https://llvm.org/LICENSE.txt for license information.
5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6//
7//===----------------------------------------------------------------------===//
8
9#ifndef _LIBCPP___ALGORITHM_NTH_ELEMENT_H
10#define _LIBCPP___ALGORITHM_NTH_ELEMENT_H
11
12#include <__algorithm/comp.h>
13#include <__algorithm/comp_ref_type.h>
14#include <__algorithm/iterator_operations.h>
15#include <__algorithm/sort.h>
16#include <__assert>
17#include <__config>
18#include <__debug_utils/randomize_range.h>
19#include <__iterator/iterator_traits.h>
20#include <__utility/move.h>
21
22#if !defined(_LIBCPP_HAS_NO_PRAGMA_SYSTEM_HEADER)
23#  pragma GCC system_header
24#endif
25
26_LIBCPP_PUSH_MACROS
27#include <__undef_macros>
28
29_LIBCPP_BEGIN_NAMESPACE_STD
30
31template <class _Compare, class _RandomAccessIterator>
32_LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX14 bool __nth_element_find_guard(
33    _RandomAccessIterator& __i, _RandomAccessIterator& __j, _RandomAccessIterator __m, _Compare __comp) {
34  // manually guard downward moving __j against __i
35  while (true) {
36    if (__i == --__j) {
37      return false;
38    }
39    if (__comp(*__j, *__m)) {
40      return true; // found guard for downward moving __j, now use unguarded partition
41    }
42  }
43}
44
45template <class _AlgPolicy, class _Compare, class _RandomAccessIterator>
46_LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX14 void
47// NOLINTNEXTLINE(readability-function-cognitive-complexity)
48__nth_element(
49    _RandomAccessIterator __first, _RandomAccessIterator __nth, _RandomAccessIterator __last, _Compare __comp) {
50  using _Ops = _IterOps<_AlgPolicy>;
51
52  // _Compare is known to be a reference type
53  typedef typename iterator_traits<_RandomAccessIterator>::difference_type difference_type;
54  const difference_type __limit = 7;
55  while (true) {
56    if (__nth == __last)
57      return;
58    difference_type __len = __last - __first;
59    switch (__len) {
60    case 0:
61    case 1:
62      return;
63    case 2:
64      if (__comp(*--__last, *__first))
65        _Ops::iter_swap(__first, __last);
66      return;
67    case 3: {
68      _RandomAccessIterator __m = __first;
69      std::__sort3<_AlgPolicy, _Compare>(__first, ++__m, --__last, __comp);
70      return;
71    }
72    }
73    if (__len <= __limit) {
74      std::__selection_sort<_AlgPolicy, _Compare>(__first, __last, __comp);
75      return;
76    }
77    // __len > __limit >= 3
78    _RandomAccessIterator __m   = __first + __len / 2;
79    _RandomAccessIterator __lm1 = __last;
80    unsigned __n_swaps          = std::__sort3<_AlgPolicy, _Compare>(__first, __m, --__lm1, __comp);
81    // *__m is median
82    // partition [__first, __m) < *__m and *__m <= [__m, __last)
83    // (this inhibits tossing elements equivalent to __m around unnecessarily)
84    _RandomAccessIterator __i = __first;
85    _RandomAccessIterator __j = __lm1;
86    // j points beyond range to be tested, *__lm1 is known to be <= *__m
87    // The search going up is known to be guarded but the search coming down isn't.
88    // Prime the downward search with a guard.
89    if (!__comp(*__i, *__m)) // if *__first == *__m
90    {
91      // *__first == *__m, *__first doesn't go in first part
92      if (std::__nth_element_find_guard<_Compare>(__i, __j, __m, __comp)) {
93        _Ops::iter_swap(__i, __j);
94        ++__n_swaps;
95      } else {
96        // *__first == *__m, *__m <= all other elements
97        // Partition instead into [__first, __i) == *__first and *__first < [__i, __last)
98        ++__i; // __first + 1
99        __j = __last;
100        if (!__comp(*__first, *--__j)) { // we need a guard if *__first == *(__last-1)
101          while (true) {
102            if (__i == __j) {
103              return; // [__first, __last) all equivalent elements
104            } else if (__comp(*__first, *__i)) {
105              _Ops::iter_swap(__i, __j);
106              ++__n_swaps;
107              ++__i;
108              break;
109            }
110            ++__i;
111          }
112        }
113        // [__first, __i) == *__first and *__first < [__j, __last) and __j == __last - 1
114        if (__i == __j) {
115          return;
116        }
117        while (true) {
118          while (!__comp(*__first, *__i)) {
119            ++__i;
120            _LIBCPP_ASSERT_VALID_ELEMENT_ACCESS(
121                __i != __last,
122                "Would read out of bounds, does your comparator satisfy the strict-weak ordering requirement?");
123          }
124          do {
125            _LIBCPP_ASSERT_VALID_ELEMENT_ACCESS(
126                __j != __first,
127                "Would read out of bounds, does your comparator satisfy the strict-weak ordering requirement?");
128            --__j;
129          } while (__comp(*__first, *__j));
130          if (__i >= __j)
131            break;
132          _Ops::iter_swap(__i, __j);
133          ++__n_swaps;
134          ++__i;
135        }
136        // [__first, __i) == *__first and *__first < [__i, __last)
137        // The first part is sorted,
138        if (__nth < __i) {
139          return;
140        }
141        // __nth_element the second part
142        // std::__nth_element<_Compare>(__i, __nth, __last, __comp);
143        __first = __i;
144        continue;
145      }
146    }
147    ++__i;
148    // j points beyond range to be tested, *__lm1 is known to be <= *__m
149    // if not yet partitioned...
150    if (__i < __j) {
151      // known that *(__i - 1) < *__m
152      while (true) {
153        // __m still guards upward moving __i
154        while (__comp(*__i, *__m)) {
155          ++__i;
156          _LIBCPP_ASSERT_VALID_ELEMENT_ACCESS(
157              __i != __last,
158              "Would read out of bounds, does your comparator satisfy the strict-weak ordering requirement?");
159        }
160        // It is now known that a guard exists for downward moving __j
161        do {
162          _LIBCPP_ASSERT_VALID_ELEMENT_ACCESS(
163              __j != __first,
164              "Would read out of bounds, does your comparator satisfy the strict-weak ordering requirement?");
165          --__j;
166        } while (!__comp(*__j, *__m));
167        if (__i >= __j)
168          break;
169        _Ops::iter_swap(__i, __j);
170        ++__n_swaps;
171        // It is known that __m != __j
172        // If __m just moved, follow it
173        if (__m == __i)
174          __m = __j;
175        ++__i;
176      }
177    }
178    // [__first, __i) < *__m and *__m <= [__i, __last)
179    if (__i != __m && __comp(*__m, *__i)) {
180      _Ops::iter_swap(__i, __m);
181      ++__n_swaps;
182    }
183    // [__first, __i) < *__i and *__i <= [__i+1, __last)
184    if (__nth == __i)
185      return;
186    if (__n_swaps == 0) {
187      // We were given a perfectly partitioned sequence.  Coincidence?
188      if (__nth < __i) {
189        // Check for [__first, __i) already sorted
190        __j = __m = __first;
191        while (true) {
192          if (++__j == __i) {
193            // [__first, __i) sorted
194            return;
195          }
196          if (__comp(*__j, *__m)) {
197            // not yet sorted, so sort
198            break;
199          }
200          __m = __j;
201        }
202      } else {
203        // Check for [__i, __last) already sorted
204        __j = __m = __i;
205        while (true) {
206          if (++__j == __last) {
207            // [__i, __last) sorted
208            return;
209          }
210          if (__comp(*__j, *__m)) {
211            // not yet sorted, so sort
212            break;
213          }
214          __m = __j;
215        }
216      }
217    }
218    // __nth_element on range containing __nth
219    if (__nth < __i) {
220      // std::__nth_element<_Compare>(__first, __nth, __i, __comp);
221      __last = __i;
222    } else {
223      // std::__nth_element<_Compare>(__i+1, __nth, __last, __comp);
224      __first = ++__i;
225    }
226  }
227}
228
229template <class _AlgPolicy, class _RandomAccessIterator, class _Compare>
230inline _LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX20 void __nth_element_impl(
231    _RandomAccessIterator __first, _RandomAccessIterator __nth, _RandomAccessIterator __last, _Compare& __comp) {
232  if (__nth == __last)
233    return;
234
235  std::__debug_randomize_range<_AlgPolicy>(__first, __last);
236
237  std::__nth_element<_AlgPolicy, __comp_ref_type<_Compare> >(__first, __nth, __last, __comp);
238
239  std::__debug_randomize_range<_AlgPolicy>(__first, __nth);
240  if (__nth != __last) {
241    std::__debug_randomize_range<_AlgPolicy>(++__nth, __last);
242  }
243}
244
245template <class _RandomAccessIterator, class _Compare>
246inline _LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX20 void
247nth_element(_RandomAccessIterator __first, _RandomAccessIterator __nth, _RandomAccessIterator __last, _Compare __comp) {
248  std::__nth_element_impl<_ClassicAlgPolicy>(std::move(__first), std::move(__nth), std::move(__last), __comp);
249}
250
251template <class _RandomAccessIterator>
252inline _LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX20 void
253nth_element(_RandomAccessIterator __first, _RandomAccessIterator __nth, _RandomAccessIterator __last) {
254  std::nth_element(std::move(__first), std::move(__nth), std::move(__last), __less<>());
255}
256
257_LIBCPP_END_NAMESPACE_STD
258
259_LIBCPP_POP_MACROS
260
261#endif // _LIBCPP___ALGORITHM_NTH_ELEMENT_H
262