[libc++] Optimize ranges::for_each for iterating over __trees

libcxx/include/CMakeLists.txt

@@ -194,6 +194,7 @@ set(files
   __algorithm/simd_utils.h
   __algorithm/sort.h
   __algorithm/sort_heap.h
+  __algorithm/specialized_algorithms.h
   __algorithm/stable_partition.h
   __algorithm/stable_sort.h
   __algorithm/swap_ranges.h

libcxx/include/__algorithm/for_each.h

@@ -11,6 +11,7 @@
 #define _LIBCPP___ALGORITHM_FOR_EACH_H
 
 #include <__algorithm/for_each_segment.h>
+#include <__algorithm/specialized_algorithms.h>
 #include <__config>
 #include <__functional/identity.h>
 #include <__iterator/segmented_iterator.h>
@@ -44,6 +45,19 @@ __for_each(_SegmentedIterator __first, _SegmentedIterator __last, _Func& __func,
   });
   return __last;
 }
+
+template <class _InputIterator,
+          class _Func,
+          class _Proj,
+          __enable_if_t<__specialized_algorithm<_Algorithm::__for_each,
+                                                __iterator_pair<_InputIterator, _InputIterator>>::__has_algorithm,
+                        int> = 0>
+_LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX20 _InputIterator
+__for_each(_InputIterator __first, _InputIterator __last, _Func& __func, _Proj& __proj) {
+  __specialized_algorithm<_Algorithm::__for_each, __iterator_pair<_InputIterator, _InputIterator>>()(
+      __first, __last, __func, __proj);
+  return __last;
+}
 #endif // !_LIBCPP_CXX03_LANG
 
 template <class _InputIterator, class _Func>

libcxx/include/__algorithm/ranges_for_each.h

@@ -12,6 +12,7 @@
 #include <__algorithm/for_each.h>
 #include <__algorithm/for_each_n.h>
 #include <__algorithm/in_fun_result.h>
+#include <__algorithm/specialized_algorithms.h>
 #include <__concepts/assignable.h>
 #include <__config>
 #include <__functional/identity.h>
@@ -20,6 +21,7 @@
 #include <__ranges/access.h>
 #include <__ranges/concepts.h>
 #include <__ranges/dangling.h>
+#include <__type_traits/remove_cvref.h>
 #include <__utility/move.h>
 
 #if !defined(_LIBCPP_HAS_NO_PRAGMA_SYSTEM_HEADER)
@@ -71,7 +73,13 @@ struct __for_each {
             indirectly_unary_invocable<projected<iterator_t<_Range>, _Proj>> _Func>
   _LIBCPP_HIDE_FROM_ABI constexpr for_each_result<borrowed_iterator_t<_Range>, _Func>
   operator()(_Range&& __range, _Func __func, _Proj __proj = {}) const {
-    return __for_each_impl(ranges::begin(__range), ranges::end(__range), __func, __proj);
+    using _SpecialAlg = __specialized_algorithm<_Algorithm::__for_each, remove_cvref_t<_Range>>;
+    if constexpr (_SpecialAlg::__has_algorithm) {
+      auto [__iter, __func2] = _SpecialAlg()(__range, std::move(__func), std::move(__proj));
+      return {std::move(__iter), std::move(__func)};
+    } else {
+      return __for_each_impl(ranges::begin(__range), ranges::end(__range), __func, __proj);
+    }
   }
 };
 

libcxx/include/__algorithm/specialized_algorithms.h

@@ -0,0 +1,35 @@
+//===----------------------------------------------------------------------===//
+//
+// 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_SPECIALIZED_ALGORITHMS_H
+#define _LIBCPP___ALGORITHM_SPECIALIZED_ALGORITHMS_H
+
+#include <__config>
+
+#if !defined(_LIBCPP_HAS_NO_PRAGMA_SYSTEM_HEADER)
+#  pragma GCC system_header
+#endif
+
+_LIBCPP_BEGIN_NAMESPACE_STD
+
+// FIXME: This should really be an enum
+namespace _Algorithm {
+  struct __for_each {};
+} // namespace _Algorithm
+
+template <class, class>
+struct __iterator_pair {};
+
+template <class _Alg, class _Range>
+struct __specialized_algorithm {
+  static const bool __has_algorithm = false;
+};
+
+_LIBCPP_END_NAMESPACE_STD
+
+#endif // _LIBCPP___ALGORITHM_SPECIALIZED_ALGORITHMS_H

libcxx/include/__tree

@@ -11,6 +11,7 @@
 #define _LIBCPP___TREE
 
 #include <__algorithm/min.h>
+#include <__algorithm/specialized_algorithms.h>
 #include <__assert>
 #include <__config>
 #include <__fwd/pair.h>
@@ -717,6 +718,59 @@ private:
   friend class __tree_const_iterator;
 };
 
+template <class _Reference, class _EndNodePtr, class _NodePtr, class _Func, class _Proj>
+_LIBCPP_HIDE_FROM_ABI bool __tree_iterate_from_root(_EndNodePtr __last, _NodePtr __root, _Func& __func, _Proj& __proj) {
+  if (__root->__left_) {
+    if (std::__tree_iterate_from_root<_Reference>(__last, static_cast<_NodePtr>(__root->__left_), __func, __proj))
+      return true;
+  }
+  if (__root == __last)
+    return true;
+  __func(static_cast<_Reference>(__root->__get_value()));
+  if (__root->__right_)
+    return std::__tree_iterate_from_root<_Reference>(__last, static_cast<_NodePtr>(__root->__right_), __func, __proj);
+  return false;
+}
+
+template <class _Reference, class _NodePtr, class _EndNodePtr, class _Func, class _Proj>
+_LIBCPP_HIDE_FROM_ABI void
+__tree_iterate_from_begin(_EndNodePtr __first, _EndNodePtr __last, _Func& __func, _Proj& __proj) {
+  while (true) {
+    if (__first == __last)
+      return;
+    auto __nfirst = static_cast<_NodePtr>(__first);
+    __func(static_cast<_Reference>(__nfirst->__get_value()));
+    if (__nfirst->__right_) {
+      if (std::__tree_iterate_from_root<_Reference>(__last, static_cast<_NodePtr>(__nfirst->__right_), __func, __proj))
+        return;
+    }
+    if (std::__tree_is_left_child(__nfirst)) {
+      __first = __nfirst->__parent_;
+    } else {
+      do {
+        __first = __nfirst->__parent_;
+      } while (!std::__tree_is_left_child(__nfirst));
+    }
+  }
+}
+
+#ifndef _LIBCPP_CXX03_LANG
+template <class _Tp, class _NodePtr, class _DiffType>
+struct __specialized_algorithm<
+    _Algorithm::__for_each,
+    __iterator_pair<__tree_iterator<_Tp, _NodePtr, _DiffType>, __tree_iterator<_Tp, _NodePtr, _DiffType>>> {
+  static const bool __has_algorithm = true;
+
+  using __iterator _LIBCPP_NODEBUG = __tree_iterator<_Tp, _NodePtr, _DiffType>;
+
+  template <class _Func, class _Proj>
+  _LIBCPP_HIDE_FROM_ABI static void operator()(__iterator __first, __iterator __last, _Func& __func, _Proj& __proj) {
+    std::__tree_iterate_from_begin<typename __iterator::reference, _NodePtr>(
+        __first.__ptr_, __last.__ptr_, __func, __proj);
+  }
+};
+#endif
+
 template <class _Tp, class _NodePtr, class _DiffType>
 class __tree_const_iterator {
   using _NodeTypes _LIBCPP_NODEBUG = __tree_node_types<_NodePtr>;
@@ -780,8 +834,28 @@ private:
 
   template <class, class, class>
   friend class __tree;
+
+  friend struct __specialized_algorithm<_Algorithm::__for_each,
+                                        __iterator_pair<__tree_const_iterator, __tree_const_iterator> >;
 };
 
+#ifndef _LIBCPP_CXX03_LANG
+template <class _Tp, class _NodePtr, class _DiffType>
+struct __specialized_algorithm<
+    _Algorithm::__for_each,
+    __iterator_pair<__tree_const_iterator<_Tp, _NodePtr, _DiffType>, __tree_const_iterator<_Tp, _NodePtr, _DiffType>>> {
+  static const bool __has_algorithm = true;
+
+  using __iterator = __tree_const_iterator<_Tp, _NodePtr, _DiffType>;
+
+  template <class _Func, class _Proj>
+  _LIBCPP_HIDE_FROM_ABI static void operator()(__iterator __first, __iterator __last, _Func& __func, _Proj& __proj) {
+    std::__tree_iterate_from_begin<typename __iterator::reference, _NodePtr>(
+        __first.__ptr_, __last.__ptr_, __func, __proj);
+  }
+};
+#endif
+
 template <class _Tp, class _Compare>
 #ifndef _LIBCPP_CXX03_LANG
 _LIBCPP_DIAGNOSE_WARNING(!__is_invocable_v<_Compare const&, _Tp const&, _Tp const&>,
@@ -1466,7 +1540,36 @@ private:
 
     return __dest;
   }
+
+  friend struct __specialized_algorithm<_Algorithm::__for_each, __tree>;
+};
+
+#if _LIBCPP_STD_VER >= 14
+template <class _Tp, class _Compare, class _Allocator>
+struct __specialized_algorithm<_Algorithm::__for_each, __tree<_Tp, _Compare, _Allocator> > {
+  static const bool __has_algorithm = true;
+
+  using __node_pointer _LIBCPP_NODEBUG = typename __tree<_Tp, _Compare, _Allocator>::__node_pointer;
+
+  template <class _Func, class _Proj>
+#ifndef _LIBCPP_COMPILER_GCC
+  _LIBCPP_HIDE_FROM_ABI
+#endif
+  static void __impl(__node_pointer __root, _Func& __func, _Proj& __proj) {
+    if (__root->__left_)
+      __impl(static_cast<__node_pointer>(__root->__left_), __func, __proj);
+    __func(__root->__get_value());
+    if (__root->__right_)
+      __impl(static_cast<__node_pointer>(__root->__right_), __func, __proj);
+  }
+
+  template <class _Tree, class _Func, class _Proj>
+  _LIBCPP_HIDE_FROM_ABI static auto operator()(_Tree&& __range, _Func __func, _Proj __proj) {
+    __impl(__range.__root(), __func, __proj);
+    return std::make_pair(__range.end(), std::move(__func));
+  }
 };
+#endif
 
 // Precondition:  __size_ != 0
 template <class _Tp, class _Compare, class _Allocator>

libcxx/include/map

@@ -577,6 +577,7 @@ erase_if(multimap<Key, T, Compare, Allocator>& c, Predicate pred);  // C++20
 #  include <__algorithm/equal.h>
 #  include <__algorithm/lexicographical_compare.h>
 #  include <__algorithm/lexicographical_compare_three_way.h>
+#  include <__algorithm/specialized_algorithms.h>
 #  include <__assert>
 #  include <__config>
 #  include <__functional/binary_function.h>
@@ -1375,6 +1376,8 @@ private:
 #  ifdef _LIBCPP_CXX03_LANG
   _LIBCPP_HIDE_FROM_ABI __node_holder __construct_node_with_key(const key_type& __k);
 #  endif
+
+  friend struct __specialized_algorithm<_Algorithm::__for_each, map>;
 };
 
 #  if _LIBCPP_STD_VER >= 17
@@ -1427,6 +1430,23 @@ map(initializer_list<pair<_Key, _Tp>>, _Allocator)
     -> map<remove_const_t<_Key>, _Tp, less<remove_const_t<_Key>>, _Allocator>;
 #  endif
 
+#  if _LIBCPP_STD_VER >= 14
+template <class _Key, class _Tp, class _Compare, class _Allocator>
+struct __specialized_algorithm<_Algorithm::__for_each, map<_Key, _Tp, _Compare, _Allocator>> {
+  using __map _LIBCPP_NODEBUG = map<_Key, _Tp, _Compare, _Allocator>;
+
+  static const bool __has_algorithm = true;
+
+  // set's begin() and end() are identical with and without const qualifiaction
+  template <class _Map, class _Func>
+  _LIBCPP_HIDE_FROM_ABI static auto operator()(_Map&& __map, _Func __func) {
+    auto [_, __func2] = __specialized_algorithm<_Algorithm::__for_each, typename __map::__base>()(
+        __map.__tree_, std::move(__func));
+    return std::make_pair(__map.end(), std::move(__func2));
+  }
+};
+#  endif
+
 #  ifndef _LIBCPP_CXX03_LANG
 template <class _Key, class _Tp, class _Compare, class _Allocator>
 map<_Key, _Tp, _Compare, _Allocator>::map(map&& __m, const allocator_type& __a)
@@ -1940,6 +1960,8 @@ private:
 
   typedef __map_node_destructor<__node_allocator> _Dp;
   typedef unique_ptr<__node, _Dp> __node_holder;
+
+  friend struct __specialized_algorithm<_Algorithm::__for_each, multimap>;
 };
 
 #  if _LIBCPP_STD_VER >= 17
@@ -1992,6 +2014,23 @@ multimap(initializer_list<pair<_Key, _Tp>>, _Allocator)
     -> multimap<remove_const_t<_Key>, _Tp, less<remove_const_t<_Key>>, _Allocator>;
 #  endif
 
+#  if _LIBCPP_STD_VER >= 14
+template <class _Key, class _Tp, class _Compare, class _Allocator>
+struct __specialized_algorithm<_Algorithm::__for_each, multimap<_Key, _Tp, _Compare, _Allocator>> {
+  using __map _LIBCPP_NODEBUG = multimap<_Key, _Tp, _Compare, _Allocator>;
+
+  static const bool __has_algorithm = true;
+
+  // set's begin() and end() are identical with and without const qualifiaction
+  template <class _Map, class _Func>
+  _LIBCPP_HIDE_FROM_ABI static auto operator()(_Map&& __map, _Func __func) {
+    auto [_, __func2] = __specialized_algorithm<_Algorithm::__for_each, typename __map::__base>()(
+        __map.__tree_, std::move(__func));
+    return std::make_pair(__map.end(), std::move(__func2));
+  }
+};
+#  endif
+
 #  ifndef _LIBCPP_CXX03_LANG
 template <class _Key, class _Tp, class _Compare, class _Allocator>
 multimap<_Key, _Tp, _Compare, _Allocator>::multimap(multimap&& __m, const allocator_type& __a)

libcxx/include/module.modulemap.in

@@ -838,6 +838,7 @@ module std [system] {
     module simd_utils                             { header "__algorithm/simd_utils.h" }
     module sort_heap                              { header "__algorithm/sort_heap.h" }
     module sort                                   { header "__algorithm/sort.h" }
+    module specialized_algorithms                 { header "__algorithm/specialized_algorithms.h" }
     module stable_partition                       { header "__algorithm/stable_partition.h" }
     module stable_sort {
       header "__algorithm/stable_sort.h"

libcxx/include/set

@@ -518,6 +518,7 @@ erase_if(multiset<Key, Compare, Allocator>& c, Predicate pred);  // C++20
 #  include <__algorithm/equal.h>
 #  include <__algorithm/lexicographical_compare.h>
 #  include <__algorithm/lexicographical_compare_three_way.h>
+#  include <__algorithm/specialized_algorithms.h>
 #  include <__assert>
 #  include <__config>
 #  include <__functional/is_transparent.h>
@@ -902,6 +903,9 @@ public:
     return __tree_.__equal_range_multi(__k);
   }
 #  endif
+
+  template <class, class>
+  friend struct __specialized_algorithm;
 };
 
 #  if _LIBCPP_STD_VER >= 17
@@ -948,6 +952,21 @@ template <class _Key, class _Allocator, class = enable_if_t<__is_allocator_v<_Al
 set(initializer_list<_Key>, _Allocator) -> set<_Key, less<_Key>, _Allocator>;
 #  endif
 
+#  if _LIBCPP_STD_VER >= 14
+template <class _Alg, class _Key, class _Compare, class _Allocator>
+struct __specialized_algorithm<_Alg, set<_Key, _Compare, _Allocator>> {
+  using __set _LIBCPP_NODEBUG = set<_Key, _Compare, _Allocator>;
+
+  static const bool __has_algorithm = __specialized_algorithm<_Alg, typename __set::__base>::__has_algorithm;
+
+  // set's begin() and end() are identical with and without const qualifiaction
+  template <class... _Args>
+  _LIBCPP_HIDE_FROM_ABI static auto operator()(const __set& __set, _Args&&... __args) {
+    return __specialized_algorithm<_Alg, typename __set::__base>()(__set.__tree_, std::forward<_Args>(__args)...);
+  }
+};
+#  endif
+
 #  ifndef _LIBCPP_CXX03_LANG
 
 template <class _Key, class _Compare, class _Allocator>
@@ -1362,6 +1381,9 @@ public:
     return __tree_.__equal_range_multi(__k);
   }
 #  endif
+
+  template <class, class>
+  friend struct __specialized_algorithm;
 };
 
 #  if _LIBCPP_STD_VER >= 17
@@ -1409,6 +1431,21 @@ template <class _Key, class _Allocator, class = enable_if_t<__is_allocator_v<_Al
 multiset(initializer_list<_Key>, _Allocator) -> multiset<_Key, less<_Key>, _Allocator>;
 #  endif
 
+#  if _LIBCPP_STD_VER >= 14
+template <class _Alg, class _Key, class _Compare, class _Allocator>
+struct __specialized_algorithm<_Alg, multiset<_Key, _Compare, _Allocator>> {
+  using __set _LIBCPP_NODEBUG = multiset<_Key, _Compare, _Allocator>;
+
+  static const bool __has_algorithm = __specialized_algorithm<_Alg, typename __set::__base>::__has_algorithm;
+
+  // set's begin() and end() are identical with and without const qualifiaction
+  template <class... _Args>
+  _LIBCPP_HIDE_FROM_ABI static auto operator()(const __set& __set, _Args&&... __args) {
+    return __specialized_algorithm<_Alg, typename __set::__base>()(__set.__tree_, std::forward<_Args>(__args)...);
+  }
+};
+#  endif
+
 #  ifndef _LIBCPP_CXX03_LANG
 
 template <class _Key, class _Compare, class _Allocator>