[libc++] Fix shrink_to_fit to swap buffer only when capacity is strictly smaller (#127321)

The current implementation of the `shrink_to_fit()` function of
`basic_string` swaps to the newly allocated buffer when the new buffer
has the same capacity as the existing one. While this is not incorrect,
it is truly unnecessary to swap to an equally-sized buffer. With equal
capacity, we should keep using the existing buffer and simply deallocate
the new one, avoiding the extra work of copying elements.

The desired behavior was documented in the following comment within the
function:


https://github.com/llvm/llvm-project/blob/61ad08792a86e62309b982189a600f4342a38d91/libcxx/include/string#L3560-L3566

However, the existing implementation did not exactly conform to this
guideline, which is a QoI matter.

This PR modifies the `shrink_to_fit()` function to ensure that the
buffer is only swapped when the new allocation is strictly smaller than
the existing one. When the capacities are equal, the new buffer will be
discarded without copying the elements. This is achieved by including
the `==` check in the above conditional logic.

Author: winner245

libcxx/include/string

@@ -3495,7 +3495,7 @@ inline _LIBCPP_CONSTEXPR_SINCE_CXX20 void basic_string<_CharT, _Traits, _Allocat
     // 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()) {
+    if (__allocation.count - 1 >= capacity()) {
       __alloc_traits::deallocate(__alloc_, __allocation.ptr, __allocation.count);
       return;
     }

libcxx/test/libcxx/strings/basic.string/string.capacity/shrink_to_fit.pass.cpp

@@ -12,12 +12,12 @@
 
 // void shrink_to_fit(); // constexpr since C++20
 
-// Make sure we use an allocation returned by allocate_at_least if it is smaller than the current allocation
-// even if it contains more bytes than we requested
-
 #include <cassert>
 #include <string>
 
+#include "asan_testing.h"
+#include "increasing_allocator.h"
+
 template <typename T>
 struct oversizing_allocator {
   using value_type       = T;
@@ -37,6 +37,9 @@ bool operator==(oversizing_allocator<T>, oversizing_allocator<U>) {
   return true;
 }
 
+// Make sure we use an allocation returned by allocate_at_least if it is smaller than the current allocation
+// even if it contains more bytes than we requested.
+// Fix issue: https://github.com/llvm/llvm-project/pull/115659
 void test_oversizing_allocator() {
   std::basic_string<char, std::char_traits<char>, oversizing_allocator<char>> s{
       "String does not fit in the internal buffer and is a bit longer"};
@@ -48,8 +51,37 @@ void test_oversizing_allocator() {
   assert(s.size() == size);
 }
 
+// Make sure libc++ shrink_to_fit does NOT swap buffer with equal allocation sizes
+void test_no_swap_with_equal_allocation_size() {
+  { // Test with custom allocator with a minimum allocation size
+    std::basic_string<char, std::char_traits<char>, min_size_allocator<128, char> > s(
+        "A long string exceeding SSO limit but within min alloc size");
+    std::size_t capacity = s.capacity();
+    std::size_t size     = s.size();
+    auto data            = s.data();
+    s.shrink_to_fit();
+    assert(s.capacity() <= capacity);
+    assert(s.size() == size);
+    assert(is_string_asan_correct(s));
+    assert(s.capacity() == capacity && s.data() == data);
+  }
+  { // Test with custom allocator with a minimum power of two allocation size
+    std::basic_string<char, std::char_traits<char>, pow2_allocator<char> > s(
+        "This is a long string that exceeds the SSO limit");
+    std::size_t capacity = s.capacity();
+    std::size_t size     = s.size();
+    auto data            = s.data();
+    s.shrink_to_fit();
+    assert(s.capacity() <= capacity);
+    assert(s.size() == size);
+    assert(is_string_asan_correct(s));
+    assert(s.capacity() == capacity && s.data() == data);
+  }
+}
+
 int main(int, char**) {
   test_oversizing_allocator();
+  test_no_swap_with_equal_allocation_size();
 
   return 0;
 }

libcxx/test/std/strings/basic.string/string.capacity/shrink_to_fit.pass.cpp

@@ -61,28 +61,25 @@ TEST_CONSTEXPR_CXX20 bool test() {
   test_string<std::basic_string<char, std::char_traits<char>, safe_allocator<char>>>();
 #endif
 
-  return true;
-}
-
 #if TEST_STD_VER >= 23
-// 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));
+  { // Make sure shrink_to_fit never increases capacity
+    // See: https://github.com/llvm/llvm-project/issues/95161
+    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
+
+  return true;
 }
-#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

libcxx/test/support/increasing_allocator.h

@@ -10,6 +10,7 @@
 #define TEST_SUPPORT_INCREASING_ALLOCATOR_H
 
 #include <cstddef>
+#include <limits>
 #include <memory>
 
 #include "test_macros.h"
@@ -49,4 +50,75 @@ TEST_CONSTEXPR_CXX20 bool operator==(increasing_allocator<T>, increasing_allocat
   return true;
 }
 
+template <std::size_t MinAllocSize, typename T>
+class min_size_allocator {
+public:
+  using value_type     = T;
+  min_size_allocator() = default;
+
+  template <typename U>
+  TEST_CONSTEXPR_CXX20 min_size_allocator(const min_size_allocator<MinAllocSize, U>&) TEST_NOEXCEPT {}
+
+  TEST_NODISCARD TEST_CONSTEXPR_CXX20 T* allocate(std::size_t n) {
+    if (n < MinAllocSize)
+      n = MinAllocSize;
+    return static_cast<T*>(::operator new(n * sizeof(T)));
+  }
+
+  TEST_CONSTEXPR_CXX20 void deallocate(T* p, std::size_t) TEST_NOEXCEPT { ::operator delete(static_cast<void*>(p)); }
+
+  template <typename U>
+  struct rebind {
+    using other = min_size_allocator<MinAllocSize, U>;
+  };
+};
+
+template <std::size_t MinAllocSize, typename T, typename U>
+TEST_CONSTEXPR_CXX20 bool
+operator==(const min_size_allocator<MinAllocSize, T>&, const min_size_allocator<MinAllocSize, U>&) {
+  return true;
+}
+
+template <std::size_t MinAllocSize, typename T, typename U>
+TEST_CONSTEXPR_CXX20 bool
+operator!=(const min_size_allocator<MinAllocSize, T>&, const min_size_allocator<MinAllocSize, U>&) {
+  return false;
+}
+
+template <typename T>
+class pow2_allocator {
+public:
+  using value_type = T;
+  pow2_allocator() = default;
+
+  template <typename U>
+  TEST_CONSTEXPR_CXX20 pow2_allocator(const pow2_allocator<U>&) TEST_NOEXCEPT {}
+
+  TEST_NODISCARD TEST_CONSTEXPR_CXX20 T* allocate(std::size_t n) {
+    return static_cast<T*>(::operator new(next_power_of_two(n) * sizeof(T)));
+  }
+
+  TEST_CONSTEXPR_CXX20 void deallocate(T* p, std::size_t) TEST_NOEXCEPT { ::operator delete(static_cast<void*>(p)); }
+
+private:
+  TEST_CONSTEXPR_CXX20 std::size_t next_power_of_two(std::size_t n) const {
+    if ((n & (n - 1)) == 0)
+      return n;
+    for (std::size_t shift = 1; shift < std::numeric_limits<std::size_t>::digits; shift <<= 1) {
+      n |= n >> shift;
+    }
+    return n + 1;
+  }
+};
+
+template <typename T, typename U>
+TEST_CONSTEXPR_CXX20 bool operator==(const pow2_allocator<T>&, const pow2_allocator<U>&) {
+  return true;
+}
+
+template <typename T, typename U>
+TEST_CONSTEXPR_CXX20 bool operator!=(const pow2_allocator<T>&, const pow2_allocator<U>&) {
+  return false;
+}
+
 #endif // TEST_SUPPORT_INCREASING_ALLOCATOR_H