[libc++] Simplify the implementation of reserve() and shrink_to_fit()

[philnik777]

Since we changed the implementation of reserve(size_type) to only ever extend,
it doesn't make a ton of sense anymore to have __shrink_or_extend, since the code
paths of reserve and shrink_to_fit are now almost completely separate.

This patch splits up __shrink_or_extend so that the individual parts are in reserve
and shrink_to_fit depending on where they are needed.

[github-actions]

⚠️ C/C++ code formatter, clang-format found issues in your code. ⚠️

You can test this locally with the following command:

git-clang-format --diff aaa0dd2f05ff957a171a87e78578dddc59fc49c2 34f59690a5135fe3b6075d83d6ba2cb219469660 --extensions  -- libcxx/include/string

View the diff from clang-format here.

diff --git a/libcxx/include/string b/libcxx/include/string
index cdc399361e..157ca6a640 100644
--- a/libcxx/include/string
+++ b/libcxx/include/string
@@ -3373,7 +3373,7 @@ _LIBCPP_CONSTEXPR_SINCE_CXX20 void basic_string<_CharT, _Traits, _Allocator>::re
 
   __annotation_guard __g(*this);
   auto __allocation = std::__allocate_at_least(__alloc_, __recommend(__requested_capacity) + 1);
-  auto __size = size();
+  auto __size       = size();
   __begin_lifetime(__allocation.ptr, __allocation.count);
   traits_type::copy(std::__to_address(__allocation.ptr), data(), __size + 1);
   if (__is_long())

[llvmbot]

@llvm/pr-subscribers-libcxx

Author: Nikolas Klauser (philnik777)

Changes

Since we changed the implementation of reserve(size_type) to only ever extend,
it doesn't make a ton of sense anymore to have __shrink_or_extend, since the code
paths of reserve and shrink_to_fit are now almost completely separate.

This patch splits up __shrink_or_extend so that the individual parts are in reserve
and shrink_to_fit depending on where they are needed.

---

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

1 Files Affected:

• (modified) libcxx/include/string (+62-62)

diff --git a/libcxx/include/string b/libcxx/include/string
index 4b5017f5e7753f..55b174f4db987c 100644
--- a/libcxx/include/string
+++ b/libcxx/include/string
@@ -1874,8 +1874,6 @@ private:
   operator==(const basic_string<char, char_traits<char>, _Alloc>& __lhs,
              const basic_string<char, char_traits<char>, _Alloc>& __rhs) _NOEXCEPT;
 
-  _LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX20 void __shrink_or_extend(size_type __target_capacity);
-
   _LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX20 _LIBCPP_STRING_INTERNAL_MEMORY_ACCESS bool
   __is_long() const _NOEXCEPT {
     if (__libcpp_is_constant_evaluated() && __builtin_constant_p(__rep_.__l.__is_long_)) {
@@ -2060,6 +2058,21 @@ private:
 #endif
   }
 
+  // Disable ASan annotations and enable them again when going out of scope. It is assumed that the string is in a valid
+  // state at that point, so `size()` can be called safely.
+  struct [[__nodiscard__]] __annotation_guard {
+    __annotation_guard(const __annotation_guard&) = delete;
+    __annotation_guard& operator=(const __annotation_guard&) = delete;
+
+    _LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX20 __annotation_guard(basic_string& __str) : __str_(__str) {
+      __str_.__annotate_delete();
+    }
+
+    _LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX20 ~__annotation_guard() { __str_.__annotate_new(__str_.size()); }
+
+    basic_string& __str_;
+  };
+
   template <size_type __a>
   static _LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX20 size_type __align_it(size_type __s) _NOEXCEPT {
     return (__s + (__a - 1)) & ~(__a - 1);
@@ -3340,7 +3353,16 @@ _LIBCPP_CONSTEXPR_SINCE_CXX20 void basic_string<_CharT, _Traits, _Allocator>::re
   if (__target_capacity == capacity())
     return;
 
-  __shrink_or_extend(__target_capacity);
+  __annotation_guard __g(*this);
+  auto __allocation = std::__allocate_at_least(__alloc_, __target_capacity + 1);
+  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_size() + 1);
+  __set_long_cap(__allocation.count);
+  __set_long_size(__size);
+  __set_long_pointer(__allocation.ptr);
 }
 
 template <class _CharT, class _Traits, class _Allocator>
@@ -3349,70 +3371,48 @@ inline _LIBCPP_CONSTEXPR_SINCE_CXX20 void basic_string<_CharT, _Traits, _Allocat
   if (__target_capacity == capacity())
     return;
 
-  __shrink_or_extend(__target_capacity);
-}
-
-template <class _CharT, class _Traits, class _Allocator>
-inline _LIBCPP_CONSTEXPR_SINCE_CXX20 void
-basic_string<_CharT, _Traits, _Allocator>::__shrink_or_extend(size_type __target_capacity) {
-  __annotate_delete();
-  size_type __cap = capacity();
-  size_type __sz  = size();
-
-  pointer __new_data, __p;
-  bool __was_long, __now_long;
   if (__fits_in_sso(__target_capacity)) {
-    __was_long = true;
-    __now_long = false;
-    __new_data = __get_short_pointer();
-    __p        = __get_long_pointer();
-  } else {
-    if (__target_capacity > __cap) {
-      // Extend
-      // - called from reserve should propagate the exception thrown.
-      auto __allocation = std::__allocate_at_least(__alloc_, __target_capacity + 1);
-      __new_data        = __allocation.ptr;
-      __target_capacity = __allocation.count - 1;
-    } else {
-      // Shrink
-      // - called from shrink_to_fit should not throw.
-      // - called from reserve may throw but is not required to.
+    if (!__is_long())
+      return;
+    __annotation_guard __g(*this);
+    auto __ptr = __get_long_pointer();
+    auto __size = __get_long_size();
+    auto __cap = __get_long_cap();
+    traits_type::copy(std::__to_address(__get_short_pointer()), data(), __size + 1);
+    __alloc_traits::deallocate(__alloc_, __ptr, __cap);
+    __set_short_size(__size);
+    return;
+  }
+
+  // Shrink
+  // - called from shrink_to_fit should not throw.
+  // - called from reserve may throw but is not required to.
 #if _LIBCPP_HAS_EXCEPTIONS
-      try {
+  try {
 #endif // _LIBCPP_HAS_EXCEPTIONS
-        auto __allocation = std::__allocate_at_least(__alloc_, __target_capacity + 1);
-
-        // 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 > __target_capacity) {
-          __alloc_traits::deallocate(__alloc_, __allocation.ptr, __allocation.count);
-          __annotate_new(__sz); // Undoes the __annotate_delete()
-          return;
-        }
-        __new_data        = __allocation.ptr;
-        __target_capacity = __allocation.count - 1;
+    __annotation_guard __g(*this);
+    auto __size = size();
+    auto __allocation = std::__allocate_at_least(__alloc_, __target_capacity + 1);
+
+    // 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 > __target_capacity) {
+      __alloc_traits::deallocate(__alloc_, __allocation.ptr, __allocation.count);
+      __annotate_new(__size); // Undoes the __annotate_delete()
+      return;
+    }
+
+    __begin_lifetime(__allocation.ptr, __allocation.count);
+    traits_type::copy(std::__to_address(__allocation.ptr), data(), size() + 1);
+    __alloc_traits::deallocate(__alloc_, __get_long_pointer(), __get_long_cap());
+    __set_long_cap(__allocation.count);
+    __set_long_pointer(__allocation.ptr);
 #if _LIBCPP_HAS_EXCEPTIONS
-      } catch (...) {
-        return;
-      }
+  } catch (...) {
+    return;
+  }
 #endif // _LIBCPP_HAS_EXCEPTIONS
-    }
-    __begin_lifetime(__new_data, __target_capacity + 1);
-    __now_long = true;
-    __was_long = __is_long();
-    __p        = __get_pointer();
-  }
-  traits_type::copy(std::__to_address(__new_data), std::__to_address(__p), size() + 1);
-  if (__was_long)
-    __alloc_traits::deallocate(__alloc_, __p, __cap + 1);
-  if (__now_long) {
-    __set_long_cap(__target_capacity + 1);
-    __set_long_size(__sz);
-    __set_long_pointer(__new_data);
-  } else
-    __set_short_size(__sz);
-  __annotate_new(__sz);
 }
 
 template <class _CharT, class _Traits, class _Allocator>

[ldionne]

This review paid off! Looks like there are many things we can improve / fix in the existing code. I think we should split this into a few PRs:

1. Introduce __scope_guard and simplify the __annotate stuff in the existing code.
2. Fix the bug with shrink_to_fit never increases capacity. This needs an additional test.
3. Remove the unnecessary capacity checks inside reserve()
4. Then, this patch.

[ldionne]

I went over the whole code again after your rebases, and this LGTM. Thanks for the simplification.

[d0k]

Breaks with sized deallocation. See inline comment, reverted this in 59f57be

[d0k]

This should pass capacity, not size.