From 18cf877a4712a91c2d8609dbaf9b38e5b33330ee Mon Sep 17 00:00:00 2001
From: Teon Banek <teon.banek@memgraph.io>
Date: Thu, 23 May 2019 12:28:36 +0200
Subject: [PATCH] Respect std::uses_allocator trait in utils::Allocator

Reviewers: mtomic, llugovic, mferencevic

Reviewed By: llugovic

Subscribers: pullbot

Differential Revision: https://phabricator.memgraph.io/D2081
---
 src/utils/memory.hpp        | 27 +++++++++++++++++
 tests/unit/utils_memory.cpp | 59 +++++++++++++++++++++++++++++++++++++
 2 files changed, 86 insertions(+)

diff --git a/src/utils/memory.hpp b/src/utils/memory.hpp
index 15d47d68d..df7c5caac 100644
--- a/src/utils/memory.hpp
+++ b/src/utils/memory.hpp
@@ -5,6 +5,7 @@
 #pragma once
 
 #include <cstddef>
+#include <memory>
 #include <type_traits>
 // Although <memory_resource> is in C++17, gcc libstdc++ still needs to
 // implement it fully. It should be available in the next major release
@@ -76,6 +77,11 @@ MemoryResource *NewDeleteResource() noexcept;
 /// The API of the Allocator is made such that it fits with the C++ STL
 /// requirements. It can be freely used in STL containers while relying on our
 /// implementations of MemoryResource.
+///
+/// Classes which have a member allocator_type typedef to this allocator will
+/// receive an instance of `this->GetMemoryResource()` upon construction if the
+/// outer container also uses this allocator. For concrete behaviour on how this
+/// is done refer to `std::uses_allocator` in C++ reference.
 template <class T>
 class Allocator {
  public:
@@ -100,6 +106,27 @@ class Allocator {
         memory_->Allocate(count_elements * sizeof(T), alignof(T)));
   }
 
+  template <class U, class... TArgs>
+  void construct(U *ptr, TArgs &&... args) {
+    if constexpr (std::uses_allocator_v<U, Allocator>) {
+      if constexpr (std::is_constructible_v<U, std::allocator_arg_t,
+                                            MemoryResource *, TArgs...>) {
+        ::new (ptr)
+            U(std::allocator_arg, memory_, std::forward<TArgs>(args)...);
+      } else if constexpr (std::is_constructible_v<U, TArgs...,
+                                                   MemoryResource *>) {
+        ::new (ptr) U(std::forward<TArgs>(args)..., memory_);
+      } else {
+        static_assert(!std::uses_allocator_v<U, Allocator>,
+                      "Class declares std::uses_allocator but has no valid "
+                      "constructor overload. Refer to 'Uses-allocator "
+                      "construction' rules in C++ reference.");
+      }
+    } else {
+      ::new (ptr) U(std::forward<TArgs>(args)...);
+    }
+  }
+
   void deallocate(T *p, size_t count_elements) {
     memory_->Deallocate(p, count_elements * sizeof(T), alignof(T));
   }
diff --git a/tests/unit/utils_memory.cpp b/tests/unit/utils_memory.cpp
index 9d0c5df6d..e4570a2e4 100644
--- a/tests/unit/utils_memory.cpp
+++ b/tests/unit/utils_memory.cpp
@@ -146,3 +146,62 @@ TEST(MonotonicBufferResource, AllocationWithSizeOverflow) {
   mem.Allocate(1, 1);
   EXPECT_THROW(mem.Allocate(max_size, 4), std::bad_alloc);
 }
+
+// NOLINTNEXTLINE(hicpp-special-member-functions)
+class ContainerWithAllocatorLast final {
+ public:
+  using allocator_type = utils::Allocator<int>;
+
+  ContainerWithAllocatorLast() = default;
+  explicit ContainerWithAllocatorLast(int value) : value_(value) {}
+  ContainerWithAllocatorLast(int value, utils::MemoryResource *memory)
+      : memory_(memory), value_(value) {}
+
+  ContainerWithAllocatorLast(const ContainerWithAllocatorLast &other)
+      : value_(other.value_) {}
+  ContainerWithAllocatorLast(const ContainerWithAllocatorLast &other,
+                             utils::MemoryResource *memory)
+      : memory_(memory), value_(other.value_) {}
+
+  utils::MemoryResource *memory_{nullptr};
+  int value_{0};
+};
+
+// NOLINTNEXTLINE(hicpp-special-member-functions)
+class ContainerWithAllocatorFirst final {
+ public:
+  using allocator_type = utils::Allocator<int>;
+
+  ContainerWithAllocatorFirst() = default;
+  explicit ContainerWithAllocatorFirst(int value) : value_(value) {}
+  ContainerWithAllocatorFirst(std::allocator_arg_t,
+                              utils::MemoryResource *memory, int value)
+      : memory_(memory), value_(value) {}
+
+  ContainerWithAllocatorFirst(const ContainerWithAllocatorFirst &other)
+      : value_(other.value_) {}
+  ContainerWithAllocatorFirst(std::allocator_arg_t,
+                              utils::MemoryResource *memory,
+                              const ContainerWithAllocatorFirst &other)
+      : memory_(memory), value_(other.value_) {}
+
+  utils::MemoryResource *memory_{nullptr};
+  int value_{0};
+};
+
+template <class T>
+class AllocatorTest : public ::testing::Test {};
+
+using ContainersWithAllocators =
+    ::testing::Types<ContainerWithAllocatorLast, ContainerWithAllocatorFirst>;
+
+TYPED_TEST_CASE(AllocatorTest, ContainersWithAllocators);
+
+TYPED_TEST(AllocatorTest, PropagatesToStdUsesAllocator) {
+  std::vector<TypeParam, utils::Allocator<TypeParam>> vec(
+      utils::NewDeleteResource());
+  vec.emplace_back(42);
+  const auto &c = vec.front();
+  EXPECT_EQ(c.value_, 42);
+  EXPECT_EQ(c.memory_, utils::NewDeleteResource());
+}