[libc++][NFC] 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 ac5a2010ad35a72de3e75a1883e2495345b92a73 cf86ee2c74422961a1b7f59ed6789cd7b63f6f09 --extensions  -- libcxx/include/string

View the diff from clang-format here.

diff --git a/libcxx/include/string b/libcxx/include/string
index 55b174f4db..8805a0ea05 100644
--- a/libcxx/include/string
+++ b/libcxx/include/string
@@ -2061,7 +2061,7 @@ private:
   // 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(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) {
@@ -3355,7 +3355,7 @@ _LIBCPP_CONSTEXPR_SINCE_CXX20 void basic_string<_CharT, _Traits, _Allocator>::re
 
   __annotation_guard __g(*this);
   auto __allocation = std::__allocate_at_least(__alloc_, __target_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())
@@ -3375,9 +3375,9 @@ inline _LIBCPP_CONSTEXPR_SINCE_CXX20 void basic_string<_CharT, _Traits, _Allocat
     if (!__is_long())
       return;
     __annotation_guard __g(*this);
-    auto __ptr = __get_long_pointer();
+    auto __ptr  = __get_long_pointer();
     auto __size = __get_long_size();
-    auto __cap = __get_long_cap();
+    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);
@@ -3391,7 +3391,7 @@ inline _LIBCPP_CONSTEXPR_SINCE_CXX20 void basic_string<_CharT, _Traits, _Allocat
   try {
 #endif // _LIBCPP_HAS_EXCEPTIONS
     __annotation_guard __g(*this);
-    auto __size = size();
+    auto __size       = size();
     auto __allocation = std::__allocate_at_least(__alloc_, __target_capacity + 1);
 
     // The Standard mandates shrink_to_fit() does not increase the capacity.

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

Let's introduce an actual __scope_guard and then use it here. It's going to be more general and it will make this code easier to understand.

[ldionne]

Above, size_type __target_capacity = std::max(__requested_capacity, size()); can be simplified to just size_type __target_capacity = __requested_capacity; since at that point we know for sure that the requested capacity is greater than our current capacity (that's line 3349 above).

And similarly, the if (__target_capacity == capacity()) can be removed entirely, since we know that we are requesting more than our current capacity.

To do this refactoring, let's add a _LIBCPP_ASSERT_INTERNAL to __recommend(__n) that the value we return is always >= __n.

[ldionne]

Suggested change

	if (__fits_in_sso(__target_capacity)) {
	// We're a long string and we're downsizing into the small buffer.
	if (__fits_in_sso(__target_capacity)) {

[ldionne]

Suggested change

	traits_type::copy(std::__to_address(__get_short_pointer()), data(), __size + 1);
	traits_type::copy(std::__to_address(__get_short_pointer()), std::__to_address(__ptr), __size + 1);

[ldionne]

IMO it makes more sense to group __set_short_size and the copy together, since this is modifying the current state of the string representation. This could also be a single "operation" if this were refactored. On the other hand, reclaiming the memory of the old long string is definitely a distinct step.

[ldionne]

We should add _LIBCPP_ASSERT_INTERNAL(__is_long()) here, and then the if (!__is_long()) below can go away.

[ldionne]

This is now handled by your guard, it should be removed.

[ldionne]

This looks broken to me. I think this check should be something like if (__allocation.count - 1 > capacity()), otherwise we're basically entering this if whenever the allocator gives us more than we requested. I think what we intended to do is to instead enter the if only when the allocator gives us more than our original capacity.

This was introduced by d0ca9f2, fortunately not released yet.

This needs a test (probably fixing libcxx/test/std/strings/basic.string/string.capacity/shrink_to_fit.pass.cpp is sufficient).

[ldionne]

Suggested change

	traits_type::copy(std::__to_address(__allocation.ptr), data(), size() + 1);
	__alloc_traits::deallocate(__alloc_, __get_long_pointer(), __get_long_cap());
	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());

[ldionne]

Let's move the try-catch around this line only, since that's the only thing we're guarding for an exception.