From 9c7f5ea10723488df5061d4381c85e523d0fa3e1 Mon Sep 17 00:00:00 2001
From: Vladislav Shchapov <vladislav@shchapov.ru>
Date: Fri, 12 Jan 2024 23:31:15 +0500
Subject: [PATCH] Implement allocator_max_size.

Signed-off-by: Vladislav Shchapov <vladislav@shchapov.ru>
---
 include/fmt/format.h     | 42 +++++++++++++++++++++++++++++++++++++---
 test/format-impl-test.cc |  7 +++++++
 test/format-test.cc      | 35 +++++++++++++++++++--------------
 3 files changed, 67 insertions(+), 17 deletions(-)

diff --git a/include/fmt/format.h b/include/fmt/format.h
index 6ffba5dd030c4..dca01e4c966da 100644
--- a/include/fmt/format.h
+++ b/include/fmt/format.h
@@ -517,6 +517,42 @@ struct allocator_size<Allocator, void_t<typename Allocator::size_type>> {
   using type = typename Allocator::size_type;
 };
 
+#if FMT_GCC_VERSION || FMT_CLANG_VERSION
+#  pragma GCC diagnostic push
+#  pragma GCC diagnostic ignored "-Wdeprecated"
+#endif
+
+// MSVC 2015 does not support of using a dependent decltype in the template
+// argument of partial specialization of a class template. Use legacy
+// implementation instead of ``void_t<decltype(std::declval<const
+// T>().max_size())>``.
+template <typename T> class has_max_size {
+  template <typename> static auto test(...) -> std::false_type;
+  template <typename U>
+  static auto test(int)
+      -> decltype(std::declval<U>().max_size(), std::true_type());
+
+ public:
+  static constexpr bool value =
+      std::is_same<decltype(test<const T>(0)), std::true_type>::value;
+};
+
+template <typename Allocator, FMT_ENABLE_IF(has_max_size<Allocator>::value)>
+FMT_CONSTEXPR20 auto allocator_max_size(const Allocator& a) ->
+    typename allocator_size<Allocator>::type {
+  return a.max_size();
+}
+template <typename Allocator, FMT_ENABLE_IF(!has_max_size<Allocator>::value)>
+FMT_CONSTEXPR20 auto allocator_max_size(const Allocator&) ->
+    typename allocator_size<Allocator>::type {
+  return max_value<typename allocator_size<Allocator>::type>() /
+         sizeof(typename Allocator::value_type);
+}
+
+#if FMT_GCC_VERSION || FMT_CLANG_VERSION
+#  pragma GCC diagnostic pop
+#endif
+
 template <typename Char, typename InputIt>
 auto copy_str(InputIt begin, InputIt end, appender out) -> appender {
   get_container(out).append(begin, end);
@@ -902,6 +938,8 @@ template <typename T, size_t SIZE = inline_buffer_size,
 class basic_memory_buffer : public detail::buffer<T> {
  private:
   T store_[SIZE];
+  static_assert(std::is_same<typename Allocator::value_type, T>::value,
+                "Invalid allocator value_type");
 
   // Don't inherit from Allocator to avoid generating type_info for it.
   FMT_NO_UNIQUE_ADDRESS Allocator alloc_;
@@ -915,9 +953,7 @@ class basic_memory_buffer : public detail::buffer<T> {
   static FMT_CONSTEXPR20 void grow(detail::buffer<T>& buf, size_t size) {
     detail::abort_fuzzing_if(size > 5000);
     auto& self = static_cast<basic_memory_buffer&>(buf);
-    constexpr size_t max_size =
-        detail::max_value<typename detail::allocator_size<Allocator>::type>() /
-        sizeof(T);
+    const size_t max_size = detail::allocator_max_size(self.alloc_);
     size_t old_capacity = buf.capacity();
     size_t new_capacity = old_capacity + old_capacity / 2;
     if (size > new_capacity)
diff --git a/test/format-impl-test.cc b/test/format-impl-test.cc
index eda1f23958024..0ad263b7e4303 100644
--- a/test/format-impl-test.cc
+++ b/test/format-impl-test.cc
@@ -20,6 +20,13 @@ using fmt::detail::bigint;
 using fmt::detail::fp;
 using fmt::detail::max_value;
 
+struct Empty {};
+struct MaxSize {
+  std::size_t max_size() const { return 0; }
+};
+static_assert(!fmt::detail::has_max_size<Empty>::value, "");
+static_assert(fmt::detail::has_max_size<MaxSize>::value, "");
+
 static_assert(!std::is_copy_constructible<bigint>::value, "");
 static_assert(!std::is_copy_assignable<bigint>::value, "");
 
diff --git a/test/format-test.cc b/test/format-test.cc
index a76c44cf76505..b38a3bd439673 100644
--- a/test/format-test.cc
+++ b/test/format-test.cc
@@ -413,30 +413,37 @@ TEST(memory_buffer_test, exception_in_deallocate) {
   EXPECT_CALL(alloc, deallocate(&mem2[0], 2 * size));
 }
 
-class smol_allocator : public std::allocator<char> {
+template <typename Allocator, size_t MaxSize>
+class max_size_allocator : public Allocator {
  public:
-  using size_type = unsigned char;
-
-  auto allocate(size_t n) -> value_type* {
-    if (n > fmt::detail::max_value<size_type>())
+  using typename Allocator::value_type;
+  size_t max_size() const noexcept { return MaxSize; }
+  value_type* allocate(size_t n) {
+    if (n > max_size()) {
       throw std::length_error("size > max_size");
-    return std::allocator<char>::allocate(n);
+    }
+    return std::allocator_traits<Allocator>::allocate(
+        *static_cast<Allocator*>(this), n);
   }
   void deallocate(value_type* p, size_t n) {
-    std::allocator<char>::deallocate(p, n);
+    std::allocator_traits<Allocator>::deallocate(*static_cast<Allocator*>(this),
+                                                 p, n);
   }
 };
 
 TEST(memory_buffer_test, max_size_allocator) {
-  basic_memory_buffer<char, 10, smol_allocator> buffer;
-  buffer.resize(200);
-  // new_capacity = 200 + 200/2 = 300 > 256
-  buffer.resize(255);  // Shouldn't throw.
+  // 160 = 128 + 32
+  using test_allocator = max_size_allocator<std::allocator<char>, 160>;
+  basic_memory_buffer<char, 10, test_allocator> buffer;
+  buffer.resize(128);
+  // new_capacity = 128 + 128/2 = 192 > 160
+  buffer.resize(160);  // Shouldn't throw.
 }
 
 TEST(memory_buffer_test, max_size_allocator_overflow) {
-  basic_memory_buffer<char, 10, smol_allocator> buffer;
-  EXPECT_THROW(buffer.resize(256), std::exception);
+  using test_allocator = max_size_allocator<std::allocator<char>, 160>;
+  basic_memory_buffer<char, 10, test_allocator> buffer;
+  EXPECT_THROW(buffer.resize(161), std::exception);
 }
 
 TEST(format_test, exception_from_lib) {
@@ -2152,7 +2159,7 @@ TEST(format_int_test, format_int) {
   EXPECT_EQ(fmt::format_int(42ul).str(), "42");
   EXPECT_EQ(fmt::format_int(-42l).str(), "-42");
   EXPECT_EQ(fmt::format_int(42ull).str(), "42");
-  EXPECT_EQ(fmt::format_int(-42ll).str(), "-42");\
+  EXPECT_EQ(fmt::format_int(-42ll).str(), "-42");
   EXPECT_EQ(fmt::format_int(max_value<int64_t>()).str(),
             std::to_string(max_value<int64_t>()));
 }