[libcxx-commits] [libcxx] [libcxx] Optimize std::generate for segmented iterators (PR #163006)

[via libcxx-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-libcxx

Author: Connector Switch (c8ef)

<details>
<summary>Changes</summary>

Part of #<!-- -->102817.

---
Full diff: https://github.com/llvm/llvm-project/pull/163006.diff


2 Files Affected:

- (modified) libcxx/include/__algorithm/generate.h (+25-2) 
- (modified) libcxx/test/std/algorithms/alg.modifying.operations/alg.generate/generate.pass.cpp (+11) 


``````````diff
diff --git a/libcxx/include/__algorithm/generate.h b/libcxx/include/__algorithm/generate.h
index c95b527402f5d..91e2ada7daf77 100644
--- a/libcxx/include/__algorithm/generate.h
+++ b/libcxx/include/__algorithm/generate.h
@@ -9,7 +9,10 @@
 #ifndef _LIBCPP___ALGORITHM_GENERATE_H
 #define _LIBCPP___ALGORITHM_GENERATE_H
 
+#include <__algorithm/for_each_segment.h>
 #include <__config>
+#include <__iterator/segmented_iterator.h>
+#include <__type_traits/enable_if.h>
 
 #if !defined(_LIBCPP_HAS_NO_PRAGMA_SYSTEM_HEADER)
 #  pragma GCC system_header
@@ -17,13 +20,33 @@
 
 _LIBCPP_BEGIN_NAMESPACE_STD
 
-template <class _ForwardIterator, class _Generator>
+template <class _ForwardIterator, class _Sent, class _Generator>
 inline _LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX20 void
-generate(_ForwardIterator __first, _ForwardIterator __last, _Generator __gen) {
+__generate(_ForwardIterator __first, _Sent __last, _Generator __gen) {
   for (; __first != __last; ++__first)
     *__first = __gen();
 }
 
+#ifndef _LIBCPP_CXX03_LANG
+template <class _SegmentedIterator,
+          class _Generator,
+          __enable_if_t<__is_segmented_iterator_v<_SegmentedIterator>, int> = 0>
+_LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX20
+_SegmentedIterator __generate(_SegmentedIterator __first, _SegmentedIterator __last, _Generator& __gen) {
+  using __local_iterator_t = typename __segmented_iterator_traits<_SegmentedIterator>::__local_iterator;
+  std::__for_each_segment(__first, __last, [&](__local_iterator_t __lfirst, __local_iterator_t __llast) {
+    std::__generate(__lfirst, __llast, __gen);
+  });
+  return __last;
+}
+#endif // !_LIBCPP_CXX03_LANG
+
+template <class _ForwardIterator, class _Generator>
+inline _LIBCPP_HIDE_FROM_ABI
+_LIBCPP_CONSTEXPR_SINCE_CXX20 void generate(_ForwardIterator __first, _ForwardIterator __last, _Generator __gen) {
+  std::__generate(__first, __last, __gen);
+}
+
 _LIBCPP_END_NAMESPACE_STD
 
 #endif // _LIBCPP___ALGORITHM_GENERATE_H
diff --git a/libcxx/test/std/algorithms/alg.modifying.operations/alg.generate/generate.pass.cpp b/libcxx/test/std/algorithms/alg.modifying.operations/alg.generate/generate.pass.cpp
index 29d32d7156742..4591d7ece4645 100644
--- a/libcxx/test/std/algorithms/alg.modifying.operations/alg.generate/generate.pass.cpp
+++ b/libcxx/test/std/algorithms/alg.modifying.operations/alg.generate/generate.pass.cpp
@@ -16,6 +16,7 @@
 
 #include <algorithm>
 #include <cassert>
+#include <deque>
 
 #include "test_macros.h"
 #include "test_iterators.h"
@@ -51,12 +52,22 @@ test()
     assert(ia[3] == 1);
 }
 
+void deque_test() {
+  int sizes[] = {0, 1, 2, 1023, 1024, 1025, 2047, 2048, 2049};
+  for (const int size : sizes) {
+    std::deque<int> d(size);
+    std::generate(d.begin(), d.end(), gen_test());
+    assert(std::all_of(d.begin(), d.end(), [](int x) { return x == 1; }));
+  }
+}
+
 int main(int, char**)
 {
     test<forward_iterator<int*> >();
     test<bidirectional_iterator<int*> >();
     test<random_access_iterator<int*> >();
     test<int*>();
+    deque_test();
 
 #if TEST_STD_VER > 17
     static_assert(test_constexpr());

``````````

</details>


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