[libc++] Simplify most of the segmented iterator optimizations

[philnik777]

This patch does two things.
(1) It replaces SFINAE with if constexpr, avoiding some overload resolution and unnecessary boilerplate.
(2) It removes an overload from __for_each_n to forward to __for_each, since __for_each doesn't provide any further optimizations.

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[llvmbot]

@llvm/pr-subscribers-libcxx

Author: Nikolas Klauser (philnik777)

Changes

This patch does two things.
(1) It replaces SFINAE with if constexpr, avoiding some overload resolution and unnecessary boilerplate.
(2) It removes an overload from __for_each_n to forward to __for_each, since __for_each doesn't provide any further optimizations.

---

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

6 Files Affected:

• (modified) libcxx/include/__algorithm/fill.h (+11-15)
• (modified) libcxx/include/__algorithm/fill_n.h (+11-30)
• (modified) libcxx/include/__algorithm/for_each.h (+12-18)
• (modified) libcxx/include/__algorithm/for_each_n.h (+13-43)
• (modified) libcxx/include/__iterator/distance.h (+15-20)
• (modified) libcxx/include/__iterator/segmented_iterator.h (-5)

diff --git a/libcxx/include/__algorithm/fill.h b/libcxx/include/__algorithm/fill.h
index 328ebb663376a..cc4a06d195bc4 100644
--- a/libcxx/include/__algorithm/fill.h
+++ b/libcxx/include/__algorithm/fill.h
@@ -15,6 +15,7 @@
 #include <__iterator/iterator_traits.h>
 #include <__iterator/segmented_iterator.h>
 #include <__type_traits/enable_if.h>
+#include <__type_traits/is_same.h>
 
 #if !defined(_LIBCPP_HAS_NO_PRAGMA_SYSTEM_HEADER)
 #  pragma GCC system_header
@@ -27,6 +28,15 @@ _LIBCPP_BEGIN_NAMESPACE_STD
 template <class _ForwardIterator, class _Sentinel, class _Tp>
 inline _LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX20 _ForwardIterator
 __fill(_ForwardIterator __first, _Sentinel __last, const _Tp& __value) {
+#ifndef _LIBCPP_CXX03_LANG
+  if constexpr (is_same<_ForwardIterator, _Sentinel>::value && __is_segmented_iterator_v<_ForwardIterator>) {
+    using __local_iterator_t = typename __segmented_iterator_traits<_ForwardIterator>::__local_iterator;
+    std::__for_each_segment(__first, __last, [&](__local_iterator_t __lfirst, __local_iterator_t __llast) {
+      std::__fill(__lfirst, __llast, __value);
+    });
+    return __last;
+  }
+#endif
   for (; __first != __last; ++__first)
     *__first = __value;
   return __first;
@@ -34,26 +44,12 @@ __fill(_ForwardIterator __first, _Sentinel __last, const _Tp& __value) {
 
 template <class _RandomAccessIterator,
           class _Tp,
-          __enable_if_t<__has_random_access_iterator_category<_RandomAccessIterator>::value &&
-                            !__is_segmented_iterator_v<_RandomAccessIterator>,
-                        int> = 0>
+          __enable_if_t<__has_random_access_iterator_category<_RandomAccessIterator>::value, int> = 0>
 inline _LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX20 _RandomAccessIterator
 __fill(_RandomAccessIterator __first, _RandomAccessIterator __last, const _Tp& __value) {
   return std::__fill_n(__first, __last - __first, __value);
 }
 
-#ifndef _LIBCPP_CXX03_LANG
-template <class _SegmentedIterator, class _Tp, __enable_if_t<__is_segmented_iterator_v<_SegmentedIterator>, int> = 0>
-inline _LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX20
-_SegmentedIterator __fill(_SegmentedIterator __first, _SegmentedIterator __last, const _Tp& __value) {
-  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::__fill(__lfirst, __llast, __value);
-  });
-  return __last;
-}
-#endif // !_LIBCPP_CXX03_LANG
-
 template <class _ForwardIterator, class _Tp>
 inline _LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX20 void
 fill(_ForwardIterator __first, _ForwardIterator __last, const _Tp& __value) {
diff --git a/libcxx/include/__algorithm/fill_n.h b/libcxx/include/__algorithm/fill_n.h
index 2bfacf3178c4e..426fe228bdabb 100644
--- a/libcxx/include/__algorithm/fill_n.h
+++ b/libcxx/include/__algorithm/fill_n.h
@@ -16,10 +16,6 @@
 #include <__iterator/iterator_traits.h>
 #include <__iterator/segmented_iterator.h>
 #include <__memory/pointer_traits.h>
-#include <__type_traits/conjunction.h>
-#include <__type_traits/enable_if.h>
-#include <__type_traits/integral_constant.h>
-#include <__type_traits/negation.h>
 #include <__utility/convert_to_integral.h>
 
 #if !defined(_LIBCPP_HAS_NO_PRAGMA_SYSTEM_HEADER)
@@ -33,39 +29,24 @@ _LIBCPP_BEGIN_NAMESPACE_STD
 
 // fill_n isn't specialized for std::memset, because the compiler already optimizes the loop to a call to std::memset.
 
-template <class _OutputIterator,
-          class _Size,
-          class _Tp
-#ifndef _LIBCPP_CXX03_LANG
-          ,
-          __enable_if_t<!_And<_BoolConstant<__is_segmented_iterator_v<_OutputIterator>>,
-                              __has_random_access_local_iterator<_OutputIterator>>::value,
-                        int> = 0
-#endif
-          >
+template <class _OutputIterator, class _Size, class _Tp>
 inline _LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX20 _OutputIterator
 __fill_n(_OutputIterator __first, _Size __n, const _Tp& __value) {
+#ifndef _LIBCPP_CXX03_LANG
+  if constexpr (__is_segmented_iterator_v<_OutputIterator>) {
+    using __local_iterator = typename __segmented_iterator_traits<_OutputIterator>::__local_iterator;
+    if constexpr (__has_random_access_iterator_category<__local_iterator>::value) {
+      return std::__for_each_n_segment(__first, __n, [&](__local_iterator __lfirst, __local_iterator __llast) {
+        std::__fill_n(__lfirst, __llast - __lfirst, __value);
+      });
+    }
+  }
+#endif
   for (; __n > 0; ++__first, (void)--__n)
     *__first = __value;
   return __first;
 }
 
-#ifndef _LIBCPP_CXX03_LANG
-template < class _OutputIterator,
-           class _Size,
-           class _Tp,
-           __enable_if_t<_And<_BoolConstant<__is_segmented_iterator_v<_OutputIterator>>,
-                              __has_random_access_local_iterator<_OutputIterator>>::value,
-                         int> = 0>
-inline _LIBCPP_HIDE_FROM_ABI
-_LIBCPP_CONSTEXPR_SINCE_CXX14 _OutputIterator __fill_n(_OutputIterator __first, _Size __n, const _Tp& __value) {
-  using __local_iterator_t = typename __segmented_iterator_traits<_OutputIterator>::__local_iterator;
-  return std::__for_each_n_segment(__first, __n, [&](__local_iterator_t __lfirst, __local_iterator_t __llast) {
-    std::__fill_n(__lfirst, __llast - __lfirst, __value);
-  });
-}
-#endif // !_LIBCPP_CXX03_LANG
-
 template <bool _FillVal, class _Cp>
 _LIBCPP_CONSTEXPR_SINCE_CXX20 _LIBCPP_HIDE_FROM_ABI void
 __fill_n_bool(__bit_iterator<_Cp, false> __first, typename __size_difference_type_traits<_Cp>::size_type __n) {
diff --git a/libcxx/include/__algorithm/for_each.h b/libcxx/include/__algorithm/for_each.h
index 6fb66d25a2462..cb26aa4d2656a 100644
--- a/libcxx/include/__algorithm/for_each.h
+++ b/libcxx/include/__algorithm/for_each.h
@@ -14,8 +14,8 @@
 #include <__config>
 #include <__functional/identity.h>
 #include <__iterator/segmented_iterator.h>
-#include <__type_traits/enable_if.h>
 #include <__type_traits/invoke.h>
+#include <__type_traits/is_same.h>
 
 #if !defined(_LIBCPP_HAS_NO_PRAGMA_SYSTEM_HEADER)
 #  pragma GCC system_header
@@ -25,27 +25,21 @@ _LIBCPP_BEGIN_NAMESPACE_STD
 
 template <class _InputIterator, class _Sent, class _Func, class _Proj>
 _LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX20 _InputIterator
-__for_each(_InputIterator __first, _Sent __last, _Func& __f, _Proj& __proj) {
+__for_each(_InputIterator __first, _Sent __last, _Func& __func, _Proj& __proj) {
+#ifndef _LIBCPP_CXX03_LANG
+  if constexpr (is_same<_InputIterator, _Sent>::value && __is_segmented_iterator_v<_InputIterator>) {
+    using __local_iterator_t = typename __segmented_iterator_traits<_InputIterator>::__local_iterator;
+    std::__for_each_segment(__first, __last, [&](__local_iterator_t __lfirst, __local_iterator_t __llast) {
+      std::__for_each(__lfirst, __llast, __func, __proj);
+    });
+    return __last;
+  }
+#endif
   for (; __first != __last; ++__first)
-    std::__invoke(__f, std::__invoke(__proj, *__first));
+    std::__invoke(__func, std::__invoke(__proj, *__first));
   return __first;
 }
 
-#ifndef _LIBCPP_CXX03_LANG
-template <class _SegmentedIterator,
-          class _Func,
-          class _Proj,
-          __enable_if_t<__is_segmented_iterator_v<_SegmentedIterator>, int> = 0>
-_LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX20 _SegmentedIterator
-__for_each(_SegmentedIterator __first, _SegmentedIterator __last, _Func& __func, _Proj& __proj) {
-  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::__for_each(__lfirst, __llast, __func, __proj);
-  });
-  return __last;
-}
-#endif // !_LIBCPP_CXX03_LANG
-
 template <class _InputIterator, class _Func>
 _LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX20 _Func
 for_each(_InputIterator __first, _InputIterator __last, _Func __f) {
diff --git a/libcxx/include/__algorithm/for_each_n.h b/libcxx/include/__algorithm/for_each_n.h
index 04650e15b6362..5c326fa825961 100644
--- a/libcxx/include/__algorithm/for_each_n.h
+++ b/libcxx/include/__algorithm/for_each_n.h
@@ -16,10 +16,7 @@
 #include <__functional/identity.h>
 #include <__iterator/iterator_traits.h>
 #include <__iterator/segmented_iterator.h>
-#include <__type_traits/disjunction.h>
-#include <__type_traits/enable_if.h>
 #include <__type_traits/invoke.h>
-#include <__type_traits/negation.h>
 #include <__utility/convert_to_integral.h>
 #include <__utility/move.h>
 
@@ -32,18 +29,23 @@ _LIBCPP_PUSH_MACROS
 
 _LIBCPP_BEGIN_NAMESPACE_STD
 
-template <class _InputIterator,
-          class _Size,
-          class _Func,
-          class _Proj,
-          __enable_if_t<!__has_random_access_iterator_category<_InputIterator>::value &&
-                            _Or<integral_constant<bool, !__is_segmented_iterator_v<_InputIterator> >,
-                                _Not<__has_random_access_local_iterator<_InputIterator> > >::value,
-                        int> = 0>
+template <class _InputIterator, class _Size, class _Func, class _Proj>
 _LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX20 _InputIterator
 __for_each_n(_InputIterator __first, _Size __orig_n, _Func& __f, _Proj& __proj) {
   typedef decltype(std::__convert_to_integral(__orig_n)) _IntegralSize;
   _IntegralSize __n = __orig_n;
+
+#ifndef _LIBCPP_CXX03_LANG
+  if constexpr (__is_segmented_iterator_v<_InputIterator>) {
+    using __local_iterator = typename __segmented_iterator_traits<_InputIterator>::__local_iterator;
+    if constexpr (__has_random_access_iterator_category<__local_iterator>::value) {
+      return std::__for_each_n_segment(__first, __orig_n, [&](__local_iterator __lfirst, __local_iterator __llast) {
+        std::__for_each(__lfirst, __llast, __f, __proj);
+      });
+    }
+  }
+#endif
+
   while (__n > 0) {
     std::__invoke(__f, std::__invoke(__proj, *__first));
     ++__first;
@@ -52,38 +54,6 @@ __for_each_n(_InputIterator __first, _Size __orig_n, _Func& __f, _Proj& __proj)
   return std::move(__first);
 }
 
-template <class _RandIter,
-          class _Size,
-          class _Func,
-          class _Proj,
-          __enable_if_t<__has_random_access_iterator_category<_RandIter>::value, int> = 0>
-_LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX20 _RandIter
-__for_each_n(_RandIter __first, _Size __orig_n, _Func& __f, _Proj& __proj) {
-  typename std::iterator_traits<_RandIter>::difference_type __n = __orig_n;
-  auto __last                                                   = __first + __n;
-  std::__for_each(__first, __last, __f, __proj);
-  return __last;
-}
-
-#ifndef _LIBCPP_CXX03_LANG
-template <class _SegmentedIterator,
-          class _Size,
-          class _Func,
-          class _Proj,
-          __enable_if_t<!__has_random_access_iterator_category<_SegmentedIterator>::value &&
-                            __is_segmented_iterator_v<_SegmentedIterator> &&
-                            __has_random_access_iterator_category<
-                                typename __segmented_iterator_traits<_SegmentedIterator>::__local_iterator>::value,
-                        int> = 0>
-_LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX20 _SegmentedIterator
-__for_each_n(_SegmentedIterator __first, _Size __orig_n, _Func& __f, _Proj& __proj) {
-  using __local_iterator_t = typename __segmented_iterator_traits<_SegmentedIterator>::__local_iterator;
-  return std::__for_each_n_segment(__first, __orig_n, [&](__local_iterator_t __lfirst, __local_iterator_t __llast) {
-    std::__for_each(__lfirst, __llast, __f, __proj);
-  });
-}
-#endif // !_LIBCPP_CXX03_LANG
-
 #if _LIBCPP_STD_VER >= 17
 
 template <class _InputIterator, class _Size, class _Func>
diff --git a/libcxx/include/__iterator/distance.h b/libcxx/include/__iterator/distance.h
index 9be9db0f0c70e..c68360af5f81a 100644
--- a/libcxx/include/__iterator/distance.h
+++ b/libcxx/include/__iterator/distance.h
@@ -11,6 +11,7 @@
 #define _LIBCPP___ITERATOR_DISTANCE_H
 
 #include <__algorithm/for_each_segment.h>
+#include <__concepts/same_as.h>
 #include <__config>
 #include <__iterator/concepts.h>
 #include <__iterator/incrementable_traits.h>
@@ -41,35 +42,29 @@ template <class _Iter>
 using __iter_distance_t _LIBCPP_NODEBUG = typename iterator_traits<_Iter>::difference_type;
 #endif
 
-template <class _InputIter, class _Sent>
-inline _LIBCPP_HIDE_FROM_ABI
-_LIBCPP_CONSTEXPR_SINCE_CXX17 __iter_distance_t<_InputIter> __distance(_InputIter __first, _Sent __last) {
-  __iter_distance_t<_InputIter> __r(0);
-  for (; __first != __last; ++__first)
-    ++__r;
-  return __r;
-}
-
 template <class _RandIter, __enable_if_t<__has_random_access_iterator_category<_RandIter>::value, int> = 0>
 inline _LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX17 __iter_distance_t<_RandIter>
 __distance(_RandIter __first, _RandIter __last) {
   return __last - __first;
 }
 
+template <class _InputIter, class _Sent>
+inline _LIBCPP_HIDE_FROM_ABI
+_LIBCPP_CONSTEXPR_SINCE_CXX17 __iter_distance_t<_InputIter> __distance(_InputIter __first, _Sent __last) {
+  __iter_distance_t<_InputIter> __r(0);
 #if _LIBCPP_STD_VER >= 20
-template <class _SegmentedIter,
-          __enable_if_t<!__has_random_access_iterator_category<_SegmentedIter>::value &&
-                            __is_segmented_iterator_v<_SegmentedIter>,
-                        int> = 0>
-inline _LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX17 __iter_distance_t<_SegmentedIter>
-__distance(_SegmentedIter __first, _SegmentedIter __last) {
-  __iter_distance_t<_SegmentedIter> __r(0);
-  std::__for_each_segment(__first, __last, [&__r](auto __lfirst, auto __llast) {
-    __r += std::__distance(__lfirst, __llast);
-  });
+  if constexpr (same_as<_InputIter, _Sent> && __is_segmented_iterator_v<_InputIter>) {
+    std::__for_each_segment(__first, __last, [&__r](auto __lfirst, auto __llast) {
+      __r += std::__distance(__lfirst, __llast);
+    });
+  } else
+#endif
+  {
+    for (; __first != __last; ++__first)
+      ++__r;
+  }
   return __r;
 }
-#endif // _LIBCPP_STD_VER >= 20
 
 template <class _InputIter>
 inline _LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX17 typename iterator_traits<_InputIter>::difference_type
diff --git a/libcxx/include/__iterator/segmented_iterator.h b/libcxx/include/__iterator/segmented_iterator.h
index 5df9737137101..dc56a740130b5 100644
--- a/libcxx/include/__iterator/segmented_iterator.h
+++ b/libcxx/include/__iterator/segmented_iterator.h
@@ -75,11 +75,6 @@ inline const bool __has_specialization_v<_Tp, sizeof(_Tp) * 0> = true;
 template <class _Iterator>
 inline const bool __is_segmented_iterator_v = __has_specialization_v<__segmented_iterator_traits<_Iterator> >;
 
-template <class _SegmentedIterator>
-struct __has_random_access_local_iterator
-    : __has_random_access_iterator_category<
-          typename __segmented_iterator_traits< _SegmentedIterator >::__local_iterator > {};
-
 _LIBCPP_END_NAMESPACE_STD
 
 #endif // _LIBCPP___SEGMENTED_ITERATOR_H

[ldionne]

Almost LGTM, modulo the functionality change.