[libcxx-commits] [libcxx] [libc++] Fix multi{map, set}::extract not returning the first matching element (PR #199703)

[via libcxx-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-libcxx

Author: Nikolas Klauser (philnik777)

<details>
<summary>Changes</summary>

According to [associative.reqmts] `extract(k)` returns the _first_ element in the container with key equivalent to k.


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


3 Files Affected:

- (modified) libcxx/include/__tree (+2-2) 
- (modified) libcxx/test/std/containers/associative/multimap/multimap.modifiers/extract_key.pass.cpp (+13) 
- (modified) libcxx/test/std/containers/associative/multiset/extract_key.pass.cpp (+13) 


``````````diff
diff --git a/libcxx/include/__tree b/libcxx/include/__tree
index bb8c4f3f8c623..5a81534156293 100644
--- a/libcxx/include/__tree
+++ b/libcxx/include/__tree
@@ -2049,8 +2049,8 @@ __tree<_Tp, _Compare, _Allocator>::__node_handle_insert_unique(const_iterator __
 template <class _Tp, class _Compare, class _Allocator>
 template <class _NodeHandle>
 _LIBCPP_HIDE_FROM_ABI _NodeHandle __tree<_Tp, _Compare, _Allocator>::__node_handle_extract(key_type const& __key) {
-  iterator __it = find(__key);
-  if (__it == end())
+  iterator __it = __lower_bound_multi(__key);
+  if (__it == end() || __value_comp_(__key, *__it))
     return _NodeHandle();
   return __node_handle_extract<_NodeHandle>(__it);
 }
diff --git a/libcxx/test/std/containers/associative/multimap/multimap.modifiers/extract_key.pass.cpp b/libcxx/test/std/containers/associative/multimap/multimap.modifiers/extract_key.pass.cpp
index b3027b3b0f727..42184edbdb8b6 100644
--- a/libcxx/test/std/containers/associative/multimap/multimap.modifiers/extract_key.pass.cpp
+++ b/libcxx/test/std/containers/associative/multimap/multimap.modifiers/extract_key.pass.cpp
@@ -50,6 +50,19 @@ int main(int, char**) {
     test(m, std::begin(keys), std::end(keys));
   }
 
+  { // Check that the first element is returned
+    std::multimap<int, int> m = {{1, 1}, {1, 2}, {1, 3}};
+    auto ptr                  = std::addressof(m.begin()->first);
+    auto res                  = m.extract(1);
+    assert(std::addressof(res.key()) == ptr);
+  }
+
+  { // Check that no element is returned if there is no match
+    std::multimap<int, int> m = {{1, 1}, {2, 2}, {3, 3}, {4, 4}, {5, 5}, {6, 6}};
+    auto res                  = m.extract(0);
+    assert(!res);
+  }
+
   {
     std::multimap<Counter<int>, Counter<int>> m = {{1, 1}, {2, 2}, {3, 3}, {4, 4}, {5, 5}, {6, 6}};
     {
diff --git a/libcxx/test/std/containers/associative/multiset/extract_key.pass.cpp b/libcxx/test/std/containers/associative/multiset/extract_key.pass.cpp
index 24d98cfaaf31a..ec75a6c5d835c 100644
--- a/libcxx/test/std/containers/associative/multiset/extract_key.pass.cpp
+++ b/libcxx/test/std/containers/associative/multiset/extract_key.pass.cpp
@@ -48,6 +48,19 @@ int main(int, char**) {
     test(m, std::begin(keys), std::end(keys));
   }
 
+  { // Check that the first element is returned
+    std::multiset<int> m = {1, 1, 1};
+    auto ptr             = std::addressof(*m.begin());
+    auto res             = m.extract(1);
+    assert(std::addressof(res.value()) == ptr);
+  }
+
+  { // Check that no element is returned if there is no match
+    std::multiset<int> m = {1, 2, 3};
+    auto res             = m.extract(0);
+    assert(!res);
+  }
+
   {
     std::multiset<Counter<int>> m = {1, 2, 3, 4, 5, 6};
     {

``````````

</details>


https://github.com/llvm/llvm-project/pull/199703