[libcxx-commits] [libcxx] [libcxx] Added segment iterator for fill (PR #104680)

[via libcxx-commits]

https://github.com/NoumanAmir657 created https://github.com/llvm/llvm-project/pull/104680

For the issue #102817.
@philnik777 

>From 2782a4cf1b62a7db53d7cb6a3543db0259f4ef5e Mon Sep 17 00:00:00 2001
From: nouman-10x <noumanamir453 at gmail.com>
Date: Sat, 17 Aug 2024 23:18:52 +0500
Subject: [PATCH] [libcxx] Added segment iterator for fill

---
 libcxx/include/__algorithm/fill.h | 53 +++++++++++++++++++++++++++----
 1 file changed, 47 insertions(+), 6 deletions(-)

diff --git a/libcxx/include/__algorithm/fill.h b/libcxx/include/__algorithm/fill.h
index 1ce3eadb013d05..3ba37e7260a4a2 100644
--- a/libcxx/include/__algorithm/fill.h
+++ b/libcxx/include/__algorithm/fill.h
@@ -21,25 +21,66 @@ _LIBCPP_BEGIN_NAMESPACE_STD
 
 // fill isn't specialized for std::memset, because the compiler already optimizes the loop to a call to std::memset.
 
-template <class _ForwardIterator, class _Tp>
+template <
+    class _ForwardIterator,
+    class _Tp,
+    __enable_if_t<
+        is_same<typename iterator_traits<_ForwardIterator>::iterator_category, forward_iterator_tag>::value ||
+            is_same<typename iterator_traits<_ForwardIterator>::iterator_category, bidirectional_iterator_tag>::value,
+        int> = 0>
 inline _LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX20 void
-__fill(_ForwardIterator __first, _ForwardIterator __last, const _Tp& __value, forward_iterator_tag) {
+__fill(_ForwardIterator __first, _ForwardIterator __last, const _Tp& __value) {
   for (; __first != __last; ++__first)
     *__first = __value;
 }
 
-template <class _RandomAccessIterator, class _Tp>
+template <class _RandomAccessIterator,
+          class _Tp,
+          __enable_if_t<(is_same<typename iterator_traits<_RandomAccessIterator>::iterator_category,
+                                 random_access_iterator_tag>::value ||
+                         is_same<typename iterator_traits<_RandomAccessIterator>::iterator_category,
+                                 contiguous_iterator_tag>::value) &&
+                            !__is_segmented_iterator<_RandomAccessIterator>::value,
+                        int> = 0>
 inline _LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX20 void
-__fill(_RandomAccessIterator __first, _RandomAccessIterator __last, const _Tp& __value, random_access_iterator_tag) {
+__fill(_RandomAccessIterator __first, _RandomAccessIterator __last, const _Tp& __value) {
   std::fill_n(__first, __last - __first, __value);
 }
 
+template <class _SegmentedIterator,
+          class _Tp,
+          __enable_if_t<__is_segmented_iterator<_SegmentedIterator>::value, int> = 0>
+inline _LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX20 void
+__fill(_SegmentedIterator __first, _SegmentedIterator __last, const _Tp& __value) {
+  using _Traits = __segmented_iterator_traits<_SegmentedIterator>;
+
+  auto __sfirst = _Traits::__segment(__first);
+  auto __slast  = _Traits::__segment(__last);
+
+  // We are in a single segment, so we might not be at the beginning or end
+  if (__sfirst == __slast) {
+    __fill(_Traits::__local(__first), _Traits::__local(__last), __value);
+    return;
+  }
+
+  // We have more than one segment. Iterate over the first segment, since we might not start at the beginning
+  __fill(_Traits::__local(__first), _Traits::__end(__sfirst), __value);
+  ++__sfirst;
+  // iterate over the segments which are guaranteed to be completely in the range
+  while (__sfirst != __slast) {
+    __fill(_Traits::__begin(__sfirst), _Traits::__end(__sfirst), __value);
+    ++__sfirst;
+  }
+  // iterate over the last segment
+  __fill(_Traits::__begin(__sfirst), _Traits::__local(__last), __value);
+}
+
 template <class _ForwardIterator, class _Tp>
 inline _LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX20 void
 fill(_ForwardIterator __first, _ForwardIterator __last, const _Tp& __value) {
-  std::__fill(__first, __last, __value, typename iterator_traits<_ForwardIterator>::iterator_category());
+  std::__fill(__first, __last, __value);
 }
 
 _LIBCPP_END_NAMESPACE_STD
 
-#endif // _LIBCPP___ALGORITHM_FILL_H
+#endif // _LIBCPP___ALGORITHM_FILL_H
\ No newline at end of file