Add support for custom allocators

format.h

@@ -238,8 +238,8 @@ inline T *make_ptr(T *ptr, std::size_t) { return ptr; }
 
 // A simple array for POD types with the first SIZE elements stored in
 // the object itself. It supports a subset of std::vector's operations.
-template <typename T, std::size_t SIZE>
-class Array {
+template <typename T, std::size_t SIZE, typename Allocator = std::allocator<T> >
+class Array : private Allocator {
  private:
   std::size_t size_;
   std::size_t capacity_;
@@ -250,7 +250,7 @@ class Array {
 
   // Free memory allocated by the array.
   void free() {
-    if (ptr_ != data_) delete [] ptr_;
+    if (ptr_ != data_) this->deallocate(ptr_, capacity_);
   }
 
   // Move data from other to this array.
@@ -271,7 +271,8 @@ class Array {
   FMT_DISALLOW_COPY_AND_ASSIGN(Array);
 
  public:
-  explicit Array() : size_(0), capacity_(SIZE), ptr_(data_) {}
+  explicit Array(const Allocator &alloc = Allocator())
+      : Allocator(alloc), size_(0), capacity_(SIZE), ptr_(data_) {}
   ~Array() { free(); }
 
 #if FMT_USE_RVALUE_REFERENCES
@@ -293,6 +294,9 @@ class Array {
   // Returns the capacity of this array.
   std::size_t capacity() const { return capacity_; }
 
+  // Returns a copy of the allocator associated with this array.
+  Allocator get_allocator() const { return *this; }
+
   // Resizes the array. If T is a POD type new elements are not initialized.
   void resize(std::size_t new_size) {
     if (new_size > capacity_)
@@ -321,18 +325,25 @@ class Array {
   const T &operator[](std::size_t index) const { return ptr_[index]; }
 };
 
-template <typename T, std::size_t SIZE>
-void Array<T, SIZE>::grow(std::size_t size) {
-  capacity_ = (std::max)(size, capacity_ + capacity_ / 2);
-  T *p = new T[capacity_];
-  std::copy(ptr_, ptr_ + size_, make_ptr(p, capacity_));
-  if (ptr_ != data_)
-    delete [] ptr_;
-  ptr_ = p;
+template <typename T, std::size_t SIZE, typename Allocator>
+void Array<T, SIZE, Allocator>::grow(std::size_t size) {
+  std::size_t new_capacity = (std::max)(size, capacity_ + capacity_ / 2);
+  T *new_ptr = this->allocate(new_capacity);
+  // The following code doesn't throw, so the raw pointer above doesn't leak.
+  std::copy(ptr_, ptr_ + size_, make_ptr(new_ptr, new_capacity));
+  std::size_t old_capacity = capacity_;
+  T *old_ptr = ptr_;
+  capacity_ = new_capacity;
+  ptr_ = new_ptr;
+  // deallocate may throw (at least in principle), but it doesn't matter since
+  // the array already uses the new storage and will deallocate it in case
+  // of exception.
+  if (old_ptr != data_)
+    this->deallocate(old_ptr, old_capacity);
 }
 
-template <typename T, std::size_t SIZE>
-void Array<T, SIZE>::append(const T *begin, const T *end) {
+template <typename T, std::size_t SIZE, typename Allocator>
+void Array<T, SIZE, Allocator>::append(const T *begin, const T *end) {
   std::ptrdiff_t num_elements = end - begin;
   if (size_ + num_elements > capacity_)
     grow(size_ + num_elements);

test/format-test.cc

@@ -34,14 +34,15 @@
 #include <iomanip>
 #include <memory>
 #include <sstream>
+#include <stdint.h>
+
+#include "gmock/gmock.h"
 
 // Include format.cc instead of format.h to test implementation-specific stuff.
 #include "format.h"
 #include "util.h"
 #include "gtest-extra.h"
 
-#include <stdint.h>
-
 #if defined(_WIN32) && !defined(__MINGW32__)
 // Fix MSVC warning about "unsafe" fopen.
 FILE *safe_fopen(const char *filename, const char *mode) {
@@ -66,6 +67,8 @@ using fmt::Writer;
 using fmt::WWriter;
 using fmt::pad;
 
+using testing::Return;
+
 namespace {
 
 // Checks if writing value to BasicWriter<Char> produces the same result
@@ -271,6 +274,97 @@ TEST(ArrayTest, AppendAllocatesEnoughStorage) {
   EXPECT_EQ(19u, array.capacity());
 }
 
+template <typename T>
+class MockAllocator {
+ public:
+  typedef T value_type;
+  MOCK_METHOD1_T(allocate, T* (std::size_t n));
+  MOCK_METHOD2_T(deallocate, void (T* p, std::size_t n));
+};
+
+template <typename Allocator>
+class AllocatorRef {
+ private:
+  Allocator *alloc_;
+
+ public:
+  typedef typename Allocator::value_type value_type;
+
+  explicit AllocatorRef(Allocator *alloc = 0) : alloc_(alloc) {}
+
+  Allocator *get() const { return alloc_; }
+
+  value_type* allocate(std::size_t n) { return alloc_->allocate(n); }
+  void deallocate(value_type* p, std::size_t n) { alloc_->deallocate(p, n); }
+};
+
+void CheckForwarding(
+    MockAllocator<int> &alloc, AllocatorRef< MockAllocator<int> > &ref) {
+  int mem;
+  // Check if value_type is properly defined.
+  AllocatorRef< MockAllocator<int> >::value_type *ptr = &mem;
+  // Check forwarding.
+  EXPECT_CALL(alloc, allocate(42)).WillOnce(Return(ptr));
+  ref.allocate(42);
+  EXPECT_CALL(alloc, deallocate(ptr, 42));
+  ref.deallocate(ptr, 42);
+}
+
+TEST(AllocatorTest, AllocatorRef) {
+  testing::StrictMock< MockAllocator<int> > alloc;
+  typedef AllocatorRef< MockAllocator<int> > TestAllocatorRef;
+  TestAllocatorRef ref(&alloc);
+  // Check if AllocatorRef forwards to the underlying allocator.
+  CheckForwarding(alloc, ref);
+  TestAllocatorRef ref2(ref);
+  CheckForwarding(alloc, ref2);
+  TestAllocatorRef ref3;
+  EXPECT_EQ(0, ref3.get());
+  ref3 = ref;
+  CheckForwarding(alloc, ref3);
+}
+
+TEST(ArrayTest, Allocator) {
+  typedef AllocatorRef< MockAllocator<char> > TestAllocator;
+  Array<char, 10, TestAllocator> array;
+  EXPECT_EQ(0, array.get_allocator().get());
+  testing::StrictMock< MockAllocator<char> > alloc;
+  char mem;
+  {
+    Array<char, 10, TestAllocator> array2((TestAllocator(&alloc)));
+    EXPECT_EQ(&alloc, array2.get_allocator().get());
+    std::size_t size = 2 * fmt::internal::INLINE_BUFFER_SIZE;
+    EXPECT_CALL(alloc, allocate(size)).WillOnce(Return(&mem));
+    array2.reserve(size);
+    EXPECT_CALL(alloc, deallocate(&mem, size));
+  }
+}
+
+TEST(ArrayTest, DeallocateException) {
+  typedef AllocatorRef< MockAllocator<char> > TestAllocator;
+  testing::StrictMock< MockAllocator<char> > alloc;
+  Array<char, 10, TestAllocator> array((TestAllocator(&alloc)));
+  std::size_t size = 2 * fmt::internal::INLINE_BUFFER_SIZE;
+  std::vector<char> mem(size);
+  {
+    EXPECT_CALL(alloc, allocate(size)).WillOnce(Return(&mem[0]));
+    array.resize(size);
+    std::fill(&array[0], &array[0] + size, 'x');
+  }
+  std::vector<char> mem2(2 * size);
+  {
+    EXPECT_CALL(alloc, allocate(2 * size)).WillOnce(Return(&mem2[0]));
+    std::exception e;
+    EXPECT_CALL(alloc, deallocate(&mem[0], size)).WillOnce(testing::Throw(e));
+    EXPECT_THROW(array.reserve(2 * size), std::exception);
+    EXPECT_EQ(&mem2[0], &array[0]);
+    // Check that the data has been copied.
+    for (std::size_t i = 0; i < size; ++i)
+      EXPECT_EQ('x', array[i]);
+  }
+  EXPECT_CALL(alloc, deallocate(&mem2[0], 2 * size));
+}
+
 TEST(WriterTest, Ctor) {
   Writer w;
   EXPECT_EQ(0u, w.size());