Wrap all stl containers with a marl::StlAllocator

CMakeLists.txt

@@ -44,6 +44,7 @@ option_if_not_defined(MARL_MSAN "Build marl with memory sanitizer" OFF)
 option_if_not_defined(MARL_TSAN "Build marl with thread sanitizer" OFF)
 option_if_not_defined(MARL_INSTALL "Create marl install target" OFF)
 option_if_not_defined(MARL_FULL_BENCHMARK "Benchmark all cores and a range of tasks" OFF)
+option_if_not_defined(MARL_DEBUG_ENABLED "Enable debug checks even in release builds" OFF)
 
 ###########################################################
 # Directories
@@ -187,6 +188,10 @@ function(marl_set_target_options target)
         target_link_libraries(${target} PUBLIC "-fsanitize=thread")
     endif()
 
+    if (MARL_DEBUG_ENABLED)
+        target_compile_definitions(${target} PRIVATE "MARL_DEBUG_ENABLED=1")
+    endif()
+
     target_include_directories(${target} PUBLIC $<BUILD_INTERFACE:${MARL_INCLUDE_DIR}>)
 endfunction(marl_set_target_options)
 

examples/primes.cpp

@@ -20,6 +20,8 @@
 #include "marl/thread.h"
 #include "marl/ticket.h"
 
+#include <vector>
+
 #include <math.h>
 
 // searchMax defines the upper limit on primes to find.

include/marl/containers.h

@@ -22,18 +22,66 @@
 #include <cstddef>    // size_t
 #include <utility>    // std::move
 
+#include <deque>
+#include <map>
+#include <set>
+#include <unordered_map>
+#include <unordered_set>
+
 namespace marl {
 namespace containers {
 
+////////////////////////////////////////////////////////////////////////////////
+// STL wrappers
+// STL containers that use a marl::StlAllocator backed by a marl::Allocator.
+// Note: These may be re-implemented to optimize for marl's usage cases.
+// See: https://github.com/google/marl/issues/129
+////////////////////////////////////////////////////////////////////////////////
+template <typename T>
+using deque = std::deque<T, StlAllocator<T>>;
+
+template <typename K, typename V, typename C = std::less<K>>
+using map = std::map<K, V, C, StlAllocator<std::pair<const K, V>>>;
+
+template <typename K, typename C = std::less<K>>
+using set = std::set<K, C, StlAllocator<K>>;
+
+template <typename K,
+          typename V,
+          typename H = std::hash<K>,
+          typename E = std::equal_to<K>>
+using unordered_map =
+    std::unordered_map<K, V, H, E, StlAllocator<std::pair<const K, V>>>;
+
+template <typename K, typename H = std::hash<K>, typename E = std::equal_to<K>>
+using unordered_set = std::unordered_set<K, H, E, StlAllocator<K>>;
+
+// take() takes and returns the front value from the deque.
+template <typename T>
+inline T take(deque<T>& queue) {
+  auto out = std::move(queue.front());
+  queue.pop_front();
+  return out;
+}
+
+// take() takes and returns the first value from the unordered_set.
+template <typename T, typename H, typename E>
+inline T take(unordered_set<T, H, E>& set) {
+  auto it = set.begin();
+  auto out = std::move(*it);
+  set.erase(it);
+  return out;
+}
+
 ////////////////////////////////////////////////////////////////////////////////
 // vector<T, BASE_CAPACITY>
 ////////////////////////////////////////////////////////////////////////////////
 
 // vector is a container of contiguously stored elements.
-// Unlike std::vector, marl::containers::vector keeps the first BASE_CAPACITY
-// elements internally, which will avoid dynamic heap allocations.
-// Once the vector exceeds BASE_CAPACITY elements, vector will allocate storage
-// from the heap.
+// Unlike std::vector, marl::containers::vector keeps the first
+// BASE_CAPACITY elements internally, which will avoid dynamic heap
+// allocations. Once the vector exceeds BASE_CAPACITY elements, vector will
+// allocate storage from the heap.
 template <typename T, int BASE_CAPACITY>
 class vector {
  public:
@@ -74,6 +122,10 @@ class vector {
   inline size_t cap() const;
   inline void resize(size_t n);
   inline void reserve(size_t n);
+  inline T* data();
+  inline const T* data() const;
+
+  Allocator* const allocator;
 
  private:
   using TStorage = typename marl::aligned_storage<sizeof(T), alignof(T)>::type;
@@ -82,7 +134,6 @@ class vector {
 
   inline void free();
 
-  Allocator* const allocator;
   size_t count = 0;
   size_t capacity = BASE_CAPACITY;
   TStorage buffer[BASE_CAPACITY];
@@ -272,6 +323,16 @@ void vector<T, BASE_CAPACITY>::reserve(size_t n) {
   }
 }
 
+template <typename T, int BASE_CAPACITY>
+T* vector<T, BASE_CAPACITY>::data() {
+  return elements;
+}
+
+template <typename T, int BASE_CAPACITY>
+const T* vector<T, BASE_CAPACITY>::data() const {
+  return elements;
+}
+
 template <typename T, int BASE_CAPACITY>
 void vector<T, BASE_CAPACITY>::free() {
   for (size_t i = 0; i < count; i++) {

include/marl/memory.h

@@ -27,6 +27,9 @@
 
 namespace marl {
 
+template <typename T>
+struct StlAllocator;
+
 // pageSize() returns the size in bytes of a virtual memory page for the host
 // system.
 size_t pageSize();
@@ -36,6 +39,19 @@ inline T alignUp(T val, T alignment) {
   return alignment * ((val + alignment - 1) / alignment);
 }
 
+// aligned_storage() is a replacement for std::aligned_storage that isn't busted
+// on older versions of MSVC.
+template <size_t SIZE, size_t ALIGNMENT>
+struct aligned_storage {
+  struct alignas(ALIGNMENT) type {
+    unsigned char data[SIZE];
+  };
+};
+
+///////////////////////////////////////////////////////////////////////////////
+// Allocation
+///////////////////////////////////////////////////////////////////////////////
+
 // Allocation holds the result of a memory allocation from an Allocator.
 struct Allocation {
   // Intended usage of the allocation. Used for allocation trackers.
@@ -45,6 +61,7 @@ struct Allocation {
     Create,  // Allocator::create(), make_unique(), make_shared()
     Vector,  // marl::containers::vector<T>
     List,    // marl::containers::list<T>
+    Stl,     // marl::StlAllocator
     Count,   // Not intended to be used as a usage type - used for upper bound.
   };
 
@@ -60,6 +77,10 @@ struct Allocation {
   Request request;      // Request used for the allocation.
 };
 
+///////////////////////////////////////////////////////////////////////////////
+// Allocator
+///////////////////////////////////////////////////////////////////////////////
+
 // Allocator is an interface to a memory allocator.
 // Marl provides a default implementation with Allocator::Default.
 class Allocator {
@@ -183,14 +204,9 @@ std::shared_ptr<T> Allocator::make_shared(ARGS&&... args) {
   return std::shared_ptr<T>(reinterpret_cast<T*>(alloc.ptr), Deleter{this});
 }
 
-// aligned_storage() is a replacement for std::aligned_storage that isn't busted
-// on older versions of MSVC.
-template <size_t SIZE, size_t ALIGNMENT>
-struct aligned_storage {
-  struct alignas(ALIGNMENT) type {
-    unsigned char data[SIZE];
-  };
-};
+///////////////////////////////////////////////////////////////////////////////
+// TrackedAllocator
+///////////////////////////////////////////////////////////////////////////////
 
 // TrackedAllocator wraps an Allocator to track the allocations made.
 class TrackedAllocator : public Allocator {
@@ -280,6 +296,141 @@ void TrackedAllocator::free(const Allocation& allocation) {
   return allocator->free(allocation);
 }
 
+///////////////////////////////////////////////////////////////////////////////
+// StlAllocator
+///////////////////////////////////////////////////////////////////////////////
+
+// StlAllocator exposes an STL-compatible allocator wrapping a marl::Allocator.
+template <typename T>
+struct StlAllocator {
+  using value_type = T;
+  using pointer = T*;
+  using const_pointer = const T*;
+  using reference = T&;
+  using const_reference = const T&;
+  using size_type = size_t;
+  using difference_type = size_t;
+
+  // An equivalent STL allocator for a different type.
+  template <class U>
+  struct rebind {
+    typedef StlAllocator<U> other;
+  };
+
+  // Constructs an StlAllocator that will allocate using allocator.
+  // allocator must remain valid until this StlAllocator has been destroyed.
+  inline StlAllocator(Allocator* allocator);
+
+  template <typename U>
+  inline StlAllocator(const StlAllocator<U>& other);
+
+  // Returns the actual address of x even in presence of overloaded operator&.
+  inline pointer address(reference x) const;
+  inline const_pointer address(const_reference x) const;
+
+  // Allocates the memory for n objects of type T.
+  // Does not actually construct the objects.
+  inline T* allocate(std::size_t n);
+
+  // Deallocates the memory for n objects of type T.
+  inline void deallocate(T* p, std::size_t n);
+
+  // Returns the maximum theoretically possible number of T stored in this
+  // allocator.
+  inline size_type max_size() const;
+
+  // Copy constructs an object of type T at the address p.
+  inline void construct(pointer p, const_reference val);
+
+  // Constructs an object of type U at the address P forwarning all other
+  // arguments to the constructor.
+  template <typename U, typename... Args>
+  inline void construct(U* p, Args&&... args);
+
+  // Deconstructs the object at p. It does not free the memory.
+  inline void destroy(pointer p);
+
+  // Deconstructs the object at p. It does not free the memory.
+  template <typename U>
+  inline void destroy(U* p);
+
+ private:
+  inline Allocation::Request request(size_t n) const;
+
+  template <typename U>
+  friend struct StlAllocator;
+  Allocator* allocator;
+};
+
+template <typename T>
+StlAllocator<T>::StlAllocator(Allocator* allocator) : allocator(allocator) {}
+
+template <typename T>
+template <typename U>
+StlAllocator<T>::StlAllocator(const StlAllocator<U>& other) {
+  allocator = other.allocator;
+}
+
+template <typename T>
+typename StlAllocator<T>::pointer StlAllocator<T>::address(reference x) const {
+  return &x;
+}
+template <typename T>
+typename StlAllocator<T>::const_pointer StlAllocator<T>::address(
+    const_reference x) const {
+  return &x;
+}
+
+template <typename T>
+T* StlAllocator<T>::allocate(std::size_t n) {
+  auto alloc = allocator->allocate(request(n));
+  return reinterpret_cast<T*>(alloc.ptr);
+}
+
+template <typename T>
+void StlAllocator<T>::deallocate(T* p, std::size_t n) {
+  Allocation alloc;
+  alloc.ptr = p;
+  alloc.request = request(n);
+  allocator->free(alloc);
+}
+
+template <typename T>
+typename StlAllocator<T>::size_type StlAllocator<T>::max_size() const {
+  return std::numeric_limits<size_type>::max() / sizeof(value_type);
+}
+
+template <typename T>
+void StlAllocator<T>::construct(pointer p, const_reference val) {
+  new (p) T(val);
+}
+
+template <typename T>
+template <typename U, typename... Args>
+void StlAllocator<T>::construct(U* p, Args&&... args) {
+  ::new ((void*)p) U(std::forward<Args>(args)...);
+}
+
+template <typename T>
+void StlAllocator<T>::destroy(pointer p) {
+  ((T*)p)->~T();
+}
+
+template <typename T>
+template <typename U>
+void StlAllocator<T>::destroy(U* p) {
+  p->~U();
+}
+
+template <typename T>
+Allocation::Request StlAllocator<T>::request(size_t n) const {
+  Allocation::Request req = {};
+  req.size = sizeof(T) * n;
+  req.alignment = alignof(T);
+  req.usage = Allocation::Usage::Stl;
+  return req;
+}
+
 }  // namespace marl
 
 #endif  // marl_memory_h

include/marl/pool.h

@@ -363,7 +363,7 @@ class UnboundedPool : public Pool<T> {
 
     Allocator* allocator;
     marl::mutex mutex;
-    std::vector<Item*> items;
+    containers::vector<Item*, 4> items;
     Item* free = nullptr;
   };
 
@@ -373,7 +373,7 @@ class UnboundedPool : public Pool<T> {
 
 template <typename T, PoolPolicy POLICY>
 UnboundedPool<T, POLICY>::Storage::Storage(Allocator* allocator)
-    : allocator(allocator) {}
+    : allocator(allocator), items(allocator) {}
 
 template <typename T, PoolPolicy POLICY>
 UnboundedPool<T, POLICY>::Storage::~Storage() {