Add mi_heap_stl_allocator

Rationale: heap-specific C++ allocator that can skip memory deallocation.

include/mimalloc.h

@@ -389,6 +389,9 @@ mi_decl_nodiscard mi_decl_export mi_decl_restrict void* mi_new_n(size_t count, s
 mi_decl_nodiscard mi_decl_export void* mi_new_realloc(void* p, size_t newsize)                mi_attr_alloc_size(2);
 mi_decl_nodiscard mi_decl_export void* mi_new_reallocn(void* p, size_t newcount, size_t size) mi_attr_alloc_size2(2, 3);
 
+mi_decl_nodiscard mi_decl_export mi_decl_restrict void* mi_heap_new_(size_t size, mi_heap_t *heap)                   mi_attr_malloc mi_attr_alloc_size(1);
+mi_decl_nodiscard mi_decl_export mi_decl_restrict void* mi_heap_new_n(size_t count, size_t size, mi_heap_t *heap)   mi_attr_malloc mi_attr_alloc_size2(1, 2);
+
 #ifdef __cplusplus
 }
 #endif
@@ -448,6 +451,77 @@ template<class T> struct mi_stl_allocator {
 
 template<class T1,class T2> bool operator==(const mi_stl_allocator<T1>& , const mi_stl_allocator<T2>& ) mi_attr_noexcept { return true; }
 template<class T1,class T2> bool operator!=(const mi_stl_allocator<T1>& , const mi_stl_allocator<T2>& ) mi_attr_noexcept { return false; }
+
+#if (__cplusplus >= 201103L) || (_MSC_VER > 1900)  // C++11
+#include <memory>
+
+template<class T> struct mi_heap_stl_allocator {
+  typedef T                 value_type;
+  typedef std::size_t       size_type;
+  typedef std::ptrdiff_t    difference_type;
+  typedef value_type&       reference;
+  typedef value_type const& const_reference;
+  typedef value_type*       pointer;
+  typedef value_type const* const_pointer;
+  template <class U> struct rebind { typedef mi_heap_stl_allocator<U> other; };
+
+  mi_heap_stl_allocator() {
+    mi_heap_t *heap = mi_heap_new();
+    this->_heap.reset(new(static_cast<managed_heap *>(mi_heap_new_(sizeof(managed_heap), heap))) managed_heap(heap), managed_heap::destroy);
+  }
+  mi_heap_stl_allocator(const mi_heap_stl_allocator&) mi_attr_noexcept = default;
+  template<class U> mi_heap_stl_allocator(const mi_heap_stl_allocator<U>& other) mi_attr_noexcept : _heap(std::reinterpret_pointer_cast<mi_heap_stl_allocator<T>::managed_heap>(other._heap)) { }
+  mi_heap_stl_allocator  select_on_container_copy_construction() const { return *this; }
+  void              deallocate(T* p, size_type) { if (_heap->free_enabled) mi_free(p); }
+
+  #if (__cplusplus >= 201703L)  // C++17
+  mi_decl_nodiscard T* allocate(size_type count) { return static_cast<T*>(mi_heap_new_n(count, sizeof(T), _heap->heap)); }
+  mi_decl_nodiscard T* allocate(size_type count, const void*) { return allocate(count); }
+  #else
+  mi_decl_nodiscard pointer allocate(size_type count, const void* = 0) { return static_cast<pointer>(mi_heap_new_n(count, sizeof(value_type), _heap->heap)); }
+  #endif
+
+  using propagate_on_container_copy_assignment = std::true_type;
+  using propagate_on_container_move_assignment = std::true_type;
+  using propagate_on_container_swap            = std::true_type;
+  using is_always_equal                        = std::true_type;
+  template <class U, class ...Args> void construct(U* p, Args&& ...args) { ::new(p) U(std::forward<Args>(args)...); }
+  template <class U> void destroy(U* p) mi_attr_noexcept { p->~U(); }
+
+  size_type     max_size() const mi_attr_noexcept { return (PTRDIFF_MAX/sizeof(value_type)); }
+  pointer       address(reference x) const        { return &x; }
+  const_pointer address(const_reference x) const  { return &x; }
+  
+  void enable_free() mi_attr_noexcept { this->_heap->free_enabled = true; }
+  void disable_free() mi_attr_noexcept { this->_heap->free_enabled = false; }
+  void collect(bool force = false) mi_attr_noexcept { mi_heap_collect(_heap->heap, force); }
+
+  protected:
+    struct managed_heap {
+      managed_heap(mi_heap_t *heap): heap(heap), free_enabled(true) { }
+      managed_heap(const managed_heap&) = delete;
+      managed_heap& operator=(managed_heap const&) = delete;
+      ~managed_heap() = delete;
+      static void destroy(managed_heap *ptr) { mi_heap_destroy(ptr->heap); }
+
+      mi_heap_t *heap;
+      bool free_enabled;
+    };
+    
+    std::shared_ptr<managed_heap> _heap;
+
+    template <typename>
+    friend struct mi_heap_stl_allocator;
+    template<class T1,class T2>
+    friend bool operator==(const mi_heap_stl_allocator<T1>& first, const mi_heap_stl_allocator<T2>& second) mi_attr_noexcept;
+    template<class T1,class T2>
+    friend bool operator!=(const mi_heap_stl_allocator<T1>& first, const mi_heap_stl_allocator<T2>& second) mi_attr_noexcept;
+};
+
+template<class T1,class T2> bool operator==(const mi_heap_stl_allocator<T1>& first, const mi_heap_stl_allocator<T2>& second) mi_attr_noexcept { return first._heap == second._heap; }
+template<class T1,class T2> bool operator!=(const mi_heap_stl_allocator<T1>& first, const mi_heap_stl_allocator<T2>& second) mi_attr_noexcept { return first._heap != second._heap; }
+#endif  // C++11
+
 #endif // __cplusplus
 
 #endif

src/alloc.c

@@ -932,3 +932,20 @@ void* mi_new_reallocn(void* p, size_t newcount, size_t size) {
     return mi_new_realloc(p, total);
   }
 }
+
+mi_decl_restrict void* mi_heap_new_(size_t size, mi_heap_t *heap) {
+  void* p = mi_heap_malloc(heap, size);
+  if (mi_unlikely(p == NULL)) return mi_try_new(size,false);
+  return p;
+}
+
+mi_decl_restrict void* mi_heap_new_n(size_t count, size_t size, mi_heap_t *heap) {
+  size_t total;
+  if (mi_unlikely(mi_count_size_overflow(count, size, &total))) {
+    mi_try_new_handler(false);  // on overflow we invoke the try_new_handler once to potentially throw std::bad_alloc
+    return NULL;
+  }
+  else {
+    return mi_heap_new_(total, heap);
+  }
+}

test/main-override.cpp

@@ -127,6 +127,32 @@ static bool test_stl_allocator2() {
   return vec.size() == 0;
 }
 
+static bool test_heap_stl_allocator1() {
+#if (__cplusplus >= 201103L) || (_MSC_VER > 1900)
+  mi_heap_stl_allocator<int> alloc;
+  std::vector<int, mi_heap_stl_allocator<int> > vec(alloc);
+  vec.push_back(1);
+  vec.pop_back();
+  return vec.size() == 0;
+#else
+  return true;
+#endif
+}
+
+static bool test_heap_stl_allocator2() {
+#if (__cplusplus >= 201103L) || (_MSC_VER > 1900)
+  mi_heap_stl_allocator<some_struct> alloc;
+  std::vector<some_struct, mi_heap_stl_allocator<some_struct> > vec(alloc);
+  alloc.disable_free();
+  vec.push_back(some_struct());
+  vec.pop_back();
+  alloc.enable_free();
+  return vec.size() == 0;
+#else
+  return true;
+#endif
+}
+
 // issue 445
 static void strdup_test() {
 #ifdef _MSC_VER

test/test-api.c

@@ -45,6 +45,10 @@ bool test_heap1(void);
 bool test_heap2(void);
 bool test_stl_allocator1(void);
 bool test_stl_allocator2(void);
+bool test_heap_stl_allocator1(void);
+bool test_heap_stl_allocator2(void);
+bool test_heap_stl_allocator3(void);
+bool test_heap_stl_allocator4(void);
 
 // ---------------------------------------------------------------------------
 // Main testing
@@ -193,6 +197,11 @@ int main(void) {
 
   CHECK("stl_allocator1", test_stl_allocator1());
   CHECK("stl_allocator2", test_stl_allocator2());
+  
+  CHECK("heap_stl_allocator1", test_heap_stl_allocator1());
+  CHECK("heap_stl_allocator2", test_heap_stl_allocator2());
+  CHECK("heap_stl_allocator3", test_heap_stl_allocator3());
+  CHECK("heap_stl_allocator3", test_heap_stl_allocator4());
 
   // ---------------------------------------------------
   // Done
@@ -247,3 +256,60 @@ bool test_stl_allocator2() {
   return true;
 #endif
 }
+
+bool test_heap_stl_allocator1() {
+#if (__cplusplus >= 201103L) || (_MSC_VER > 1900)
+  mi_heap_stl_allocator<int> alloc;
+  std::vector<int, mi_heap_stl_allocator<int> > vec(alloc);
+  vec.push_back(1);
+  vec.pop_back();
+  return vec.size() == 0;
+#else
+  return true;
+#endif
+}
+
+bool test_heap_stl_allocator2() {
+#if (__cplusplus >= 201103L) || (_MSC_VER > 1900)
+  mi_heap_stl_allocator<some_struct> alloc;
+  std::vector<some_struct, mi_heap_stl_allocator<some_struct> > vec(alloc);
+  vec.push_back(some_struct());
+  vec.pop_back();
+  return vec.size() == 0;
+#else
+  return true;
+#endif
+}
+
+bool test_heap_stl_allocator3() {
+#if (__cplusplus >= 201103L) || (_MSC_VER > 1900)
+  mi_heap_stl_allocator<int> alloc;
+  alloc.disable_free();
+  std::vector<int, mi_heap_stl_allocator<int> > vec(alloc);
+  for (int i = 0; i < 1000; i++) {
+    vec.push_back(i);
+  }
+  return vec.size() == 1000;
+#else
+  return true;
+#endif
+}
+
+bool test_heap_stl_allocator4() {
+#if (__cplusplus >= 201103L) || (_MSC_VER > 1900)
+  mi_heap_stl_allocator<int> alloc;
+  alloc.disable_free();
+  std::vector<int, mi_heap_stl_allocator<int> > vec(alloc);
+  for (int i = 0; i < 100; i++) {
+    vec.push_back(i);
+  }
+  alloc.enable_free();
+  for (int i = 0; i < 1000; i++) {
+    vec.push_back(i);
+  }
+  alloc.collect();
+  return vec.size() == 1100;
+#else
+  return true;
+#endif
+}