[libcxx-commits] [libcxx] [libc++] Refactor memory allocation in basic_string (PR #128423)

Nikolas Klauser via libcxx-commits libcxx-commits at lists.llvm.org
Sun Feb 23 07:27:45 PST 2025


https://github.com/philnik777 created https://github.com/llvm/llvm-project/pull/128423

None

>From e0ab1dc2ddf5307c22dae81e4b68c286e0617f5a Mon Sep 17 00:00:00 2001
From: Nikolas Klauser <nikolasklauser at berlin.de>
Date: Fri, 7 Feb 2025 10:56:08 +0100
Subject: [PATCH] [libc++] Refactor memory allocation in basic_string

---
 libcxx/include/__memory/allocate_at_least.h   |  15 +-
 libcxx/include/string                         | 265 +++++++-----------
 libcxx/src/string.cpp                         |  19 +-
 .../string.cons/copy_alloc.pass.cpp           |  17 +-
 4 files changed, 113 insertions(+), 203 deletions(-)

diff --git a/libcxx/include/__memory/allocate_at_least.h b/libcxx/include/__memory/allocate_at_least.h
index 9b5a8bcbd4596..eb9aa5d37fbd5 100644
--- a/libcxx/include/__memory/allocate_at_least.h
+++ b/libcxx/include/__memory/allocate_at_least.h
@@ -19,21 +19,22 @@
 
 _LIBCPP_BEGIN_NAMESPACE_STD
 
+template <class _Pointer>
+struct __allocation_result {
+  _Pointer ptr;
+  size_t count;
+};
+
 #if _LIBCPP_STD_VER >= 23
 
 template <class _Alloc>
 [[nodiscard]] _LIBCPP_HIDE_FROM_ABI constexpr auto __allocate_at_least(_Alloc& __alloc, size_t __n) {
-  return std::allocator_traits<_Alloc>::allocate_at_least(__alloc, __n);
+  auto __res = std::allocator_traits<_Alloc>::allocate_at_least(__alloc, __n);
+  return __allocation_result{__res.ptr, __res.count};
 }
 
 #else
 
-template <class _Pointer>
-struct __allocation_result {
-  _Pointer ptr;
-  size_t count;
-};
-
 template <class _Alloc>
 [[__nodiscard__]] _LIBCPP_HIDE_FROM_ABI
 _LIBCPP_CONSTEXPR __allocation_result<typename allocator_traits<_Alloc>::pointer>
diff --git a/libcxx/include/string b/libcxx/include/string
index 46bf13a5c900f..d0d9cd2b131e1 100644
--- a/libcxx/include/string
+++ b/libcxx/include/string
@@ -600,7 +600,6 @@ basic_string<char32_t> operator""s( const char32_t *str, size_t len );
 #  include <__functional/hash.h>
 #  include <__functional/unary_function.h>
 #  include <__fwd/string.h>
-#  include <__ios/fpos.h>
 #  include <__iterator/bounded_iter.h>
 #  include <__iterator/distance.h>
 #  include <__iterator/iterator_traits.h>
@@ -643,7 +642,6 @@ basic_string<char32_t> operator""s( const char32_t *str, size_t len );
 #  include <__utility/move.h>
 #  include <__utility/scope_guard.h>
 #  include <__utility/swap.h>
-#  include <__utility/unreachable.h>
 #  include <climits>
 #  include <cstdio> // EOF
 #  include <cstring>
@@ -814,6 +812,10 @@ private:
 #  ifdef _LIBCPP_ABI_ALTERNATE_STRING_LAYOUT
 
   struct __long {
+    __long() = default;
+    _LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX20 __long(size_type __cap, size_type __size, pointer __data)
+        : __data_(__data), __size_(__size), __cap_(__cap), __is_long_(true) {}
+
     pointer __data_;
     size_type __size_;
     size_type __cap_ : sizeof(size_type) * CHAR_BIT - 1;
@@ -861,6 +863,13 @@ private:
   // some platforms bit fields have a default size rather than the actual
   // size used, e.g., it is 4 bytes on AIX. See D128285 for details.
   struct __long {
+    __long() = default;
+
+    _LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX20 __long(__allocation_result<pointer> __alloc, size_type __size)
+        : __is_long_(true), __cap_(__alloc.count / __endian_factor), __size_(__size), __data_(__alloc.ptr) {
+      _LIBCPP_ASSERT_INTERNAL(!__fits_in_sso(__alloc.count), "Long capacity should always be larger than the SSO");
+    }
+
     struct _LIBCPP_PACKED {
       size_type __is_long_ : 1;
       size_type __cap_ : sizeof(size_type) * CHAR_BIT - 1;
@@ -906,20 +915,7 @@ private:
   _LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX20 explicit basic_string(
       __uninitialized_size_tag, size_type __size, const allocator_type& __a)
       : __alloc_(__a) {
-    if (__size > max_size())
-      this->__throw_length_error();
-    if (__fits_in_sso(__size)) {
-      __rep_ = __rep();
-      __set_short_size(__size);
-    } else {
-      auto __capacity   = __recommend(__size) + 1;
-      auto __allocation = __alloc_traits::allocate(__alloc_, __capacity);
-      __begin_lifetime(__allocation, __capacity);
-      __set_long_cap(__capacity);
-      __set_long_pointer(__allocation);
-      __set_long_size(__size);
-    }
-    __annotate_new(__size);
+    __init_internal_buffer(__size);
   }
 
   template <class _Iter, class _Sent>
@@ -1181,11 +1177,7 @@ public:
   }
 #  endif // _LIBCPP_CXX03_LANG
 
-  inline _LIBCPP_CONSTEXPR_SINCE_CXX20 ~basic_string() {
-    __annotate_delete();
-    if (__is_long())
-      __alloc_traits::deallocate(__alloc_, __get_long_pointer(), __get_long_cap());
-  }
+  inline _LIBCPP_CONSTEXPR_SINCE_CXX20 ~basic_string() { __reset_internal_buffer(); }
 
   _LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX20 operator __self_view() const _NOEXCEPT {
     return __self_view(typename __self_view::__assume_valid(), data(), size());
@@ -1995,18 +1987,6 @@ private:
     return __rep_.__s.__is_long_;
   }
 
-  static _LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX20 void __begin_lifetime(pointer __begin, size_type __n) {
-#  if _LIBCPP_STD_VER >= 20
-    if (__libcpp_is_constant_evaluated()) {
-      for (size_type __i = 0; __i != __n; ++__i)
-        std::construct_at(std::addressof(__begin[__i]));
-    }
-#  else
-    (void)__begin;
-    (void)__n;
-#  endif // _LIBCPP_STD_VER >= 20
-  }
-
   _LIBCPP_CONSTEXPR _LIBCPP_HIDE_FROM_ABI static bool __fits_in_sso(size_type __sz) { return __sz < __min_cap; }
 
   template <class _Iterator, class _Sentinel>
@@ -2094,21 +2074,11 @@ private:
       __set_short_size(__s);
   }
 
-  _LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX20 void __set_long_cap(size_type __s) _NOEXCEPT {
-    _LIBCPP_ASSERT_INTERNAL(!__fits_in_sso(__s), "Long capacity should always be larger than the SSO");
-    __rep_.__l.__cap_     = __s / __endian_factor;
-    __rep_.__l.__is_long_ = true;
-  }
-
   _LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX20 size_type __get_long_cap() const _NOEXCEPT {
     _LIBCPP_ASSERT_INTERNAL(__rep_.__l.__is_long_, "String has to be long when trying to get the long capacity");
     return __rep_.__l.__cap_ * __endian_factor;
   }
 
-  _LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX20 void __set_long_pointer(pointer __p) _NOEXCEPT {
-    __rep_.__l.__data_ = __p;
-  }
-
   _LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX20 pointer __get_long_pointer() _NOEXCEPT {
     _LIBCPP_ASSERT_INTERNAL(__rep_.__l.__is_long_, "String has to be long when trying to get the long pointer");
     return _LIBCPP_ASAN_VOLATILE_WRAPPER(__rep_.__l.__data_);
@@ -2136,6 +2106,58 @@ private:
     return __is_long() ? __get_long_pointer() : __get_short_pointer();
   }
 
+  _LIBCPP_HIDE_FROM_ABI static _LIBCPP_CONSTEXPR_SINCE_CXX20 __allocation_result<pointer>
+  __allocate_long_buffer(_Allocator& __alloc, size_type __capacity) {
+    auto __buffer = std::__allocate_at_least(__alloc, __recommend(__capacity) + 1);
+
+    if (__libcpp_is_constant_evaluated()) {
+      for (size_type __i = 0; __i != __buffer.count; ++__i)
+        std::__construct_at(std::addressof(__buffer.ptr[__i]));
+    }
+
+    return __buffer;
+  }
+
+  _LIBCPP_HIDE_FROM_ABI static _LIBCPP_CONSTEXPR_SINCE_CXX20 void
+  __deallocate_long_buffer(_Allocator& __alloc, __allocation_result<pointer> __allocation) {
+    __alloc_traits::deallocate(__alloc, __allocation.ptr, __allocation.count);
+  }
+
+  _LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX20 void __reset_internal_buffer() {
+    __annotate_delete();
+    if (__is_long())
+      __deallocate_long_buffer(__alloc_, __get_internal_long_buffer());
+    __rep_.__s = __short();
+  }
+
+  _LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX20 __allocation_result<pointer> __get_internal_long_buffer() {
+    _LIBCPP_ASSERT_INTERNAL(__is_long(), "Trying to get buffer which doesn't exist!");
+    return __allocation_result<pointer>{__get_long_pointer(), __get_long_cap()};
+  }
+
+  _LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX20 void
+  __replace_internal_buffer(__allocation_result<pointer> __alloc, size_type __size) {
+    __reset_internal_buffer();
+    __rep_.__l = __long(__alloc, __size);
+  }
+
+  _LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX20 pointer __init_internal_buffer(size_type __size) {
+    if (__libcpp_is_constant_evaluated())
+      __rep_ = __rep();
+
+    if (__size > max_size())
+      __throw_length_error();
+
+    if (__fits_in_sso(__size)) {
+      __set_short_size(__size);
+      __annotate_new(__size);
+      return __get_short_pointer();
+    }
+    __rep_.__l = __long(__allocate_long_buffer(__alloc_, __size), __size);
+    __annotate_new(__size);
+    return __get_long_pointer();
+  }
+
   // The following functions are no-ops outside of AddressSanitizer mode.
   _LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX20 void
   __annotate_contiguous_container(const void* __old_mid, const void* __new_mid) const {
@@ -2293,24 +2315,15 @@ private:
       __alloc_ = __str.__alloc_;
     else {
       if (!__str.__is_long()) {
-        if (__is_long()) {
-          __annotate_delete();
-          __alloc_traits::deallocate(__alloc_, __get_long_pointer(), capacity() + 1);
-          __rep_ = __rep();
-        }
+        __reset_internal_buffer();
         __alloc_ = __str.__alloc_;
       } else {
         __annotate_delete();
-        auto __guard       = std::__make_scope_guard(__annotate_new_size(*this));
-        allocator_type __a = __str.__alloc_;
-        auto __allocation  = std::__allocate_at_least(__a, __str.__get_long_cap());
-        __begin_lifetime(__allocation.ptr, __allocation.count);
-        if (__is_long())
-          __alloc_traits::deallocate(__alloc_, __get_long_pointer(), __get_long_cap());
-        __alloc_ = std::move(__a);
-        __set_long_pointer(__allocation.ptr);
-        __set_long_cap(__allocation.count);
-        __set_long_size(__str.size());
+        auto __guard      = std::__make_scope_guard(__annotate_new_size(*this));
+        auto __alloc      = __str.__alloc_;
+        auto __allocation = __allocate_long_buffer(__alloc, __str.size());
+        __replace_internal_buffer(__allocation, __str.size());
+        __alloc_   = std::move(__alloc);
       }
     }
   }
@@ -2452,73 +2465,23 @@ basic_string(from_range_t, _Range&&, _Allocator = _Allocator())
 template <class _CharT, class _Traits, class _Allocator>
 _LIBCPP_CONSTEXPR_SINCE_CXX20 void
 basic_string<_CharT, _Traits, _Allocator>::__init(const value_type* __s, size_type __sz) {
-  if (__libcpp_is_constant_evaluated())
-    __rep_ = __rep();
-  if (__sz > max_size())
-    this->__throw_length_error();
-  pointer __p;
-  if (__fits_in_sso(__sz)) {
-    __set_short_size(__sz);
-    __p = __get_short_pointer();
-  } else {
-    auto __allocation = std::__allocate_at_least(__alloc_, __recommend(__sz) + 1);
-    __p               = __allocation.ptr;
-    __begin_lifetime(__p, __allocation.count);
-    __set_long_pointer(__p);
-    __set_long_cap(__allocation.count);
-    __set_long_size(__sz);
-  }
+  pointer __p = __init_internal_buffer(__sz);
   traits_type::copy(std::__to_address(__p), __s, __sz);
   traits_type::assign(__p[__sz], value_type());
-  __annotate_new(__sz);
 }
 
 template <class _CharT, class _Traits, class _Allocator>
 _LIBCPP_CONSTEXPR_SINCE_CXX20 _LIBCPP_NOINLINE void
 basic_string<_CharT, _Traits, _Allocator>::__init_copy_ctor_external(const value_type* __s, size_type __sz) {
-  if (__libcpp_is_constant_evaluated())
-    __rep_ = __rep();
-
-  pointer __p;
-  if (__fits_in_sso(__sz)) {
-    __p = __get_short_pointer();
-    __set_short_size(__sz);
-  } else {
-    if (__sz > max_size())
-      this->__throw_length_error();
-    auto __allocation = std::__allocate_at_least(__alloc_, __recommend(__sz) + 1);
-    __p               = __allocation.ptr;
-    __begin_lifetime(__p, __allocation.count);
-    __set_long_pointer(__p);
-    __set_long_cap(__allocation.count);
-    __set_long_size(__sz);
-  }
+  pointer __p = __init_internal_buffer(__sz);
   traits_type::copy(std::__to_address(__p), __s, __sz + 1);
-  __annotate_new(__sz);
 }
 
 template <class _CharT, class _Traits, class _Allocator>
 _LIBCPP_CONSTEXPR_SINCE_CXX20 void basic_string<_CharT, _Traits, _Allocator>::__init(size_type __n, value_type __c) {
-  if (__libcpp_is_constant_evaluated())
-    __rep_ = __rep();
-
-  if (__n > max_size())
-    this->__throw_length_error();
-  pointer __p;
-  if (__fits_in_sso(__n)) {
-    __set_short_size(__n);
-    __p = __get_short_pointer();
-  } else {
-    auto __allocation = std::__allocate_at_least(__alloc_, __recommend(__n) + 1);
-    __p               = __allocation.ptr;
-    __begin_lifetime(__p, __allocation.count);
-    __set_long_pointer(__p);
-    __set_long_cap(__allocation.count);
-    __set_long_size(__n);
-  }
+  pointer __p = __init_internal_buffer(__n);
   traits_type::assign(std::__to_address(__p), __n, __c);
   traits_type::assign(__p[__n], value_type());
-  __annotate_new(__n);
 }
 
 template <class _CharT, class _Traits, class _Allocator>
@@ -2542,9 +2505,7 @@ basic_string<_CharT, _Traits, _Allocator>::__init_with_sentinel(_InputIterator _
       push_back(*__first);
 #  if _LIBCPP_HAS_EXCEPTIONS
   } catch (...) {
-    __annotate_delete();
-    if (__is_long())
-      __alloc_traits::deallocate(__alloc_, __get_long_pointer(), __get_long_cap());
+    __reset_internal_buffer();
     throw;
   }
 #  endif // _LIBCPP_HAS_EXCEPTIONS
@@ -2562,25 +2523,7 @@ template <class _CharT, class _Traits, class _Allocator>
 template <class _InputIterator, class _Sentinel>
 _LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX20 void
 basic_string<_CharT, _Traits, _Allocator>::__init_with_size(_InputIterator __first, _Sentinel __last, size_type __sz) {
-  if (__libcpp_is_constant_evaluated())
-    __rep_ = __rep();
-
-  if (__sz > max_size())
-    this->__throw_length_error();
-
-  pointer __p;
-  if (__fits_in_sso(__sz)) {
-    __set_short_size(__sz);
-    __p = __get_short_pointer();
-
-  } else {
-    auto __allocation = std::__allocate_at_least(__alloc_, __recommend(__sz) + 1);
-    __p               = __allocation.ptr;
-    __begin_lifetime(__p, __allocation.count);
-    __set_long_pointer(__p);
-    __set_long_cap(__allocation.count);
-    __set_long_size(__sz);
-  }
+  pointer __p = __init_internal_buffer(__sz);
 
 #  if _LIBCPP_HAS_EXCEPTIONS
   try {
@@ -2589,12 +2532,10 @@ basic_string<_CharT, _Traits, _Allocator>::__init_with_size(_InputIterator __fir
     traits_type::assign(*__end, value_type());
 #  if _LIBCPP_HAS_EXCEPTIONS
   } catch (...) {
-    if (__is_long())
-      __alloc_traits::deallocate(__alloc_, __get_long_pointer(), __get_long_cap());
+    __reset_internal_buffer();
     throw;
   }
 #  endif                       // _LIBCPP_HAS_EXCEPTIONS
-  __annotate_new(__sz);
 }
 
 template <class _CharT, class _Traits, class _Allocator>
@@ -2614,9 +2555,8 @@ _LIBCPP_CONSTEXPR_SINCE_CXX20 void basic_string<_CharT, _Traits, _Allocator>::__
       __old_cap < __ms / 2 - __alignment ? __recommend(std::max(__old_cap + __delta_cap, 2 * __old_cap)) : __ms - 1;
   __annotate_delete();
   auto __guard      = std::__make_scope_guard(__annotate_new_size(*this));
-  auto __allocation = std::__allocate_at_least(__alloc_, __cap + 1);
+  auto __allocation = __allocate_long_buffer(__alloc_, __cap);
   pointer __p       = __allocation.ptr;
-  __begin_lifetime(__p, __allocation.count);
   if (__n_copy != 0)
     traits_type::copy(std::__to_address(__p), std::__to_address(__old_p), __n_copy);
   if (__n_add != 0)
@@ -2625,12 +2565,8 @@ _LIBCPP_CONSTEXPR_SINCE_CXX20 void basic_string<_CharT, _Traits, _Allocator>::__
   if (__sec_cp_sz != 0)
     traits_type::copy(
         std::__to_address(__p) + __n_copy + __n_add, std::__to_address(__old_p) + __n_copy + __n_del, __sec_cp_sz);
-  if (__old_cap + 1 != __min_cap)
-    __alloc_traits::deallocate(__alloc_, __old_p, __old_cap + 1);
-  __set_long_pointer(__p);
-  __set_long_cap(__allocation.count);
   __old_sz = __n_copy + __n_add + __sec_cp_sz;
-  __set_long_size(__old_sz);
+  __replace_internal_buffer(__allocation, __old_sz);
   traits_type::assign(__p[__old_sz], value_type());
 }
 
@@ -2655,19 +2591,17 @@ _LIBCPP_DEPRECATED_("use __grow_by_without_replace") basic_string<_CharT, _Trait
   pointer __old_p = __get_pointer();
   size_type __cap =
       __old_cap < __ms / 2 - __alignment ? __recommend(std::max(__old_cap + __delta_cap, 2 * __old_cap)) : __ms - 1;
-  auto __allocation = std::__allocate_at_least(__alloc_, __cap + 1);
+  auto __allocation = __allocate_long_buffer(__alloc_, __cap);
   pointer __p       = __allocation.ptr;
-  __begin_lifetime(__p, __allocation.count);
   if (__n_copy != 0)
     traits_type::copy(std::__to_address(__p), std::__to_address(__old_p), __n_copy);
   size_type __sec_cp_sz = __old_sz - __n_del - __n_copy;
   if (__sec_cp_sz != 0)
     traits_type::copy(
         std::__to_address(__p) + __n_copy + __n_add, std::__to_address(__old_p) + __n_copy + __n_del, __sec_cp_sz);
-  if (__old_cap + 1 != __min_cap)
-    __alloc_traits::deallocate(__alloc_, __old_p, __old_cap + 1);
-  __set_long_pointer(__p);
-  __set_long_cap(__allocation.count);
+  // This is -1 to make sure the caller sets the size properly, since old versions of this function didn't set the size
+  // at all.
+  __replace_internal_buffer(__allocation, -1);
 }
 
 template <class _CharT, class _Traits, class _Allocator>
@@ -2684,6 +2618,7 @@ basic_string<_CharT, _Traits, _Allocator>::__grow_by_without_replace(
   _LIBCPP_SUPPRESS_DEPRECATED_PUSH
   __grow_by(__old_cap, __delta_cap, __old_sz, __n_copy, __n_del, __n_add);
   _LIBCPP_SUPPRESS_DEPRECATED_POP
+  // Due to the ABI of __grow_by we have to set the size after calling it.
   __set_long_size(__old_sz - __n_del + __n_add);
 }
 
@@ -2818,7 +2753,7 @@ basic_string<_CharT, _Traits, _Allocator>::__move_assign(basic_string& __str, tr
 {
   __annotate_delete();
   if (__is_long()) {
-    __alloc_traits::deallocate(__alloc_, __get_long_pointer(), __get_long_cap());
+    __reset_internal_buffer();
 #    if _LIBCPP_STD_VER <= 14
     if (!is_nothrow_move_assignable<allocator_type>::value) {
       __set_short_size(0);
@@ -3454,15 +3389,10 @@ _LIBCPP_CONSTEXPR_SINCE_CXX20 void basic_string<_CharT, _Traits, _Allocator>::re
     return;
 
   __annotation_guard __g(*this);
-  auto __allocation = std::__allocate_at_least(__alloc_, __recommend(__requested_capacity) + 1);
+  auto __allocation = __allocate_long_buffer(__alloc_, __requested_capacity);
   auto __size       = size();
-  __begin_lifetime(__allocation.ptr, __allocation.count);
   traits_type::copy(std::__to_address(__allocation.ptr), data(), __size + 1);
-  if (__is_long())
-    __alloc_traits::deallocate(__alloc_, __get_long_pointer(), __get_long_cap());
-  __set_long_cap(__allocation.count);
-  __set_long_size(__size);
-  __set_long_pointer(__allocation.ptr);
+  __replace_internal_buffer(__allocation, __size);
 }
 
 template <class _CharT, class _Traits, class _Allocator>
@@ -3476,12 +3406,11 @@ inline _LIBCPP_CONSTEXPR_SINCE_CXX20 void basic_string<_CharT, _Traits, _Allocat
   // We're a long string and we're shrinking into the small buffer.
   if (__fits_in_sso(__target_capacity)) {
     __annotation_guard __g(*this);
-    auto __ptr  = __get_long_pointer();
+    auto __allocation = __get_internal_long_buffer();
     auto __size = __get_long_size();
-    auto __cap  = __get_long_cap();
-    traits_type::copy(std::__to_address(__get_short_pointer()), std::__to_address(__ptr), __size + 1);
+    traits_type::copy(std::__to_address(__get_short_pointer()), std::__to_address(__allocation.ptr), __size + 1);
     __set_short_size(__size);
-    __alloc_traits::deallocate(__alloc_, __ptr, __cap);
+    __deallocate_long_buffer(__alloc_, __allocation);
     return;
   }
 
@@ -3490,22 +3419,18 @@ inline _LIBCPP_CONSTEXPR_SINCE_CXX20 void basic_string<_CharT, _Traits, _Allocat
 #  endif // _LIBCPP_HAS_EXCEPTIONS
     __annotation_guard __g(*this);
     auto __size       = size();
-    auto __allocation = std::__allocate_at_least(__alloc_, __target_capacity + 1);
+    auto __allocation = __allocate_long_buffer(__alloc_, size());
 
     // The Standard mandates shrink_to_fit() does not increase the capacity.
     // With equal capacity keep the existing buffer. This avoids extra work
     // due to swapping the elements.
-    if (__allocation.count - 1 >= capacity()) {
-      __alloc_traits::deallocate(__alloc_, __allocation.ptr, __allocation.count);
+    if (__allocation.count - 1 > capacity()) {
+      __deallocate_long_buffer(__alloc_, __allocation);
       return;
     }
 
-    __begin_lifetime(__allocation.ptr, __allocation.count);
-    auto __ptr = __get_long_pointer();
-    traits_type::copy(std::__to_address(__allocation.ptr), std::__to_address(__ptr), __size + 1);
-    __alloc_traits::deallocate(__alloc_, __ptr, __get_long_cap());
-    __set_long_cap(__allocation.count);
-    __set_long_pointer(__allocation.ptr);
+    traits_type::copy(std::__to_address(__allocation.ptr), std::__to_address(__get_long_pointer()), __size + 1);
+    __replace_internal_buffer(__allocation, size());
 #  if _LIBCPP_HAS_EXCEPTIONS
   } catch (...) {
     return;
diff --git a/libcxx/src/string.cpp b/libcxx/src/string.cpp
index e335639883dba..49e5f8f73445a 100644
--- a/libcxx/src/string.cpp
+++ b/libcxx/src/string.cpp
@@ -44,22 +44,9 @@ void __basic_string_common<true>::__throw_out_of_range() const { std::__throw_ou
 
 template <class _CharT, class _Traits, class _Allocator>
 void basic_string<_CharT, _Traits, _Allocator>::__init(const value_type* __s, size_type __sz, size_type __reserve) {
-  if (__libcpp_is_constant_evaluated())
-    __rep_ = __rep();
-  if (__reserve > max_size())
-    __throw_length_error();
-  pointer __p;
-  if (__fits_in_sso(__reserve)) {
-    __set_short_size(__sz);
-    __p = __get_short_pointer();
-  } else {
-    auto __allocation = std::__allocate_at_least(__alloc_, __recommend(__reserve) + 1);
-    __p               = __allocation.ptr;
-    __begin_lifetime(__p, __allocation.count);
-    __set_long_pointer(__p);
-    __set_long_cap(__allocation.count);
-    __set_long_size(__sz);
-  }
+  pointer __p = __init_internal_buffer(__reserve);
+  __annotate_delete();
+  __set_size(__sz);
   traits_type::copy(std::__to_address(__p), __s, __sz);
   traits_type::assign(__p[__sz], value_type());
   __annotate_new(__sz);
diff --git a/libcxx/test/std/strings/basic.string/string.cons/copy_alloc.pass.cpp b/libcxx/test/std/strings/basic.string/string.cons/copy_alloc.pass.cpp
index b0045cb4afbba..f074ce0bc4ce0 100644
--- a/libcxx/test/std/strings/basic.string/string.cons/copy_alloc.pass.cpp
+++ b/libcxx/test/std/strings/basic.string/string.cons/copy_alloc.pass.cpp
@@ -66,15 +66,6 @@ bool operator!=(const poca_alloc<T>& lhs, const poca_alloc<U>& rhs) {
   return lhs.imp != rhs.imp;
 }
 
-template <class S>
-TEST_CONSTEXPR_CXX20 void test_assign(S& s1, const S& s2) {
-  try {
-    s1 = s2;
-  } catch (std::bad_alloc&) {
-    return;
-  }
-  assert(false);
-}
 #endif
 
 template <class S>
@@ -122,7 +113,13 @@ TEST_CONSTEXPR_CXX20 bool test() {
     assert(s2 == p2);
 
     imp2.deactivate();
-    test_assign(s1, s2);
+
+    try {
+      s1 = s2;
+      assert(false);
+    } catch (std::bad_alloc&) {
+    }
+
     assert(s1 == p1);
     assert(s2 == p2);
   }



More information about the libcxx-commits mailing list