[libc++] Introduce __for_each_segment and use it in copy/move

This simplifies the code inside copy/move and makes it easier to apply the optimization to other algorithms.

Reviewed By: ldionne, Mordante, #libc

Spies: arichardson, libcxx-commits

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

libcxx/include/CMakeLists.txt

@@ -24,6 +24,7 @@ set(files
   __algorithm/find_if_not.h
   __algorithm/for_each.h
   __algorithm/for_each_n.h
+  __algorithm/for_each_segment.h
   __algorithm/generate.h
   __algorithm/generate_n.h
   __algorithm/half_positive.h

libcxx/include/__algorithm/copy.h

@@ -10,6 +10,7 @@
 #define _LIBCPP___ALGORITHM_COPY_H
 
 #include <__algorithm/copy_move_common.h>
+#include <__algorithm/for_each_segment.h>
 #include <__algorithm/iterator_operations.h>
 #include <__algorithm/min.h>
 #include <__config>
@@ -44,26 +45,24 @@ struct __copy_loop {
     return std::make_pair(std::move(__first), std::move(__result));
   }
 
+  template <class _InIter, class _OutIter>
+  struct _CopySegment {
+    using _Traits = __segmented_iterator_traits<_InIter>;
+
+    _OutIter& __result_;
+
+    _LIBCPP_HIDE_FROM_ABI _CopySegment(_OutIter& __result) : __result_(__result) {}
+
+    _LIBCPP_HIDE_FROM_ABI void
+    operator()(typename _Traits::__local_iterator __lfirst, typename _Traits::__local_iterator __llast) {
+      __result_ = std::__copy<_AlgPolicy>(__lfirst, __llast, std::move(__result_)).second;
+    }
+  };
+
   template <class _InIter, class _OutIter, __enable_if_t<__is_segmented_iterator<_InIter>::value, int> = 0>
   _LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX14 pair<_InIter, _OutIter>
   operator()(_InIter __first, _InIter __last, _OutIter __result) const {
-    using _Traits = __segmented_iterator_traits<_InIter>;
-    auto __sfirst = _Traits::__segment(__first);
-    auto __slast  = _Traits::__segment(__last);
-    if (__sfirst == __slast) {
-      auto __iters = std::__copy<_AlgPolicy>(_Traits::__local(__first), _Traits::__local(__last), std::move(__result));
-      return std::make_pair(__last, std::move(__iters.second));
-    }
-
-    __result = std::__copy<_AlgPolicy>(_Traits::__local(__first), _Traits::__end(__sfirst), std::move(__result)).second;
-    ++__sfirst;
-    while (__sfirst != __slast) {
-      __result =
-          std::__copy<_AlgPolicy>(_Traits::__begin(__sfirst), _Traits::__end(__sfirst), std::move(__result)).second;
-      ++__sfirst;
-    }
-    __result =
-        std::__copy<_AlgPolicy>(_Traits::__begin(__sfirst), _Traits::__local(__last), std::move(__result)).second;
+    std::__for_each_segment(__first, __last, _CopySegment<_InIter, _OutIter>(__result));
     return std::make_pair(__last, std::move(__result));
   }
 
@@ -98,8 +97,7 @@ struct __copy_loop {
 
 struct __copy_trivial {
   // At this point, the iterators have been unwrapped so any `contiguous_iterator` has been unwrapped to a pointer.
-  template <class _In, class _Out,
-            __enable_if_t<__can_lower_copy_assignment_to_memmove<_In, _Out>::value, int> = 0>
+  template <class _In, class _Out, __enable_if_t<__can_lower_copy_assignment_to_memmove<_In, _Out>::value, int> = 0>
   _LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX14 pair<_In*, _Out*>
   operator()(_In* __first, _In* __last, _Out* __result) const {
     return std::__copy_trivial_impl(__first, __last, __result);

libcxx/include/__algorithm/for_each_segment.h

@@ -0,0 +1,53 @@
+//===----------------------------------------------------------------------===//
+//
+// 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_FOR_EACH_SEGMENT_H
+#define _LIBCPP___ALGORITHM_FOR_EACH_SEGMENT_H
+
+#include <__config>
+#include <__iterator/segmented_iterator.h>
+
+#if !defined(_LIBCPP_HAS_NO_PRAGMA_SYSTEM_HEADER)
+#  pragma GCC system_header
+#endif
+
+_LIBCPP_BEGIN_NAMESPACE_STD
+
+// __for_each_segment is a utility function for optimizing iterating over segmented iterators linearly.
+// __first and __last are expected to be a segmented range. __func is expected to take a range of local iterators.
+// Anything that is returned from __func is ignored.
+
+template <class _SegmentedIterator, class _Functor>
+_LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX14 void
+__for_each_segment(_SegmentedIterator __first, _SegmentedIterator __last, _Functor __func) {
+  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) {
+    __func(_Traits::__local(__first), _Traits::__local(__last));
+    return;
+  }
+
+  // We have more than one segment. Iterate over the first segment, since we might not start at the beginning
+  __func(_Traits::__local(__first), _Traits::__end(__sfirst));
+  ++__sfirst;
+  // iterate over the segments which are guaranteed to be completely in the range
+  while (__sfirst != __slast) {
+    __func(_Traits::__begin(__sfirst), _Traits::__end(__sfirst));
+    ++__sfirst;
+  }
+  // iterate over the last segment
+  __func(_Traits::__begin(__sfirst), _Traits::__local(__last));
+}
+
+_LIBCPP_END_NAMESPACE_STD
+
+#endif // _LIBCPP___ALGORITHM_FOR_EACH_SEGMENT_H

libcxx/include/__algorithm/move.h

@@ -10,6 +10,7 @@
 #define _LIBCPP___ALGORITHM_MOVE_H
 
 #include <__algorithm/copy_move_common.h>
+#include <__algorithm/for_each_segment.h>
 #include <__algorithm/iterator_operations.h>
 #include <__algorithm/min.h>
 #include <__config>
@@ -45,26 +46,24 @@ struct __move_loop {
     return std::make_pair(std::move(__first), std::move(__result));
   }
 
+  template <class _InIter, class _OutIter>
+  struct _MoveSegment {
+    using _Traits = __segmented_iterator_traits<_InIter>;
+
+    _OutIter& __result_;
+
+    _LIBCPP_HIDE_FROM_ABI _MoveSegment(_OutIter& __result) : __result_(__result) {}
+
+    _LIBCPP_HIDE_FROM_ABI void
+    operator()(typename _Traits::__local_iterator __lfirst, typename _Traits::__local_iterator __llast) {
+      __result_ = std::__move<_AlgPolicy>(__lfirst, __llast, std::move(__result_)).second;
+    }
+  };
+
   template <class _InIter, class _OutIter, __enable_if_t<__is_segmented_iterator<_InIter>::value, int> = 0>
   _LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX14 pair<_InIter, _OutIter>
   operator()(_InIter __first, _InIter __last, _OutIter __result) const {
-    using _Traits = __segmented_iterator_traits<_InIter>;
-    auto __sfirst = _Traits::__segment(__first);
-    auto __slast  = _Traits::__segment(__last);
-    if (__sfirst == __slast) {
-      auto __iters = std::__move<_AlgPolicy>(_Traits::__local(__first), _Traits::__local(__last), std::move(__result));
-      return std::make_pair(__last, std::move(__iters.second));
-    }
-
-    __result = std::__move<_AlgPolicy>(_Traits::__local(__first), _Traits::__end(__sfirst), std::move(__result)).second;
-    ++__sfirst;
-    while (__sfirst != __slast) {
-      __result =
-          std::__move<_AlgPolicy>(_Traits::__begin(__sfirst), _Traits::__end(__sfirst), std::move(__result)).second;
-      ++__sfirst;
-    }
-    __result =
-        std::__move<_AlgPolicy>(_Traits::__begin(__sfirst), _Traits::__local(__last), std::move(__result)).second;
+    std::__for_each_segment(__first, __last, _MoveSegment<_InIter, _OutIter>(__result));
     return std::make_pair(__last, std::move(__result));
   }
 

libcxx/include/module.modulemap.in

@@ -267,6 +267,7 @@ module std [system] {
       module find_if_not                     { private header "__algorithm/find_if_not.h" }
       module for_each                        { private header "__algorithm/for_each.h" }
       module for_each_n                      { private header "__algorithm/for_each_n.h" }
+      module for_each_segment                { private header "__algorithm/for_each_segment.h" }
       module generate                        { private header "__algorithm/generate.h" }
       module generate_n                      { private header "__algorithm/generate_n.h" }
       module half_positive                   { private header "__algorithm/half_positive.h" }

libcxx/test/libcxx/private_headers.verify.cpp

@@ -69,6 +69,7 @@ END-SCRIPT
 #include <__algorithm/find_if_not.h> // expected-error@*:* {{use of private header from outside its module: '__algorithm/find_if_not.h'}}
 #include <__algorithm/for_each.h> // expected-error@*:* {{use of private header from outside its module: '__algorithm/for_each.h'}}
 #include <__algorithm/for_each_n.h> // expected-error@*:* {{use of private header from outside its module: '__algorithm/for_each_n.h'}}
+#include <__algorithm/for_each_segment.h> // expected-error@*:* {{use of private header from outside its module: '__algorithm/for_each_segment.h'}}
 #include <__algorithm/generate.h> // expected-error@*:* {{use of private header from outside its module: '__algorithm/generate.h'}}
 #include <__algorithm/generate_n.h> // expected-error@*:* {{use of private header from outside its module: '__algorithm/generate_n.h'}}
 #include <__algorithm/half_positive.h> // expected-error@*:* {{use of private header from outside its module: '__algorithm/half_positive.h'}}

libcxx/utils/data/ignore_format.txt

@@ -40,7 +40,6 @@ libcxx/include/__algorithm/clamp.h
 libcxx/include/__algorithm/comp.h
 libcxx/include/__algorithm/comp_ref_type.h
 libcxx/include/__algorithm/copy_backward.h
-libcxx/include/__algorithm/copy.h
 libcxx/include/__algorithm/copy_if.h
 libcxx/include/__algorithm/copy_move_common.h
 libcxx/include/__algorithm/copy_n.h

libcxx/utils/libcxx/test/params.py

@@ -21,7 +21,6 @@ _warningFlags = [
     "-Wno-unused-command-line-argument",
     "-Wno-attributes",
     "-Wno-pessimizing-move",
-    "-Wno-c++11-extensions",
     "-Wno-noexcept-type",
     "-Wno-aligned-allocation-unavailable",
     "-Wno-atomic-alignment",
@@ -47,6 +46,10 @@ _warningFlags = [
     "-Wunused-parameter",
     "-Wunreachable-code",
     "-Wno-unused-local-typedef",
+
+    # Disable warnings for extensions used in C++03
+    "-Wno-local-type-template-args",
+    "-Wno-c++11-extensions",
 ]
 
 _allStandards = ["c++03", "c++11", "c++14", "c++17", "c++20", "c++23", "c++26"]