[libc++] Optimize __tree::__erase_unique

Author: philnik777

libcxx/include/__tree

@@ -766,6 +766,14 @@ public:
   typedef __rebind_alloc<__alloc_traits, __node> __node_allocator;
   typedef allocator_traits<__node_allocator> __node_traits;
 
+  enum class __uniqueness {
+    __unique,
+    __multi,
+  };
+
+  using __unique_tag = integral_constant<__uniqueness, __uniqueness::__unique>;
+  using __multi_tag  = integral_constant<__uniqueness, __uniqueness::__multi>;
+
 // TODO(LLVM 22): Remove this check
 #ifndef _LIBCPP_ABI_TREE_REMOVE_NODE_POINTER_UB
   static_assert(sizeof(__node_base_pointer) == sizeof(__end_node_pointer) && _LIBCPP_ALIGNOF(__node_base_pointer) ==
@@ -1009,9 +1017,27 @@ public:
   __insert_node_at(__end_node_pointer __parent, __node_base_pointer& __child, __node_base_pointer __new_node) _NOEXCEPT;
 
   template <class _Key>
-  _LIBCPP_HIDE_FROM_ABI iterator find(const _Key& __v);
+  _LIBCPP_HIDE_FROM_ABI iterator find(__multi_tag, const _Key& __v);
+  template <class _Key>
+  _LIBCPP_HIDE_FROM_ABI const_iterator find(__multi_tag, const _Key& __v) const;
+
+  template <class _Key>
+  _LIBCPP_HIDE_FROM_ABI iterator find(__unique_tag, const _Key& __key) {
+    __end_node_pointer __parent;
+    __node_base_pointer __match = __find_equal(__parent, __key);
+    if (__match == nullptr)
+      return end();
+    return iterator(static_cast<__node_pointer>(__match));
+  }
+
   template <class _Key>
-  _LIBCPP_HIDE_FROM_ABI const_iterator find(const _Key& __v) const;
+  _LIBCPP_HIDE_FROM_ABI const_iterator find(__unique_tag, const _Key& __key) const {
+    __end_node_pointer __parent;
+    __node_base_pointer __match = __find_equal(__parent, __key);
+    if (__match == nullptr)
+      return end();
+    return const_iterator(static_cast<__node_pointer>(__match));
+  }
 
   template <class _Key>
   _LIBCPP_HIDE_FROM_ABI size_type __count_unique(const _Key& __k) const;
@@ -1914,7 +1940,7 @@ __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);
+  iterator __it = find(__multi_tag(), __key);
   if (__it == end())
     return _NodeHandle();
   return __node_handle_extract<_NodeHandle>(__it);
@@ -2013,7 +2039,7 @@ template <class _Tp, class _Compare, class _Allocator>
 template <class _Key>
 typename __tree<_Tp, _Compare, _Allocator>::size_type
 __tree<_Tp, _Compare, _Allocator>::__erase_unique(const _Key& __k) {
-  iterator __i = find(__k);
+  iterator __i = find(__unique_tag(), __k);
   if (__i == end())
     return 0;
   erase(__i);
@@ -2033,7 +2059,8 @@ __tree<_Tp, _Compare, _Allocator>::__erase_multi(const _Key& __k) {
 
 template <class _Tp, class _Compare, class _Allocator>
 template <class _Key>
-typename __tree<_Tp, _Compare, _Allocator>::iterator __tree<_Tp, _Compare, _Allocator>::find(const _Key& __v) {
+typename __tree<_Tp, _Compare, _Allocator>::iterator
+__tree<_Tp, _Compare, _Allocator>::find(__multi_tag, const _Key& __v) {
   iterator __p = __lower_bound(__v, __root(), __end_node());
   if (__p != end() && !value_comp()(__v, *__p))
     return __p;
@@ -2043,7 +2070,7 @@ typename __tree<_Tp, _Compare, _Allocator>::iterator __tree<_Tp, _Compare, _Allo
 template <class _Tp, class _Compare, class _Allocator>
 template <class _Key>
 typename __tree<_Tp, _Compare, _Allocator>::const_iterator
-__tree<_Tp, _Compare, _Allocator>::find(const _Key& __v) const {
+__tree<_Tp, _Compare, _Allocator>::find(__multi_tag, const _Key& __v) const {
   const_iterator __p = __lower_bound(__v, __root(), __end_node());
   if (__p != end() && !value_comp()(__v, *__p))
     return __p;

libcxx/include/map

@@ -903,6 +903,8 @@ private:
 
   __base __tree_;
 
+  using __unique_tag = __base::__unique_tag;
+
 public:
   typedef typename __alloc_traits::pointer pointer;
   typedef typename __alloc_traits::const_pointer const_pointer;
@@ -1243,8 +1245,8 @@ public:
 
   _LIBCPP_HIDE_FROM_ABI void swap(map& __m) _NOEXCEPT_(__is_nothrow_swappable_v<__base>) { __tree_.swap(__m.__tree_); }
 
-  _LIBCPP_HIDE_FROM_ABI iterator find(const key_type& __k) { return __tree_.find(__k); }
-  _LIBCPP_HIDE_FROM_ABI const_iterator find(const key_type& __k) const { return __tree_.find(__k); }
+  _LIBCPP_HIDE_FROM_ABI iterator find(const key_type& __k) { return __tree_.find(__unique_tag(), __k); }
+  _LIBCPP_HIDE_FROM_ABI const_iterator find(const key_type& __k) const { return __tree_.find(__unique_tag(), __k); }
 #  if _LIBCPP_STD_VER >= 14
   template <typename _K2, enable_if_t<__is_transparent_v<_Compare, _K2>, int> = 0>
   _LIBCPP_HIDE_FROM_ABI iterator find(const _K2& __k) {

libcxx/include/set

@@ -595,6 +595,8 @@ private:
 
   __base __tree_;
 
+  using __unique_tag = __base::__unique_tag;
+
 public:
   typedef typename __base::pointer pointer;
   typedef typename __base::const_pointer const_pointer;
@@ -826,8 +828,8 @@ public:
   _LIBCPP_HIDE_FROM_ABI value_compare value_comp() const { return __tree_.value_comp(); }
 
   // set operations:
-  _LIBCPP_HIDE_FROM_ABI iterator find(const key_type& __k) { return __tree_.find(__k); }
-  _LIBCPP_HIDE_FROM_ABI const_iterator find(const key_type& __k) const { return __tree_.find(__k); }
+  _LIBCPP_HIDE_FROM_ABI iterator find(const key_type& __k) { return __tree_.find(__unique_tag(), __k); }
+  _LIBCPP_HIDE_FROM_ABI const_iterator find(const key_type& __k) const { return __tree_.find(__unique_tag(), __k); }
 #  if _LIBCPP_STD_VER >= 14
   template <typename _K2, enable_if_t<__is_transparent_v<_Compare, _K2>, int> = 0>
   _LIBCPP_HIDE_FROM_ABI iterator find(const _K2& __k) {