From 86fbf0905eb85ded4f7f16f8f5eff61e81f2f67d Mon Sep 17 00:00:00 2001
From: Nikolas Klauser <nikolasklauser@berlin.de>
Date: Sun, 10 Nov 2024 14:24:53 +0100
Subject: [PATCH] [libc++] Fix throwing away smaller allocations in
 string::shrink_to_fit

---
 libcxx/include/string                         |  2 +-
 .../string.capacity/shrink_to_fit.pass.cpp    | 55 +++++++++++++++++++
 2 files changed, 56 insertions(+), 1 deletion(-)
 create mode 100644 libcxx/test/libcxx/strings/basic.string/string.capacity/shrink_to_fit.pass.cpp

diff --git a/libcxx/include/string b/libcxx/include/string
index b1cedbe68f795..1034be8fdad8b 100644
--- a/libcxx/include/string
+++ b/libcxx/include/string
@@ -3393,7 +3393,7 @@ basic_string<_CharT, _Traits, _Allocator>::__shrink_or_extend(size_type __target
         // 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) {
+        if (__allocation.count - 1 > capacity()) {
           __alloc_traits::deallocate(__alloc_, __allocation.ptr, __allocation.count);
           __annotate_new(__sz); // Undoes the __annotate_delete()
           return;
diff --git a/libcxx/test/libcxx/strings/basic.string/string.capacity/shrink_to_fit.pass.cpp b/libcxx/test/libcxx/strings/basic.string/string.capacity/shrink_to_fit.pass.cpp
new file mode 100644
index 0000000000000..73b70d6f10bd5
--- /dev/null
+++ b/libcxx/test/libcxx/strings/basic.string/string.capacity/shrink_to_fit.pass.cpp
@@ -0,0 +1,55 @@
+//===----------------------------------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+// UNSUPPORTED: c++03, c++11, c++14, c++17, c++20
+
+// <string>
+
+// 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>
+
+template <typename T>
+struct oversizing_allocator {
+  using value_type       = T;
+  oversizing_allocator() = default;
+  template <typename U>
+  oversizing_allocator(const oversizing_allocator<U>&) noexcept {}
+  std::allocation_result<T*> allocate_at_least(std::size_t n) {
+    ++n;
+    return {static_cast<T*>(::operator new(n * sizeof(T))), n};
+  }
+  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==(oversizing_allocator<T>, oversizing_allocator<U>) {
+  return true;
+}
+
+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"};
+  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);
+}
+
+int main(int, char**) {
+  test_oversizing_allocator();
+
+  return 0;
+}