[libc++][string] Fixes shrink_to_fit.

[mordante]

This assures shrink_to_fit does not increase the allocated size.

Partly addresses #95161

[llvmbot]

@llvm/pr-subscribers-libcxx

Author: Mark de Wever (mordante)

Changes

This assures shrink_to_fit does not increase the allocated size.

Partly addresses #95161

---

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

2 Files Affected:

• (modified) libcxx/include/string (+18-3)
• (modified) libcxx/test/std/strings/basic.string/string.capacity/shrink_to_fit.pass.cpp (+41)

diff --git a/libcxx/include/string b/libcxx/include/string
index 9a52ab6aef41e8..22d11b99ed693d 100644
--- a/libcxx/include/string
+++ b/libcxx/include/string
@@ -3265,23 +3265,38 @@ basic_string<_CharT, _Traits, _Allocator>::__shrink_or_extend(size_type __target
     __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.
 #ifndef _LIBCPP_HAS_NO_EXCEPTIONS
       try {
 #endif // _LIBCPP_HAS_NO_EXCEPTIONS
         auto __allocation = std::__allocate_at_least(__alloc(), __target_capacity + 1);
+
+#ifdef _LIBCPP_HAS_NO_EXCEPTIONS
+        if (__allocation.ptr == nullptr)
+          return;
+#endif // _LIBCPP_HAS_NO_EXCEPTIONS
+
+        // 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);
+          return;
+        }
         __new_data        = __allocation.ptr;
         __target_capacity = __allocation.count - 1;
 #ifndef _LIBCPP_HAS_NO_EXCEPTIONS
       } catch (...) {
         return;
       }
-#else  // _LIBCPP_HAS_NO_EXCEPTIONS
-      if (__new_data == nullptr)
-        return;
 #endif // _LIBCPP_HAS_NO_EXCEPTIONS
     }
     __begin_lifetime(__new_data, __target_capacity + 1);
diff --git a/libcxx/test/std/strings/basic.string/string.capacity/shrink_to_fit.pass.cpp b/libcxx/test/std/strings/basic.string/string.capacity/shrink_to_fit.pass.cpp
index 057050cdcf7fa3..6f5e43d1341f53 100644
--- a/libcxx/test/std/strings/basic.string/string.capacity/shrink_to_fit.pass.cpp
+++ b/libcxx/test/std/strings/basic.string/string.capacity/shrink_to_fit.pass.cpp
@@ -63,8 +63,49 @@ TEST_CONSTEXPR_CXX20 bool test() {
   return true;
 }
 
+#if TEST_STD_VER >= 23
+std::size_t min_bytes = 1000;
+
+template <typename T>
+struct increasing_allocator {
+  using value_type       = T;
+  increasing_allocator() = default;
+  template <typename U>
+  increasing_allocator(const increasing_allocator<U>&) noexcept {}
+  std::allocation_result<T*> allocate_at_least(std::size_t n) {
+    std::size_t allocation_amount = n * sizeof(T);
+    if (allocation_amount < min_bytes)
+      allocation_amount = min_bytes;
+    min_bytes += 1000;
+    return {static_cast<T*>(::operator new(allocation_amount)), allocation_amount / sizeof(T)};
+  }
+  T* allocate(std::size_t n) { return allocate_at_least(n).ptr; }
+  void deallocate(T* p, std::size_t) noexcept { ::operator delete(static_cast<void*>(p)); }
+};
+
+template <typename T, typename U>
+bool operator==(increasing_allocator<T>, increasing_allocator<U>) {
+  return true;
+}
+
+// https://github.com/llvm/llvm-project/issues/95161
+void test_increasing_allocator() {
+  std::basic_string<char, std::char_traits<char>, increasing_allocator<char>> s{
+      "String does not fit in the internal buffer"};
+  std::size_t capacity = s.capacity();
+  std::size_t size     = s.size();
+  s.shrink_to_fit();
+  assert(s.capacity() <= capacity);
+  assert(s.size() == size);
+  LIBCPP_ASSERT(is_string_asan_correct(s));
+}
+#endif // TEST_STD_VER >= 23
+
 int main(int, char**) {
   test();
+#if TEST_STD_VER >= 23
+  test_increasing_allocator();
+#endif
 #if TEST_STD_VER > 17
   static_assert(test());
 #endif

[ldionne]

Adding this here really assumes that we're only getting into this code path in shrink_to_fit. Otherwise, we may be discarding the larger block of memory we just allocated without really needing to. Would it make sense to instead rewrite shrink_to_fit to avoid using this __shrink_or_extend method which seems kinda convoluted?

[mordante]

It would require quite a bit of copy pasting. Both resize and shrink_to_fit need to handle the SSO buffer part. IMO it's no issue to keep this in one function. I expect the growing buffer is very unlikely to happen with typical allocators. I also expect the shrinking reserve and shrink_to_fit are not very popular functions unless you know it will safe a considerable amount of memory.

[ldionne]

LGTM. Changes of this nature to std::string are always really tricky. Let's keep an eye out for unintended fallout, and if anything happens, we should likely revert and try landing this into LLVM 20 instead, even though it means that vector and string will behave differently w.r.t. shrink_to_fit.

[ldionne]

/cherry-pick d0ca9f2

[llvmbot]

/pull-request #100149