Add tests to ensure no buffer swapping with equal alloc size

Author: winner245

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

@@ -61,28 +61,54 @@ 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
+
+  {   // Ensure that the libc++ implementation of shrink_to_fit does NOT swap buffer with equal allocation sizes
+    { // 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);
+      LIBCPP_ASSERT(is_string_asan_correct(s));
+      if (s.capacity() == capacity)
+        LIBCPP_ASSERT(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);
+      LIBCPP_ASSERT(is_string_asan_correct(s));
+      if (s.capacity() == capacity)
+        LIBCPP_ASSERT(s.data() == data);
+    }
+  }
+
+  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,89 @@ 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 {}
+
+#if TEST_STD_VER >= 23
+  TEST_CONSTEXPR_CXX23 std::allocation_result<T*> allocate_at_least(std::size_t n) {
+    if (n < MinAllocSize)
+      n = MinAllocSize;
+    return std::allocator<T>{}.allocate_at_least(n);
+  }
+#endif // TEST_STD_VER >= 23
+
+  TEST_NODISCARD TEST_CONSTEXPR_CXX20 T* allocate(std::size_t n) {
+    if (n < MinAllocSize)
+      n = MinAllocSize;
+    return std::allocator<T>().allocate(n);
+  }
+
+  TEST_CONSTEXPR_CXX20 void deallocate(T* p, std::size_t n) TEST_NOEXCEPT { std::allocator<T>().deallocate(p, n); }
+
+  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 {}
+
+#if TEST_STD_VER >= 23
+  TEST_CONSTEXPR_CXX23 std::allocation_result<T*> allocate_at_least(std::size_t n) {
+    return std::allocator<T>{}.allocate_at_least(next_power_of_two(n));
+  }
+#endif // TEST_STD_VER >= 23
+
+  TEST_NODISCARD TEST_CONSTEXPR_CXX20 T* allocate(std::size_t n) {
+    return std::allocator<T>().allocate(next_power_of_two(n));
+  }
+
+  TEST_CONSTEXPR_CXX20 void deallocate(T* p, std::size_t n) TEST_NOEXCEPT { std::allocator<T>().deallocate(p, n); }
+
+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