[libc++][vector] Fixes shrink_to_fit.

[mordante]

This assures shrink_to_fit does not increase the allocated size.

Partly addresses #95161

[llvmbot]

@llvm/pr-subscribers-libcxx

Author: Mark de Wever (mordante)

Changes

This assures shrink_to_fit does not increase the allocated size.

Partly addresses #95161

---

Full diff: https://github.com/llvm/llvm-project/pull/97895.diff

2 Files Affected:

• (modified) libcxx/include/vector (+5-1)
• (modified) libcxx/test/std/containers/sequences/vector/vector.capacity/shrink_to_fit.pass.cpp (+37)

diff --git a/libcxx/include/vector b/libcxx/include/vector
index 299ad8c9b23f2..7c7adb0da7da5 100644
--- a/libcxx/include/vector
+++ b/libcxx/include/vector
@@ -1443,7 +1443,11 @@ _LIBCPP_CONSTEXPR_SINCE_CXX20 void vector<_Tp, _Allocator>::shrink_to_fit() _NOE
 #endif // _LIBCPP_HAS_NO_EXCEPTIONS
       allocator_type& __a = this->__alloc();
       __split_buffer<value_type, allocator_type&> __v(size(), size(), __a);
-      __swap_out_circular_buffer(__v);
+      // 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 (__v.capacity() < capacity())
+        __swap_out_circular_buffer(__v);
 #ifndef _LIBCPP_HAS_NO_EXCEPTIONS
     } catch (...) {
     }
diff --git a/libcxx/test/std/containers/sequences/vector/vector.capacity/shrink_to_fit.pass.cpp b/libcxx/test/std/containers/sequences/vector/vector.capacity/shrink_to_fit.pass.cpp
index 8851e2a9ed0c7..facb724c713af 100644
--- a/libcxx/test/std/containers/sequences/vector/vector.capacity/shrink_to_fit.pass.cpp
+++ b/libcxx/test/std/containers/sequences/vector/vector.capacity/shrink_to_fit.pass.cpp
@@ -71,9 +71,46 @@ TEST_CONSTEXPR_CXX20 bool tests() {
     return true;
 }
 
+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};
+  }
+  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::vector<int, increasing_allocator<int>> v;
+  v.push_back(1);
+  assert(is_contiguous_container_asan_correct(v));
+  std::size_t capacity = v.capacity();
+  v.shrink_to_fit();
+  assert(v.capacity() == capacity);
+  assert(v.size() == 1);
+  assert(is_contiguous_container_asan_correct(v));
+}
+
 int main(int, char**)
 {
     tests();
+    test_increasing_allocator();
 #if TEST_STD_VER > 17
     static_assert(tests());
 #endif

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[ldionne]

It is somewhat frustrating to throw away the allocation we just made because it ends up being larger than the one we already had, but it makes us conforming and it does save us from having to copy the elements over to the new buffer. So technically we are saving some work (depending on the relative cost of everything, of course).

What we'd really like is a way to ask the allocator for at most n bytes. That way the allocator could instead report that it can only give us more than n bytes, and we'd likely prevent the system allocator from having to do additional work.

Anyway, this LGTM with a suggestion for the tests.

[mordante]

What we'd really like is a way to ask the allocator for at most n bytes. That way the allocator could instead report that it can only give us more than n bytes, and we'd likely prevent the system allocator from having to do additional work.

That is what allocate does, but that secretly throws away bytes. I expect the case where the allocator gives more bytes would be rare, the arena allocator in the example seems a good example.