[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)