[libc++] Optimize {std,ranges}::for_each for iterating over __trees (#164405)

This patch optimizes how `for_each` iterates over trees by using
recursion and storing pointers to the next nodes on the stack. This
avoids pointer chasing through the `__parent_` pointer, reducing cache
misses. It also makes use of the compiler being able tail-call optimize
the recursive function, removing back-tracking the iterators have to do.

```
Benchmark                                                      old             new    Difference    % Difference
--------------------------------------------------  --------------  --------------  ------------  --------------
rng::for_each(map<int>)/32                                   35.19           26.67         -8.52         -24.21%
rng::for_each(map<int>)/50                                   64.13           40.68        -23.45         -36.57%
rng::for_each(map<int>)/8                                     5.06            6.49          1.43          28.21%
rng::for_each(map<int>)/8192                              22893.89         9266.68     -13627.21         -59.52%
rng::for_each(map<int>::iterator)/32                         35.51           26.88         -8.63         -24.31%
rng::for_each(map<int>::iterator)/50                         64.39           41.24        -23.15         -35.95%
rng::for_each(map<int>::iterator)/8                           5.12            5.93          0.81          15.80%
rng::for_each(map<int>::iterator)/8192                    21283.14         9736.83     -11546.31         -54.25%
rng::for_each(multimap<int>)/32                              35.22           26.61         -8.61         -24.45%
rng::for_each(multimap<int>)/50                              64.10           40.07        -24.03         -37.49%
rng::for_each(multimap<int>)/8                                5.08            6.69          1.61          31.70%
rng::for_each(multimap<int>)/8192                         23130.44         9026.16     -14104.28         -60.98%
rng::for_each(multimap<int>::iterator)/32                    35.40           25.08        -10.32         -29.15%
rng::for_each(multimap<int>::iterator)/50                    64.19           38.15        -26.04         -40.56%
rng::for_each(multimap<int>::iterator)/8                      5.04            5.25          0.22           4.31%
rng::for_each(multimap<int>::iterator)/8192               22875.97         9392.08     -13483.89         -58.94%
rng::for_each(multiset<int>)/32                              35.82           27.11         -8.72         -24.33%
rng::for_each(multiset<int>)/50                              62.92           41.59        -21.32         -33.89%
rng::for_each(multiset<int>)/8                                4.79            6.79          2.00          41.70%
rng::for_each(multiset<int>)/8192                         22642.68         9280.95     -13361.73         -59.01%
rng::for_each(multiset<int>::iterator)/32                    35.76           26.71         -9.04         -25.28%
rng::for_each(multiset<int>::iterator)/50                    63.44           39.00        -24.44         -38.53%
rng::for_each(multiset<int>::iterator)/8                      4.90            5.21          0.30           6.18%
rng::for_each(multiset<int>::iterator)/8192               19930.45         9867.60     -10062.85         -50.49%
rng::for_each(set<int>)/32                                   35.90           27.30         -8.60         -23.96%
rng::for_each(set<int>)/50                                   63.15           40.75        -22.40         -35.47%
rng::for_each(set<int>)/8                                     4.77            6.83          2.06          43.23%
rng::for_each(set<int>)/8192                              20262.77         9381.57     -10881.20         -53.70%
rng::for_each(set<int>::iterator)/32                         36.02           26.42         -9.60         -26.64%
rng::for_each(set<int>::iterator)/50                         63.29           37.97        -25.32         -40.01%
rng::for_each(set<int>::iterator)/8                           4.72            5.22          0.50          10.50%
rng::for_each(set<int>::iterator)/8192                    20041.91         9831.91     -10210.00         -50.94%
```

Author: philnik777

libcxx/docs/ReleaseNotes/22.rst

@@ -86,6 +86,8 @@ Improvements and New Features
   This is guarded behind the ABI Macro ``_LIBCPP_ABI_ATOMIC_WAIT_NATIVE_BY_SIZE``.
 - The performance of ``vector<bool>::reserve()`` has been improved by up to 2x.
 
+- ``std::for_each`` and ``ranges::for_each`` have been optimized to iterate more efficiently over the associative
+  containers, resulting in performance improvements of up to 2x.
 - The ``num_get::do_get`` integral overloads have been optimized, resulting in a performance improvement of up to 2.8x.
 
 Deprecations and Removals

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>
@@ -27,7 +28,12 @@ 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& __func, _Proj& __proj) {
 #ifndef _LIBCPP_CXX03_LANG
-  if constexpr (is_same<_InputIterator, _Sent>::value && __is_segmented_iterator_v<_InputIterator>) {
+  if constexpr (using _SpecialAlg =
+                    __specialized_algorithm<_Algorithm::__for_each, __iterator_pair<_InputIterator, _Sent>>;
+                _SpecialAlg::__has_algorithm) {
+    _SpecialAlg()(__first, __last, __func, __proj);
+    return __last;
+  } else 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);

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, __single_range<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

@@ -19,11 +19,18 @@ _LIBCPP_BEGIN_NAMESPACE_STD
 
 namespace _Algorithm {
 struct __fill_n {};
+struct __for_each {};
 } // namespace _Algorithm
 
 template <class>
 struct __single_iterator;
 
+template <class, class>
+struct __iterator_pair;
+
+template <class>
+struct __single_range;
+
 // This struct allows specializing algorithms for specific arguments. This is useful when we know a more efficient
 // algorithm implementation for e.g. library-defined iterators. _Alg is one of tags defined inside the _Algorithm
 // namespace above. _Ranges is an essentially arbitrary subset of the arguments to the algorithm that are used for

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>
@@ -36,6 +37,7 @@
 #include <__type_traits/is_swappable.h>
 #include <__type_traits/make_transparent.h>
 #include <__type_traits/remove_const.h>
+#include <__type_traits/remove_cvref.h>
 #include <__utility/forward.h>
 #include <__utility/lazy_synth_three_way_comparator.h>
 #include <__utility/move.h>
@@ -656,6 +658,53 @@ struct __generic_container_node_destructor<__tree_node<_Tp, _VoidPtr>, _Alloc> :
 };
 #endif
 
+// Do an in-order traversal of the tree until `__break` returns true. Takes the root node of the tree.
+template <class _Reference, class _Break, class _NodePtr, class _Func, class _Proj>
+#ifndef _LIBCPP_COMPILER_GCC // This function is recursive, so GCC complains about always_inline.
+_LIBCPP_HIDE_FROM_ABI
+#endif
+bool __tree_iterate_from_root(_Break __break, _NodePtr __root, _Func& __func, _Proj& __proj) {
+  if (__root->__left_) {
+    if (std::__tree_iterate_from_root<_Reference>(__break, static_cast<_NodePtr>(__root->__left_), __func, __proj))
+      return true;
+  }
+  if (__break(__root))
+    return true;
+  __func(static_cast<_Reference>(__root->__get_value()));
+  if (__root->__right_)
+    return std::__tree_iterate_from_root<_Reference>(__break, static_cast<_NodePtr>(__root->__right_), __func, __proj);
+  return false;
+}
+
+// Do an in-order traversal of the tree from __first to __last.
+template <class _NodeIter, class _Func, class _Proj>
+_LIBCPP_HIDE_FROM_ABI void
+__tree_iterate_subrange(_NodeIter __first_it, _NodeIter __last_it, _Func& __func, _Proj& __proj) {
+  using _NodePtr   = typename _NodeIter::__node_pointer;
+  using _Reference = typename _NodeIter::reference;
+
+  auto __first = __first_it.__ptr_;
+  auto __last  = __last_it.__ptr_;
+
+  while (true) {
+    if (__first == __last)
+      return;
+    const auto __nfirst = static_cast<_NodePtr>(__first);
+    __func(static_cast<_Reference>(__nfirst->__get_value()));
+    if (__nfirst->__right_) {
+      if (std::__tree_iterate_from_root<_Reference>(
+              [&](_NodePtr __node) -> bool { return __node == __last; },
+              static_cast<_NodePtr>(__nfirst->__right_),
+              __func,
+              __proj))
+        return;
+    }
+    while (!std::__tree_is_left_child(static_cast<_NodePtr>(__first)))
+      __first = static_cast<_NodePtr>(__first)->__parent_;
+    __first = static_cast<_NodePtr>(__first)->__parent_;
+  }
+}
+
 template <class _Tp, class _NodePtr, class _DiffType>
 class __tree_iterator {
   using _NodeTypes _LIBCPP_NODEBUG = __tree_node_types<_NodePtr>;
@@ -715,7 +764,27 @@ private:
   friend class __tree;
   template <class, class, class>
   friend class __tree_const_iterator;
+
+  template <class _NodeIter, class _Func, class _Proj>
+  friend void __tree_iterate_subrange(_NodeIter, _NodeIter, _Func&, _Proj&);
+};
+
+#ifndef _LIBCPP_CXX03_LANG
+// This also handles {multi,}set::iterator, since they're just aliases to __tree::iterator
+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_subrange(__first, __last, __func, __proj);
+  }
 };
+#endif
 
 template <class _Tp, class _NodePtr, class _DiffType>
 class __tree_const_iterator {
@@ -780,8 +849,28 @@ private:
 
   template <class, class, class>
   friend class __tree;
+
+  template <class _NodeIter, class _Func, class _Proj>
+  friend void __tree_iterate_subrange(_NodeIter, _NodeIter, _Func&, _Proj&);
 };
 
+#ifndef _LIBCPP_CXX03_LANG
+// This also handles {multi,}set::const_iterator, since they're just aliases to __tree::iterator
+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 _LIBCPP_NODEBUG = __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_subrange(__first, __last, __func, __proj);
+  }
+};
+#endif
+
 template <class _Tp, class _Compare>
 #ifndef _LIBCPP_CXX03_LANG
 _LIBCPP_DIAGNOSE_WARNING(!__is_invocable_v<_Compare const&, _Tp const&, _Tp const&>,
@@ -1440,7 +1529,25 @@ private:
         [](value_type& __lhs, value_type& __rhs) { __assign_value(__lhs, std::move(__rhs)); },
         [this](__node_pointer __nd) { return __move_construct_tree(__nd); });
   }
+
+  friend struct __specialized_algorithm<_Algorithm::__for_each, __single_range<__tree> >;
+};
+
+#if _LIBCPP_STD_VER >= 14
+template <class _Tp, class _Compare, class _Allocator>
+struct __specialized_algorithm<_Algorithm::__for_each, __single_range<__tree<_Tp, _Compare, _Allocator> > > {
+  static const bool __has_algorithm = true;
+
+  using __node_pointer _LIBCPP_NODEBUG = typename __tree<_Tp, _Compare, _Allocator>::__node_pointer;
+
+  template <class _Tree, class _Func, class _Proj>
+  _LIBCPP_HIDE_FROM_ABI static auto operator()(_Tree&& __range, _Func __func, _Proj __proj) {
+    std::__tree_iterate_from_root<__copy_cvref_t<_Tree, typename __remove_cvref_t<_Tree>::value_type>>(
+        [](__node_pointer) { return false; }, __range.__root(), __func, __proj);
+    return std::make_pair(__range.end(), std::move(__func));
+  }
 };
+#endif
 
 template <class _Tp, class _Compare, class _Allocator>
 __tree<_Tp, _Compare, _Allocator>& __tree<_Tp, _Compare, _Allocator>::operator=(const __tree& __t) {

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>
@@ -818,7 +819,26 @@ public:
   friend class multimap;
   template <class>
   friend class __map_const_iterator;
+
+  template <class, class...>
+  friend struct __specialized_algorithm;
+};
+
+#  ifndef _LIBCPP_CXX03_LANG
+template <class _Alg, class _TreeIterator>
+struct __specialized_algorithm<_Alg, __iterator_pair<__map_iterator<_TreeIterator>, __map_iterator<_TreeIterator>>> {
+  using __base _LIBCPP_NODEBUG = __specialized_algorithm<_Alg, __iterator_pair<_TreeIterator, _TreeIterator>>;
+
+  static const bool __has_algorithm = __base::__has_algorithm;
+
+  using __iterator _LIBCPP_NODEBUG = __map_iterator<_TreeIterator>;
+
+  template <class... _Args>
+  _LIBCPP_HIDE_FROM_ABI static void operator()(__iterator __first, __iterator __last, _Args&&... __args) {
+    __base()(__first.__i_, __last.__i_, std::forward<_Args>(__args)...);
+  }
 };
+#  endif
 
 template <class _TreeIterator>
 class __map_const_iterator {
@@ -873,7 +893,28 @@ public:
   friend class multimap;
   template <class, class, class>
   friend class __tree_const_iterator;
+
+  template <class, class...>
+  friend struct __specialized_algorithm;
+};
+
+#  ifndef _LIBCPP_CXX03_LANG
+template <class _Alg, class _TreeIterator>
+struct __specialized_algorithm<
+    _Alg,
+    __iterator_pair<__map_const_iterator<_TreeIterator>, __map_const_iterator<_TreeIterator>>> {
+  using __base _LIBCPP_NODEBUG = __specialized_algorithm<_Alg, __iterator_pair<_TreeIterator, _TreeIterator>>;
+
+  static const bool __has_algorithm = __base::__has_algorithm;
+
+  using __iterator _LIBCPP_NODEBUG = __map_const_iterator<_TreeIterator>;
+
+  template <class... _Args>
+  _LIBCPP_HIDE_FROM_ABI static void operator()(__iterator __first, __iterator __last, _Args&&... __args) {
+    __base()(__first.__i_, __last.__i_, std::forward<_Args>(__args)...);
+  }
 };
+#  endif
 
 template <class _Key, class _Tp, class _Compare = less<_Key>, class _Allocator = allocator<pair<const _Key, _Tp> > >
 class multimap;
@@ -1371,6 +1412,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, __single_range<map> >;
 };
 
 #  if _LIBCPP_STD_VER >= 17
@@ -1423,6 +1466,22 @@ 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, __single_range<map<_Key, _Tp, _Compare, _Allocator>>> {
+  using __map _LIBCPP_NODEBUG = map<_Key, _Tp, _Compare, _Allocator>;
+
+  static const bool __has_algorithm = true;
+
+  template <class _Map, class _Func, class _Proj>
+  _LIBCPP_HIDE_FROM_ABI static auto operator()(_Map&& __map, _Func __func, _Proj __proj) {
+    auto [_, __func2] = __specialized_algorithm<_Algorithm::__for_each, __single_range<typename __map::__base>>()(
+        __map.__tree_, std::move(__func), std::move(__proj));
+    return std::make_pair(__map.end(), std::move(__func2));
+  }
+};
+#  endif
+
 #  ifndef _LIBCPP_CXX03_LANG
 template <class _Key, class _Tp, class _Compare, class _Allocator>
 _Tp& map<_Key, _Tp, _Compare, _Allocator>::operator[](const key_type& __k) {
@@ -1922,6 +1981,8 @@ private:
 
   typedef __map_node_destructor<__node_allocator> _Dp;
   typedef unique_ptr<__node, _Dp> __node_holder;
+
+  friend struct __specialized_algorithm<_Algorithm::__for_each, __single_range<multimap> >;
 };
 
 #  if _LIBCPP_STD_VER >= 17
@@ -1974,6 +2035,22 @@ 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, __single_range<multimap<_Key, _Tp, _Compare, _Allocator>>> {
+  using __map _LIBCPP_NODEBUG = multimap<_Key, _Tp, _Compare, _Allocator>;
+
+  static const bool __has_algorithm = true;
+
+  template <class _Map, class _Func, class _Proj>
+  _LIBCPP_HIDE_FROM_ABI static auto operator()(_Map&& __map, _Func __func, _Proj __proj) {
+    auto [_, __func2] = __specialized_algorithm<_Algorithm::__for_each, __single_range<typename __map::__base>>()(
+        __map.__tree_, std::move(__func), std::move(__proj));
+    return std::make_pair(__map.end(), std::move(__func2));
+  }
+};
+#  endif
+
 template <class _Key, class _Tp, class _Compare, class _Allocator>
 inline _LIBCPP_HIDE_FROM_ABI bool
 operator==(const multimap<_Key, _Tp, _Compare, _Allocator>& __x, const multimap<_Key, _Tp, _Compare, _Allocator>& __y) {