[libc++] Fix shrink_to_fit to swap buffer only when capacity is strictly smaller

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