[libc++] Fix `ranges::for_each` taking whole associative containers

[frederick-vs-ja]

Currently, the version of ranges::for_each which takes whole associative containers treats __root->__get_value() as an existing element even when the container is empty.

We should exit earlier in __specialized_algorithm<...>::operator() when the tree is empty.

[llvmbot]

@llvm/pr-subscribers-libcxx

Author: A. Jiang (frederick-vs-ja)

Changes

Currently, the version of ranges::for_each which takes whole associative containers treats __root->__get_value() as an existing element even when the container is empty.

We should exit earlier in __specialized_algorithm<...>::operator() when the tree is empty.

---

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

2 Files Affected:

• (modified) libcxx/include/__tree (+3-2)
• (modified) libcxx/test/std/algorithms/alg.nonmodifying/alg.foreach/ranges.for_each.associative.pass.cpp (+27)

diff --git a/libcxx/include/__tree b/libcxx/include/__tree
index 0f6b2a84418b9..fbb48f8196964 100644
--- a/libcxx/include/__tree
+++ b/libcxx/include/__tree
@@ -1540,8 +1540,9 @@ struct __specialized_algorithm<_Algorithm::__for_each, __single_range<__tree<_Tp
 
   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);
+    if (__range.size() != 0)
+      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));
   }
 };
diff --git a/libcxx/test/std/algorithms/alg.nonmodifying/alg.foreach/ranges.for_each.associative.pass.cpp b/libcxx/test/std/algorithms/alg.nonmodifying/alg.foreach/ranges.for_each.associative.pass.cpp
index 637a2eda76529..120308852b994 100644
--- a/libcxx/test/std/algorithms/alg.nonmodifying/alg.foreach/ranges.for_each.associative.pass.cpp
+++ b/libcxx/test/std/algorithms/alg.nonmodifying/alg.foreach/ranges.for_each.associative.pass.cpp
@@ -37,6 +37,14 @@ void test_node_container(Converter conv) {
     });
     assert(invoke_count == 0);
   }
+  { // Check that an empty container works, taking the whole range
+    Container c;
+    int invoke_count = 0;
+    std::ranges::for_each(c, [&c, &invoke_count](const value_type& i) {
+      assert(&i == &*std::next(c.begin(), invoke_count++));
+    });
+    assert(invoke_count == 0);
+  }
   { // Check that a single-element container works
     Container c;
     c.insert(conv(0));
@@ -46,6 +54,15 @@ void test_node_container(Converter conv) {
     });
     assert(invoke_count == 1);
   }
+  { // Check that a single-element container works, taking the whole range
+    Container c;
+    c.insert(conv(0));
+    int invoke_count = 0;
+    std::ranges::for_each(c, [&c, &invoke_count](const value_type& i) {
+      assert(&i == &*std::next(c.begin(), invoke_count++));
+    });
+    assert(invoke_count == 1);
+  }
   { // Check that a two-element container works
     Container c;
     c.insert(conv(0));
@@ -56,6 +73,16 @@ void test_node_container(Converter conv) {
     });
     assert(invoke_count == 2);
   }
+  { // Check that a two-element container works, taking the whole range
+    Container c;
+    c.insert(conv(0));
+    c.insert(conv(1));
+    int invoke_count = 0;
+    std::ranges::for_each(c, [&c, &invoke_count](const value_type& i) {
+      assert(&i == &*std::next(c.begin(), invoke_count++));
+    });
+    assert(invoke_count == 2);
+  }
 
   Container c;
   for (int i = 0; i != 10; ++i)

[nico]

Thanks!

Have you considered making __tree_iterate_from_root check the root pointer? That allows removing a few other checks and looks like it might be less code (and you don't need the size() call). Something like this (untested):

diff --git a/include/__tree b/include/__tree
index 2b93ea6603..ec04b2c2e5 100644
--- a/include/__tree
+++ b/include/__tree
@@ -664,16 +664,14 @@ template <class _Reference, class _Break, class _NodePtr, class _Func, class _Pr
 _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 (!__root)
+    return false;
+  if (std::__tree_iterate_from_root<_Reference>(__break, static_cast<_NodePtr>(__root->__left_), __func, __proj))
+    return true;
   if (__break(__root))
     return true;
   std::__invoke(__func, std::__invoke(__proj, 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;
+  return std::__tree_iterate_from_root<_Reference>(__break, static_cast<_NodePtr>(__root->__right_), __func, __proj);
 }
 
 // Do an in-order traversal of the tree from __first to __last.
@@ -691,14 +689,12 @@ __tree_iterate_subrange(_NodeIter __first_it, _NodeIter __last_it, _Func& __func
       return;
     const auto __nfirst = static_cast<_NodePtr>(__first);
     std::__invoke(__func, std::__invoke(__proj, 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;
-    }
+    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_;

[nico]

I hit "merge" on this, given it's been approved and fixes a regression.

Thanks for the fix :)