helper methods for over-aligned allocation with C++ allocators

Summary:
This diff adds helper methods for allocating
possibly-over-aligned types using C++ allocators.  It uses one of three
strategies:

* REBINDING - When there is a type that is not over-aligned with the
requested alignment, the allocator is rebound to that a type with the
correct alignment (this allows allocating 1 char with alignment 8,
for example).

* BYPASSING - When the requested alignment is larger than max_align_t
but the standard allocator is in use, the allocator is bypassed and
aligned_malloc/aligned_free are used directly to manage the allocation.

* OVER-ALLOCATING - When the type is truly over-aligned and the allocator
is custom, extra memory is requested and a delta is stored before the
returned memory to allow the original base location to be recovered.

Reviewed By: shixiao

Differential Revision: D8592767

fbshipit-source-id: 487b5ff7f7f6b1b29a3f2a80d07129afd50bda24

folly/Memory.h

@@ -77,6 +77,123 @@ inline void aligned_free(void* aligned_ptr) {
 
 #endif
 
+// Works like std::allocator_traits<Alloc>::allocate, but handles
+// over-aligned types.  Feel free to manually specify any power of two as
+// the Align template arg.  Must be matched with deallocateOverAligned.
+// allocationBytesForOverAligned will give you the number of bytes that
+// this function actually requests.
+template <
+    typename Alloc,
+    size_t kAlign = alignof(typename std::allocator_traits<Alloc>::value_type)>
+typename std::allocator_traits<Alloc>::pointer allocateOverAligned(
+    Alloc const& alloc,
+    size_t n) {
+  static_assert((kAlign & (kAlign - 1)) == 0, "Align must be a power of 2");
+
+  using AllocTraits = std::allocator_traits<Alloc>;
+  using T = typename AllocTraits::value_type;
+
+  constexpr size_t kBaseAlign = constexpr_min(kAlign, alignof(max_align_t));
+  using BaseType = std::aligned_storage_t<kBaseAlign, kBaseAlign>;
+  using BaseAllocTraits =
+      typename AllocTraits::template rebind_traits<BaseType>;
+  using BaseAlloc = typename BaseAllocTraits::allocator_type;
+  static_assert(
+      sizeof(BaseType) == kBaseAlign && alignof(BaseType) == kBaseAlign, "");
+
+  void* rv;
+  if (kAlign > kBaseAlign && std::is_same<Alloc, std::allocator<T>>::value) {
+    // allocating as BaseType isn't sufficient, but it is safe to bypass
+    // std::allocator
+    rv = aligned_malloc(n * sizeof(T), kAlign);
+  } else {
+    BaseAlloc a(alloc);
+    size_t quanta = (n * sizeof(T) + kBaseAlign - 1) / kBaseAlign;
+    if (kAlign <= kBaseAlign) {
+      // not actually over-aligned, happy path
+      rv = static_cast<void*>(
+          std::addressof(*BaseAllocTraits::allocate(a, quanta)));
+    } else {
+      static_assert(kAlign == kBaseAlign || kBaseAlign >= sizeof(size_t), "");
+
+      // if we give ourselves kAlign extra bytes, then since kBaseAlign
+      // divides kAlign we can meet alignment while getting a prefix of
+      // one kBaseAlign.  We wouldn't be over-aligned unless kBaseAlign
+      // was already as big as possible, so this is enough space to store
+      // a size_t.
+      char* base = reinterpret_cast<char*>(std::addressof(
+          *BaseAllocTraits::allocate(a, quanta + kAlign / kBaseAlign)));
+      size_t delta =
+          kAlign - (reinterpret_cast<uintptr_t>(base) & (kAlign - 1));
+      rv = static_cast<void*>(base + delta);
+      static_cast<size_t*>(rv)[-1] = delta;
+    }
+  }
+  return std::pointer_traits<typename AllocTraits::pointer>::pointer_to(
+      *static_cast<T*>(rv));
+}
+
+template <
+    typename Alloc,
+    size_t kAlign = alignof(typename std::allocator_traits<Alloc>::value_type)>
+void deallocateOverAligned(
+    Alloc const& alloc,
+    typename std::allocator_traits<Alloc>::pointer ptr,
+    size_t n) {
+  static_assert((kAlign & (kAlign - 1)) == 0, "Align must be a power of 2");
+
+  using AllocTraits = std::allocator_traits<Alloc>;
+  using T = typename AllocTraits::value_type;
+
+  constexpr size_t kBaseAlign = constexpr_min(kAlign, alignof(max_align_t));
+  using BaseType = std::aligned_storage_t<kBaseAlign, kBaseAlign>;
+  using BaseAllocTraits =
+      typename AllocTraits::template rebind_traits<BaseType>;
+  using BaseAlloc = typename BaseAllocTraits::allocator_type;
+  static_assert(
+      sizeof(BaseType) == kBaseAlign && alignof(BaseType) == kBaseAlign, "");
+
+  void* raw = static_cast<void*>(std::addressof(*ptr));
+  if (kAlign > kBaseAlign && std::is_same<Alloc, std::allocator<T>>::value) {
+    aligned_free(raw);
+  } else {
+    BaseAlloc a(alloc);
+    size_t quanta = (n * sizeof(T) + kBaseAlign - 1) / kBaseAlign;
+    if (kAlign > kBaseAlign) {
+      size_t delta = static_cast<size_t*>(raw)[-1];
+      raw = static_cast<void*>(static_cast<char*>(raw) - delta);
+      quanta += kAlign / kBaseAlign;
+    }
+    BaseAllocTraits::deallocate(
+        a,
+        std::pointer_traits<typename BaseAllocTraits::pointer>::pointer_to(
+            *static_cast<BaseType*>(raw)),
+        quanta);
+  }
+}
+
+template <
+    typename Alloc,
+    size_t kAlign = alignof(typename std::allocator_traits<Alloc>::value_type)>
+size_t allocationBytesForOverAligned(size_t n) {
+  static_assert((kAlign & (kAlign - 1)) == 0, "Align must be a power of 2");
+
+  using AllocTraits = std::allocator_traits<Alloc>;
+  using T = typename AllocTraits::value_type;
+
+  constexpr size_t kBaseAlign = constexpr_min(kAlign, alignof(max_align_t));
+
+  if (kAlign > kBaseAlign && std::is_same<Alloc, std::allocator<T>>::value) {
+    return n * sizeof(T);
+  } else {
+    size_t quanta = (n * sizeof(T) + kBaseAlign - 1) / kBaseAlign;
+    if (kAlign > kBaseAlign) {
+      quanta += kAlign / kBaseAlign;
+    }
+    return quanta * kBaseAlign;
+  }
+}
+
 /**
  * For exception safety and consistency with make_shared. Erase me when
  * we have std::make_unique().

folly/test/MemoryTest.cpp

@@ -17,6 +17,7 @@
 #include <folly/Memory.h>
 
 #include <limits>
+#include <memory>
 #include <type_traits>
 #include <utility>
 
@@ -316,3 +317,79 @@ TEST(enable_shared_from_this, compatible_with_std_enable_shared_from_this) {
   test_enable_shared_from_this(std::make_shared<C_folly>());
   test_enable_shared_from_this(std::make_shared<C_folly const>());
 }
+
+template <typename T>
+struct ExpectingAlloc {
+  using value_type = T;
+
+  ExpectingAlloc(std::size_t expectedSize, std::size_t expectedCount)
+      : expectedSize_{expectedSize}, expectedCount_{expectedCount} {}
+
+  template <class U>
+  explicit ExpectingAlloc(ExpectingAlloc<U> const& other) noexcept
+      : expectedSize_{other.expectedSize_},
+        expectedCount_{other.expectedCount_} {}
+
+  T* allocate(std::size_t n) {
+    EXPECT_EQ(expectedSize_, sizeof(T));
+    EXPECT_EQ(expectedSize_, alignof(T));
+    EXPECT_EQ(expectedCount_, n);
+    return std::allocator<T>{}.allocate(n);
+  }
+
+  void deallocate(T* p, std::size_t n) {
+    std::allocator<T>{}.deallocate(p, n);
+  }
+
+  std::size_t expectedSize_;
+  std::size_t expectedCount_;
+};
+
+struct alignas(64) OverAlignedType {
+  std::array<char, 64> raw_;
+};
+
+TEST(allocateOverAligned, notActuallyOver) {
+  // allocates 6 bytes with alignment 4, should get turned into an
+  // allocation of 2 elements of something of size and alignment 4
+  ExpectingAlloc<char> a(4, 2);
+  auto p = folly::allocateOverAligned<decltype(a), 4>(a, 6);
+  EXPECT_EQ((reinterpret_cast<uintptr_t>(p) % 4), 0);
+  folly::deallocateOverAligned<decltype(a), 4>(a, p, 6);
+  EXPECT_EQ((folly::allocationBytesForOverAligned<decltype(a), 4>(6)), 8);
+}
+
+TEST(allocateOverAligned, manualOverStdAlloc) {
+  // allocates 6 bytes with alignment 64 using std::allocator, which will
+  // result in a call to aligned_malloc underneath.  We free one directly
+  // to check that this is not the padding path
+  std::allocator<short> a;
+  auto p = folly::allocateOverAligned<decltype(a), 64>(a, 3);
+  auto p2 = folly::allocateOverAligned<decltype(a), 64>(a, 3);
+  EXPECT_EQ((reinterpret_cast<uintptr_t>(p) % 64), 0);
+  folly::deallocateOverAligned<decltype(a), 64>(a, p, 3);
+  aligned_free(p2);
+  EXPECT_EQ((folly::allocationBytesForOverAligned<decltype(a), 64>(3)), 6);
+}
+
+TEST(allocateOverAligned, manualOverCustomAlloc) {
+  // allocates 6 byte with alignment 64 using non-standard allocator, which
+  // will result in an allocation of 64 + alignof(max_align_t) underneath.
+  ExpectingAlloc<short> a(
+      alignof(folly::max_align_t), 64 / alignof(folly::max_align_t) + 1);
+  auto p = folly::allocateOverAligned<decltype(a), 64>(a, 3);
+  EXPECT_EQ((reinterpret_cast<uintptr_t>(p) % 64), 0);
+  folly::deallocateOverAligned<decltype(a), 64>(a, p, 3);
+  EXPECT_EQ(
+      (folly::allocationBytesForOverAligned<decltype(a), 64>(3)),
+      64 + alignof(folly::max_align_t));
+}
+
+TEST(allocateOverAligned, defaultOverCustomAlloc) {
+  ExpectingAlloc<OverAlignedType> a(
+      alignof(folly::max_align_t), 128 / alignof(folly::max_align_t));
+  auto p = folly::allocateOverAligned(a, 1);
+  EXPECT_EQ((reinterpret_cast<uintptr_t>(p) % 64), 0);
+  folly::deallocateOverAligned(a, p, 1);
+  EXPECT_EQ(folly::allocationBytesForOverAligned<decltype(a)>(1), 128);
+}