[libc++][ranges] Implement `ranges::sample`.

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

NOKEYCHECK=True
GitOrigin-RevId: 6bdb64223473585f783572c9fbf0673b4b324a35
diff --git a/include/CMakeLists.txt b/include/CMakeLists.txt
index 1193f68..297e902 100644
--- a/include/CMakeLists.txt
+++ b/include/CMakeLists.txt
@@ -133,6 +133,7 @@
   __algorithm/ranges_reverse.h
   __algorithm/ranges_reverse_copy.h
   __algorithm/ranges_rotate_copy.h
+  __algorithm/ranges_sample.h
   __algorithm/ranges_search.h
   __algorithm/ranges_search_n.h
   __algorithm/ranges_set_difference.h
@@ -178,6 +179,7 @@
   __algorithm/stable_sort.h
   __algorithm/swap_ranges.h
   __algorithm/transform.h
+  __algorithm/uniform_random_bit_generator_adaptor.h
   __algorithm/unique.h
   __algorithm/unique_copy.h
   __algorithm/unwrap_iter.h
diff --git a/include/__algorithm/iterator_operations.h b/include/__algorithm/iterator_operations.h
index e060574..2222b98 100644
--- a/include/__algorithm/iterator_operations.h
+++ b/include/__algorithm/iterator_operations.h
@@ -10,9 +10,11 @@
 #define _LIBCPP___ALGORITHM_ITERATOR_OPERATIONS_H
 
 #include <__algorithm/iter_swap.h>
+#include <__algorithm/ranges_iterator_concept.h>
 #include <__config>
 #include <__iterator/advance.h>
 #include <__iterator/distance.h>
+#include <__iterator/incrementable_traits.h>
 #include <__iterator/iter_move.h>
 #include <__iterator/iter_swap.h>
 #include <__iterator/iterator_traits.h>
@@ -40,6 +42,12 @@
   template <class _Iter>
   using __value_type = iter_value_t<_Iter>;
 
+  template <class _Iter>
+  using __iterator_category = ranges::__iterator_concept<_Iter>;
+
+  template <class _Iter>
+  using __difference_type = iter_difference_t<_Iter>;
+
   static constexpr auto advance = ranges::advance;
   static constexpr auto distance = ranges::distance;
   static constexpr auto __iter_move = ranges::iter_move;
@@ -58,6 +66,12 @@
   template <class _Iter>
   using __value_type = typename iterator_traits<_Iter>::value_type;
 
+  template <class _Iter>
+  using __iterator_category = typename iterator_traits<_Iter>::iterator_category;
+
+  template <class _Iter>
+  using __difference_type = typename iterator_traits<_Iter>::difference_type;
+
   // advance
   template <class _Iter, class _Distance>
   _LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_AFTER_CXX11
diff --git a/include/__algorithm/ranges_sample.h b/include/__algorithm/ranges_sample.h
new file mode 100644
index 0000000..a8477f8
--- /dev/null
+++ b/include/__algorithm/ranges_sample.h
@@ -0,0 +1,74 @@
+//===----------------------------------------------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef _LIBCPP___ALGORITHM_RANGES_SAMPLE_H
+#define _LIBCPP___ALGORITHM_RANGES_SAMPLE_H
+
+#include <__algorithm/iterator_operations.h>
+#include <__algorithm/sample.h>
+#include <__algorithm/uniform_random_bit_generator_adaptor.h>
+#include <__config>
+#include <__iterator/concepts.h>
+#include <__iterator/incrementable_traits.h>
+#include <__iterator/iterator_traits.h>
+#include <__random/uniform_random_bit_generator.h>
+#include <__ranges/access.h>
+#include <__ranges/concepts.h>
+#include <__utility/forward.h>
+#include <__utility/move.h>
+#include <type_traits>
+
+#if !defined(_LIBCPP_HAS_NO_PRAGMA_SYSTEM_HEADER)
+#  pragma GCC system_header
+#endif
+
+#if _LIBCPP_STD_VER > 17 && !defined(_LIBCPP_HAS_NO_INCOMPLETE_RANGES)
+
+_LIBCPP_BEGIN_NAMESPACE_STD
+
+namespace ranges {
+namespace __sample {
+
+struct __fn {
+
+  template <input_iterator _Iter, sentinel_for<_Iter> _Sent, weakly_incrementable _OutIter, class _Gen>
+  requires (forward_iterator<_Iter> || random_access_iterator<_OutIter>) &&
+           indirectly_copyable<_Iter, _OutIter> &&
+           uniform_random_bit_generator<remove_reference_t<_Gen>>
+  _LIBCPP_HIDE_FROM_ABI
+  _OutIter operator()(_Iter __first, _Sent __last,
+                      _OutIter __out_first, iter_difference_t<_Iter> __n, _Gen&& __gen) const {
+    _ClassicGenAdaptor<_Gen> __adapted_gen(__gen);
+    return std::__sample<_RangeAlgPolicy>(
+        std::move(__first), std::move(__last), std::move(__out_first), __n, __adapted_gen);
+  }
+
+  template <input_range _Range, weakly_incrementable _OutIter, class _Gen>
+  requires (forward_range<_Range> || random_access_iterator<_OutIter>) &&
+           indirectly_copyable<iterator_t<_Range>, _OutIter> &&
+           uniform_random_bit_generator<remove_reference_t<_Gen>>
+  _LIBCPP_HIDE_FROM_ABI
+  _OutIter operator()(_Range&& __range, _OutIter __out_first, range_difference_t<_Range> __n, _Gen&& __gen) const {
+    return (*this)(ranges::begin(__range), ranges::end(__range),
+                   std::move(__out_first), __n, std::forward<_Gen>(__gen));
+  }
+
+};
+
+} // namespace __sample
+
+inline namespace __cpo {
+  inline constexpr auto sample = __sample::__fn{};
+} // namespace __cpo
+} // namespace ranges
+
+_LIBCPP_END_NAMESPACE_STD
+
+#endif // _LIBCPP_STD_VER > 17 && !defined(_LIBCPP_HAS_NO_INCOMPLETE_RANGES)
+
+#endif // _LIBCPP___ALGORITHM_RANGES_SAMPLE_H
diff --git a/include/__algorithm/ranges_shuffle.h b/include/__algorithm/ranges_shuffle.h
index b101a85..9b7f81e 100644
--- a/include/__algorithm/ranges_shuffle.h
+++ b/include/__algorithm/ranges_shuffle.h
@@ -11,6 +11,7 @@
 
 #include <__algorithm/iterator_operations.h>
 #include <__algorithm/shuffle.h>
+#include <__algorithm/uniform_random_bit_generator_adaptor.h>
 #include <__config>
 #include <__functional/invoke.h>
 #include <__functional/ranges_operations.h>
@@ -32,43 +33,12 @@
 
 #if _LIBCPP_STD_VER > 17 && !defined(_LIBCPP_HAS_NO_INCOMPLETE_RANGES)
 
-_LIBCPP_PUSH_MACROS
-#include <__undef_macros>
-
 _LIBCPP_BEGIN_NAMESPACE_STD
 
 namespace ranges {
 namespace __shuffle {
 
 struct __fn {
-  // `std::shuffle` is more constrained than `std::ranges::shuffle`. `std::ranges::shuffle` only requires the given
-  // generator to satisfy the `std::uniform_random_bit_generator` concept. `std::shuffle` requires the given
-  // generator to meet the uniform random bit generator requirements; these requirements include satisfying
-  // `std::uniform_random_bit_generator` and add a requirement for the generator to provide a nested `result_type`
-  // typedef (see `[rand.req.urng]`).
-  //
-  // To reuse the implementation from `std::shuffle`, make the given generator meet the classic requirements by wrapping
-  // it into an adaptor type that forwards all of its interface and adds the required typedef.
-  template <class _Gen>
-  class _ClassicGenAdaptor {
-  private:
-    // The generator is not required to be copyable or movable, so it has to be stored as a reference.
-    _Gen& __gen;
-
-  public:
-    using result_type = invoke_result_t<_Gen&>;
-
-    _LIBCPP_HIDE_FROM_ABI
-    static constexpr auto min() { return __uncvref_t<_Gen>::min(); }
-    _LIBCPP_HIDE_FROM_ABI
-    static constexpr auto max() { return __uncvref_t<_Gen>::max(); }
-
-    _LIBCPP_HIDE_FROM_ABI
-    constexpr explicit _ClassicGenAdaptor(_Gen& __g) : __gen(__g) {}
-
-    _LIBCPP_HIDE_FROM_ABI
-    constexpr auto operator()() const { return __gen(); }
-  };
 
   template <random_access_iterator _Iter, sentinel_for<_Iter> _Sent, class _Gen>
   requires permutable<_Iter> && uniform_random_bit_generator<remove_reference_t<_Gen>>
@@ -96,8 +66,6 @@
 
 _LIBCPP_END_NAMESPACE_STD
 
-_LIBCPP_POP_MACROS
-
 #endif // _LIBCPP_STD_VER > 17 && !defined(_LIBCPP_HAS_NO_INCOMPLETE_RANGES)
 
 #endif // _LIBCPP___ALGORITHM_RANGES_SHUFFLE_H
diff --git a/include/__algorithm/sample.h b/include/__algorithm/sample.h
index e04466a..f403ba6 100644
--- a/include/__algorithm/sample.h
+++ b/include/__algorithm/sample.h
@@ -9,12 +9,14 @@
 #ifndef _LIBCPP___ALGORITHM_SAMPLE_H
 #define _LIBCPP___ALGORITHM_SAMPLE_H
 
+#include <__algorithm/iterator_operations.h>
 #include <__algorithm/min.h>
 #include <__assert>
 #include <__config>
 #include <__iterator/distance.h>
 #include <__iterator/iterator_traits.h>
 #include <__random/uniform_int_distribution.h>
+#include <__utility/move.h>
 #include <type_traits>
 
 #if !defined(_LIBCPP_HAS_NO_PRAGMA_SYSTEM_HEADER)
@@ -26,13 +28,14 @@
 
 _LIBCPP_BEGIN_NAMESPACE_STD
 
-template <class _PopulationIterator, class _SampleIterator, class _Distance,
+template <class _AlgPolicy,
+          class _PopulationIterator, class _PopulationSentinel, class _SampleIterator, class _Distance,
           class _UniformRandomNumberGenerator>
 _LIBCPP_INLINE_VISIBILITY
 _SampleIterator __sample(_PopulationIterator __first,
-                         _PopulationIterator __last, _SampleIterator __output_iter,
+                         _PopulationSentinel __last, _SampleIterator __output_iter,
                          _Distance __n,
-                         _UniformRandomNumberGenerator & __g,
+                         _UniformRandomNumberGenerator& __g,
                          input_iterator_tag) {
 
   _Distance __k = 0;
@@ -47,15 +50,16 @@
   return __output_iter + _VSTD::min(__n, __k);
 }
 
-template <class _PopulationIterator, class _SampleIterator, class _Distance,
+template <class _AlgPolicy,
+          class _PopulationIterator, class _PopulationSentinel, class _SampleIterator, class _Distance,
           class _UniformRandomNumberGenerator>
 _LIBCPP_INLINE_VISIBILITY
 _SampleIterator __sample(_PopulationIterator __first,
-                         _PopulationIterator __last, _SampleIterator __output_iter,
+                         _PopulationSentinel __last, _SampleIterator __output_iter,
                          _Distance __n,
                          _UniformRandomNumberGenerator& __g,
                          forward_iterator_tag) {
-  _Distance __unsampled_sz = _VSTD::distance(__first, __last);
+  _Distance __unsampled_sz = _IterOps<_AlgPolicy>::distance(__first, __last);
   for (__n = _VSTD::min(__n, __unsampled_sz); __n != 0; ++__first) {
     _Distance __r = uniform_int_distribution<_Distance>(0, --__unsampled_sz)(__g);
     if (__r < __n) {
@@ -66,24 +70,22 @@
   return __output_iter;
 }
 
-template <class _PopulationIterator, class _SampleIterator, class _Distance,
+template <class _AlgPolicy,
+          class _PopulationIterator, class _PopulationSentinel, class _SampleIterator, class _Distance,
           class _UniformRandomNumberGenerator>
 _LIBCPP_INLINE_VISIBILITY
 _SampleIterator __sample(_PopulationIterator __first,
-                         _PopulationIterator __last, _SampleIterator __output_iter,
+                         _PopulationSentinel __last, _SampleIterator __output_iter,
                          _Distance __n, _UniformRandomNumberGenerator& __g) {
-  typedef typename iterator_traits<_PopulationIterator>::iterator_category
-        _PopCategory;
-  typedef typename iterator_traits<_PopulationIterator>::difference_type
-        _Difference;
-  static_assert(__is_cpp17_forward_iterator<_PopulationIterator>::value ||
-                __is_cpp17_random_access_iterator<_SampleIterator>::value,
-                "SampleIterator must meet the requirements of RandomAccessIterator");
-  typedef typename common_type<_Distance, _Difference>::type _CommonType;
   _LIBCPP_ASSERT(__n >= 0, "N must be a positive number.");
-  return _VSTD::__sample(
-      __first, __last, __output_iter, _CommonType(__n),
-      __g, _PopCategory());
+
+  using _PopIterCategory = typename _IterOps<_AlgPolicy>::template __iterator_category<_PopulationIterator>;
+  using _Difference = typename _IterOps<_AlgPolicy>::template __difference_type<_PopulationIterator>;
+  using _CommonType = typename common_type<_Distance, _Difference>::type;
+
+  return std::__sample<_AlgPolicy>(
+      std::move(__first), std::move(__last), std::move(__output_iter), _CommonType(__n),
+      __g, _PopIterCategory());
 }
 
 #if _LIBCPP_STD_VER > 14
@@ -93,8 +95,14 @@
 _SampleIterator sample(_PopulationIterator __first,
                        _PopulationIterator __last, _SampleIterator __output_iter,
                        _Distance __n, _UniformRandomNumberGenerator&& __g) {
-    return _VSTD::__sample(__first, __last, __output_iter, __n, __g);
+  static_assert(__is_cpp17_forward_iterator<_PopulationIterator>::value ||
+                __is_cpp17_random_access_iterator<_SampleIterator>::value,
+                "SampleIterator must meet the requirements of RandomAccessIterator");
+
+  return std::__sample<_ClassicAlgPolicy>(
+      std::move(__first), std::move(__last), std::move(__output_iter), __n, __g);
 }
+
 #endif // _LIBCPP_STD_VER > 14
 
 _LIBCPP_END_NAMESPACE_STD
diff --git a/include/__algorithm/uniform_random_bit_generator_adaptor.h b/include/__algorithm/uniform_random_bit_generator_adaptor.h
new file mode 100644
index 0000000..04a6775
--- /dev/null
+++ b/include/__algorithm/uniform_random_bit_generator_adaptor.h
@@ -0,0 +1,62 @@
+//===----------------------------------------------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef _LIBCPP___ALGORITHM_RANGES_UNIFORM_RANDOM_BIT_GENERATOR_ADAPTOR_H
+#define _LIBCPP___ALGORITHM_RANGES_UNIFORM_RANDOM_BIT_GENERATOR_ADAPTOR_H
+
+#include <__config>
+#include <__functional/invoke.h>
+#include <type_traits>
+
+#if !defined(_LIBCPP_HAS_NO_PRAGMA_SYSTEM_HEADER)
+#  pragma GCC system_header
+#endif
+
+#if _LIBCPP_STD_VER > 17
+
+_LIBCPP_PUSH_MACROS
+#include <__undef_macros>
+
+_LIBCPP_BEGIN_NAMESPACE_STD
+
+// Range versions of random algorithms (e.g. `std::shuffle`) are less constrained than their classic counterparts.
+// Range algorithms only require the given generator to satisfy the `std::uniform_random_bit_generator` concept.
+// Classic algorithms require the given generator to meet the uniform random bit generator requirements; these
+// requirements include satisfying `std::uniform_random_bit_generator` and add a requirement for the generator to
+// provide a nested `result_type` typedef (see `[rand.req.urng]`).
+//
+// To be able to reuse classic implementations, make the given generator meet the classic requirements by wrapping
+// it into an adaptor type that forwards all of its interface and adds the required typedef.
+template <class _Gen>
+class _ClassicGenAdaptor {
+private:
+  // The generator is not required to be copyable or movable, so it has to be stored as a reference.
+  _Gen& __gen;
+
+public:
+  using result_type = invoke_result_t<_Gen&>;
+
+  _LIBCPP_HIDE_FROM_ABI
+  static constexpr auto min() { return __uncvref_t<_Gen>::min(); }
+  _LIBCPP_HIDE_FROM_ABI
+  static constexpr auto max() { return __uncvref_t<_Gen>::max(); }
+
+  _LIBCPP_HIDE_FROM_ABI
+  constexpr explicit _ClassicGenAdaptor(_Gen& __g) : __gen(__g) {}
+
+  _LIBCPP_HIDE_FROM_ABI
+  constexpr auto operator()() const { return __gen(); }
+};
+
+_LIBCPP_END_NAMESPACE_STD
+
+_LIBCPP_POP_MACROS
+
+#endif // _LIBCPP_STD_VER > 17
+
+#endif // _LIBCPP___ALGORITHM_RANGES_UNIFORM_RANDOM_BIT_GENERATOR_ADAPTOR_H
diff --git a/include/algorithm b/include/algorithm
index 0a40778..15ba340 100644
--- a/include/algorithm
+++ b/include/algorithm
@@ -758,6 +758,18 @@
     constexpr ranges::rotate_copy_result<borrowed_iterator_t<R>, O>
       ranges::rotate_copy(R&& r, iterator_t<R> middle, O result);                                   // since C++20
 
+  template<input_iterator I, sentinel_for<I> S, weakly_incrementable O, class Gen>
+    requires (forward_iterator<I> || random_access_iterator<O>) &&
+            indirectly_copyable<I, O> &&
+            uniform_random_bit_generator<remove_reference_t<Gen>>
+    O sample(I first, S last, O out, iter_difference_t<I> n, Gen&& g);                              // Since C++20
+
+  template<input_range R, weakly_incrementable O, class Gen>
+    requires (forward_range<R> || random_access_iterator<O>) &&
+            indirectly_copyable<iterator_t<R>, O> &&
+            uniform_random_bit_generator<remove_reference_t<Gen>>
+    O sample(R&& r, O out, range_difference_t<R> n, Gen&& g);                                       // Since C++20
+
   template<random_access_iterator I, sentinel_for<I> S, class Gen>
     requires permutable<I> &&
             uniform_random_bit_generator<remove_reference_t<Gen>>
@@ -1773,6 +1785,7 @@
 #include <__algorithm/ranges_reverse.h>
 #include <__algorithm/ranges_reverse_copy.h>
 #include <__algorithm/ranges_rotate_copy.h>
+#include <__algorithm/ranges_sample.h>
 #include <__algorithm/ranges_search.h>
 #include <__algorithm/ranges_search_n.h>
 #include <__algorithm/ranges_set_difference.h>
diff --git a/include/module.modulemap.in b/include/module.modulemap.in
index 8777db3..94da66c 100644
--- a/include/module.modulemap.in
+++ b/include/module.modulemap.in
@@ -372,6 +372,7 @@
       module ranges_reverse                  { private header "__algorithm/ranges_reverse.h" }
       module ranges_reverse_copy             { private header "__algorithm/ranges_reverse_copy.h" }
       module ranges_rotate_copy              { private header "__algorithm/ranges_rotate_copy.h" }
+      module ranges_sample                   { private header "__algorithm/ranges_sample.h" }
       module ranges_search                   { private header "__algorithm/ranges_search.h" }
       module ranges_search_n                 { private header "__algorithm/ranges_search_n.h" }
       module ranges_set_difference           { private header "__algorithm/ranges_set_difference.h" }
@@ -385,6 +386,9 @@
       module ranges_stable_sort              { private header "__algorithm/ranges_stable_sort.h" }
       module ranges_swap_ranges              { private header "__algorithm/ranges_swap_ranges.h" }
       module ranges_transform                { private header "__algorithm/ranges_transform.h" }
+      module uniform_random_bit_generator_adaptor {
+        private header "__algorithm/uniform_random_bit_generator_adaptor.h"
+      }
       module ranges_unique                   { private header "__algorithm/ranges_unique.h" }
       module ranges_unique_copy              { private header "__algorithm/ranges_unique_copy.h" }
       module ranges_upper_bound              { private header "__algorithm/ranges_upper_bound.h" }