[libc++] Slightly simplify max_size and add new tests for vector

[winner245]

This PR slightly simplifies the implementation of vector<bool>::max_size and adds extensive tests for the max_size() function for both vector and std::vector. The main purposes of the new tests include:

• Verify correctness of max_size() under various size_type and difference_type definitions: check that max_size() works properly with allocators that have custom size_type and difference_type. This is particularly useful for vector<bool>, as different size_type lead to different __storage_type of different word lengths, resulting in varying max_size() values for vector<bool>. Additionally, different difference_type also sets different upper limit of max_size() for both vector<bool> and std::vector. These tests were previously missing.

• Eliminate incorrect implementations: Special tests are added to identify and reject incorrect implementations of vector<bool>::max_size that unconditionally return std::min<size_type>(__nmax, __internal_cap_to_external(__amax)). This can cause overflow in the __internal_cap_to_external() call and lead to incorrect results. The new tests ensure that such incorrect implementations are identified.

[github-actions]

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

[llvmbot]

@llvm/pr-subscribers-libcxx

Author: Peng Liu (winner245)

Changes

This PR slightly simplifies the implementation of vector<bool>::max_size and adds extensive tests for the max_size() function for both vector<bool> and std::vector. The main purposes of the new tests include:

• Verify correctness of max_size() under various size_type and difference_type definitions: check that max_size() works properly with allocators that have custom size_type and difference_type. This is particularly useful for vector<bool>, as different size_type lead to different __storage_type of different word lengths, resulting in varying max_size() values for vector<bool>. Additionally, different difference_type also sets different upper limit of max_size() for both vector<bool> and std::vector. These tests were previously missing.

• Eliminate incorrect implementations: Special tests are added to identify and reject incorrect implementations of vector<bool>::max_size that unconditionally return std::min<size_type>(__nmax, __internal_cap_to_external(__amax)). This can cause overflow in the __internal_cap_to_external() call and lead to incorrect results. The new tests ensure that such incorrect implementations are identified.

---

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

4 Files Affected:

• (modified) libcxx/include/__vector/vector_bool.h (+2-4)
• (added) libcxx/test/std/containers/sequences/vector.bool/max_size.pass.cpp (+105)
• (modified) libcxx/test/std/containers/sequences/vector/vector.capacity/max_size.pass.cpp (+45-7)
• (added) libcxx/test/support/sized_allocator.h (+58)

diff --git a/libcxx/include/__vector/vector_bool.h b/libcxx/include/__vector/vector_bool.h
index 36eb7f350ac406..84e0167e71a6cb 100644
--- a/libcxx/include/__vector/vector_bool.h
+++ b/libcxx/include/__vector/vector_bool.h
@@ -511,10 +511,8 @@ template <class _Allocator>
 _LIBCPP_CONSTEXPR_SINCE_CXX20 typename vector<bool, _Allocator>::size_type
 vector<bool, _Allocator>::max_size() const _NOEXCEPT {
   size_type __amax = __storage_traits::max_size(__alloc_);
-  size_type __nmax = numeric_limits<size_type>::max() / 2; // end() >= begin(), always
-  if (__nmax / __bits_per_word <= __amax)
-    return __nmax;
-  return __internal_cap_to_external(__amax);
+  size_type __nmax = numeric_limits<difference_type>::max();
+  return __nmax / __bits_per_word <= __amax ? __nmax : __internal_cap_to_external(__amax);
 }
 
 //  Precondition:  __new_size > capacity()
diff --git a/libcxx/test/std/containers/sequences/vector.bool/max_size.pass.cpp b/libcxx/test/std/containers/sequences/vector.bool/max_size.pass.cpp
new file mode 100644
index 00000000000000..84cf2a7de539e0
--- /dev/null
+++ b/libcxx/test/std/containers/sequences/vector.bool/max_size.pass.cpp
@@ -0,0 +1,105 @@
+//===----------------------------------------------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+// <vector>
+
+// size_type max_size() const;
+
+#include <cassert>
+#include <climits>
+#include <cstdint>
+#include <limits>
+#include <memory>
+#include <type_traits>
+#include <vector>
+
+#include "min_allocator.h"
+#include "sized_allocator.h"
+#include "test_allocator.h"
+#include "test_macros.h"
+
+#if TEST_STD_VER >= 11
+
+template <typename Alloc>
+TEST_CONSTEXPR_CXX20 void test(const std::vector<bool, Alloc>& v) {
+  using Vector              = std::vector<bool, Alloc>;
+  using size_type           = typename Vector::size_type;
+  using difference_type     = typename Vector::difference_type;
+  using storage_type        = typename Vector::__storage_type;
+  using storage_alloc       = typename std::allocator_traits<Alloc>::template rebind_alloc<storage_type>;
+  using storage_traits      = typename std::allocator_traits<Alloc>::template rebind_traits<storage_type>;
+  const size_type max_dist  = static_cast<size_type>(std::numeric_limits<difference_type>::max());
+  const size_type max_alloc = storage_traits::max_size(storage_alloc(v.get_allocator()));
+  std::size_t bits_per_word = sizeof(storage_type) * CHAR_BIT;
+  const size_type max_size  = max_dist / bits_per_word < max_alloc ? max_dist : max_alloc * bits_per_word;
+  assert(v.max_size() <= max_dist);
+  assert(v.max_size() / bits_per_word <= max_alloc); // max_alloc * bits_per_word may overflow
+  LIBCPP_ASSERT(v.max_size() == max_size);
+}
+
+#endif
+
+TEST_CONSTEXPR_CXX20 bool tests() {
+  // Test with limited_allocator where v.max_size() is determined by allocator::max_size()
+  {
+    using Alloc        = limited_allocator<bool, 10>;
+    using Vector       = std::vector<bool, Alloc>;
+    using storage_type = Vector::__storage_type;
+    Vector v;
+    std::size_t bits_per_word = sizeof(storage_type) * CHAR_BIT;
+    assert(v.max_size() <= 10 * bits_per_word);
+    LIBCPP_ASSERT(v.max_size() == 10 * bits_per_word);
+  }
+  {
+    using Alloc        = limited_allocator<double, 10>;
+    using Vector       = std::vector<bool, Alloc>;
+    using storage_type = Vector::__storage_type;
+    Vector v;
+    std::size_t bits_per_word = sizeof(storage_type) * CHAR_BIT;
+    assert(v.max_size() <= 10 * bits_per_word);
+    LIBCPP_ASSERT(v.max_size() == 10 * bits_per_word);
+  }
+
+#if TEST_STD_VER >= 11
+
+  // Test with various allocators and diffrent size_type
+  {
+    test(std::vector<bool>());
+    test(std::vector<bool, std::allocator<int> >());
+    test(std::vector<bool, min_allocator<bool> >());
+    test(std::vector<bool, test_allocator<bool> >(test_allocator<bool>(1)));
+    test(std::vector<bool, other_allocator<bool> >(other_allocator<bool>(5)));
+    test(std::vector<bool, sized_allocator<bool, std::uint8_t, std::int8_t> >());
+    test(std::vector<bool, sized_allocator<bool, std::uint16_t, std::int16_t> >());
+    test(std::vector<bool, sized_allocator<bool, std::uint32_t, std::int32_t> >());
+    test(std::vector<bool, sized_allocator<bool, std::uint64_t, std::int64_t> >());
+    test(std::vector<bool, limited_allocator<bool, static_cast<std::size_t>(-1)> >());
+  }
+
+  // Test cases to identify incorrect implementations that unconditionally return:
+  //   std::min<size_type>(__nmax, __internal_cap_to_external(__amax))
+  // This can cause overflow in __internal_cap_to_external and lead to incorrect results.
+  {
+    test(std::vector<bool, limited_allocator<bool, static_cast<std::size_t>(-1) / 61> >());
+    test(std::vector<bool, limited_allocator<bool, static_cast<std::size_t>(-1) / 63> >());
+  }
+
+#endif
+
+  return true;
+}
+
+int main(int, char**) {
+  tests();
+
+#if TEST_STD_VER >= 20
+  static_assert(tests());
+#endif
+
+  return 0;
+}
diff --git a/libcxx/test/std/containers/sequences/vector/vector.capacity/max_size.pass.cpp b/libcxx/test/std/containers/sequences/vector/vector.capacity/max_size.pass.cpp
index 33abaa04b12079..bc4d4e3fc24ad2 100644
--- a/libcxx/test/std/containers/sequences/vector/vector.capacity/max_size.pass.cpp
+++ b/libcxx/test/std/containers/sequences/vector/vector.capacity/max_size.pass.cpp
@@ -11,15 +11,34 @@
 // size_type max_size() const;
 
 #include <cassert>
+#include <cstdint>
 #include <limits>
 #include <type_traits>
 #include <vector>
 
+#include "min_allocator.h"
+#include "sized_allocator.h"
 #include "test_allocator.h"
 #include "test_macros.h"
 
+#if TEST_STD_VER >= 11
 
-TEST_CONSTEXPR_CXX20 bool test() {
+template <typename T, typename Alloc>
+TEST_CONSTEXPR_CXX20 void test(const std::vector<T, Alloc>& v) {
+  using Vector              = std::vector<T, Alloc>;
+  using alloc_traits        = typename Vector::__alloc_traits;
+  using size_type           = typename Vector::size_type;
+  using difference_type     = typename Vector::difference_type;
+  const size_type max_dist  = static_cast<size_type>(std::numeric_limits<difference_type>::max());
+  const size_type max_alloc = alloc_traits::max_size(v.get_allocator());
+  assert(v.max_size() <= max_dist);
+  assert(v.max_size() <= max_alloc);
+  LIBCPP_ASSERT(v.max_size() == std::min<size_type>(max_dist, max_alloc));
+}
+
+#endif
+
+TEST_CONSTEXPR_CXX20 bool tests() {
   {
     typedef limited_allocator<int, 10> A;
     typedef std::vector<int, A> C;
@@ -30,29 +49,48 @@ TEST_CONSTEXPR_CXX20 bool test() {
   {
     typedef limited_allocator<int, (std::size_t)-1> A;
     typedef std::vector<int, A> C;
-    const C::size_type max_dist =
-        static_cast<C::size_type>(std::numeric_limits<C::difference_type>::max());
+    const C::size_type max_dist = static_cast<C::size_type>(std::numeric_limits<C::difference_type>::max());
     C c;
     assert(c.max_size() <= max_dist);
     LIBCPP_ASSERT(c.max_size() == max_dist);
   }
   {
     typedef std::vector<char> C;
-    const C::size_type max_dist =
-        static_cast<C::size_type>(std::numeric_limits<C::difference_type>::max());
+    const C::size_type max_dist = static_cast<C::size_type>(std::numeric_limits<C::difference_type>::max());
     C c;
     assert(c.max_size() <= max_dist);
     assert(c.max_size() <= alloc_max_size(c.get_allocator()));
+    LIBCPP_ASSERT(c.max_size() == std::min(max_dist, alloc_max_size(c.get_allocator())));
+  }
+
+#if TEST_STD_VER >= 11
+
+  // Test with various allocators and diffrent size_type
+  {
+    test(std::vector<int>());
+    test(std::vector<short, std::allocator<short> >());
+    test(std::vector<unsigned, min_allocator<unsigned> >());
+    test(std::vector<char, test_allocator<char> >(test_allocator<char>(1)));
+    test(std::vector<std::size_t, other_allocator<std::size_t> >(other_allocator<std::size_t>(5)));
+    test(std::vector<int, sized_allocator<int, std::uint8_t, std::int8_t> >());
+    test(std::vector<int, sized_allocator<int, std::uint16_t, std::int16_t> >());
+    test(std::vector<int, sized_allocator<int, std::uint32_t, std::int32_t> >());
+    test(std::vector<int, sized_allocator<int, std::uint64_t, std::int64_t> >());
+    test(std::vector<int, limited_allocator<int, static_cast<std::size_t>(-1)> >());
+    test(std::vector<int, limited_allocator<int, static_cast<std::size_t>(-1) / 2> >());
+    test(std::vector<int, limited_allocator<int, static_cast<std::size_t>(-1) / 4> >());
   }
 
+#endif
+
   return true;
 }
 
 int main(int, char**) {
-  test();
+  tests();
 
 #if TEST_STD_VER > 17
-  static_assert(test());
+  static_assert(tests());
 #endif
 
   return 0;
diff --git a/libcxx/test/support/sized_allocator.h b/libcxx/test/support/sized_allocator.h
new file mode 100644
index 00000000000000..a877252e82962c
--- /dev/null
+++ b/libcxx/test/support/sized_allocator.h
@@ -0,0 +1,58 @@
+//===----------------------------------------------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef TEST_SUPPORT_SIZED_ALLOCATOR_H
+#define TEST_SUPPORT_SIZED_ALLOCATOR_H
+
+#include <cstddef>
+#include <limits>
+#include <memory>
+#include <new>
+
+#include "test_macros.h"
+
+template <typename T, typename SIZE_TYPE = std::size_t, typename DIFF_TYPE = std::ptrdiff_t>
+class sized_allocator {
+  template <typename U, typename Sz, typename Diff>
+  friend class sized_allocator;
+
+public:
+  using value_type                  = T;
+  using size_type                   = SIZE_TYPE;
+  using difference_type             = DIFF_TYPE;
+  using propagate_on_container_swap = std::true_type;
+
+  TEST_CONSTEXPR_CXX20 explicit sized_allocator(int d = 0) : data_(d) {}
+
+  template <typename U, typename Sz, typename Diff>
+  TEST_CONSTEXPR_CXX20 sized_allocator(const sized_allocator<U, Sz, Diff>& a) TEST_NOEXCEPT : data_(a.data_) {}
+
+  TEST_CONSTEXPR_CXX20 T* allocate(size_type n) {
+    if (n > max_size())
+      TEST_THROW(std::bad_array_new_length());
+    return std::allocator<T>().allocate(n);
+  }
+
+  TEST_CONSTEXPR_CXX20 void deallocate(T* p, size_type n) TEST_NOEXCEPT { std::allocator<T>().deallocate(p, n); }
+
+  TEST_CONSTEXPR size_type max_size() const TEST_NOEXCEPT {
+    return std::numeric_limits<size_type>::max() / sizeof(value_type);
+  }
+
+private:
+  int data_;
+
+  TEST_CONSTEXPR friend bool operator==(const sized_allocator& a, const sized_allocator& b) {
+    return a.data_ == b.data_;
+  }
+  TEST_CONSTEXPR friend bool operator!=(const sized_allocator& a, const sized_allocator& b) {
+    return a.data_ != b.data_;
+  }
+};
+
+#endif