[libcxx-commits] [libcxx] [libc++][string] Fixes shrink_to_fit. (PR #97961)

[Mark de Wever via libcxx-commits]

https://github.com/mordante created https://github.com/llvm/llvm-project/pull/97961

This assures shrink_to_fit does not increase the allocated size.

Partly addresses #95161

>From 0de62f1743855f337c70db28aba11afffaabcda7 Mon Sep 17 00:00:00 2001
From: Mark de Wever <koraq at xs4all.nl>
Date: Sat, 6 Jul 2024 16:28:59 +0200
Subject: [PATCH] [libc++][string] Fixes shrink_to_fit.

This assures shrink_to_fit does not increase the allocated size.

Partly addresses #95161
---
 libcxx/include/string                         | 21 ++++++++--
 .../string.capacity/shrink_to_fit.pass.cpp    | 41 +++++++++++++++++++
 2 files changed, 59 insertions(+), 3 deletions(-)

diff --git a/libcxx/include/string b/libcxx/include/string
index 9a52ab6aef41e..22d11b99ed693 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 057050cdcf7fa..6f5e43d1341f5 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