From 7ea9cf8d1fd7d9892b91df893814e113e93c8036 Mon Sep 17 00:00:00 2001
From: daan <daanl@outlook.com>
Date: Mon, 15 Jul 2019 14:53:03 -0700
Subject: [PATCH] add region collection

---
 include/mimalloc-internal.h |  2 ++
 src/heap.c                  |  5 +++++
 src/memory.c                | 25 +++++++++++++++++++++++++
 test/test-stress.c          |  1 +
 4 files changed, 33 insertions(+)

diff --git a/include/mimalloc-internal.h b/include/mimalloc-internal.h
index 13d777ef..3b45ada4 100644
--- a/include/mimalloc-internal.h
+++ b/include/mimalloc-internal.h
@@ -54,6 +54,8 @@ bool       _mi_mem_commit(void* p, size_t size, mi_stats_t* stats);
 bool       _mi_mem_protect(void* addr, size_t size);
 bool       _mi_mem_unprotect(void* addr, size_t size);
 
+void        _mi_mem_collect(mi_stats_t* stats);
+
 // "segment.c"
 mi_page_t* _mi_segment_page_alloc(size_t block_wsize, mi_segments_tld_t* tld, mi_os_tld_t* os_tld);
 void       _mi_segment_page_free(mi_page_t* page, bool force, mi_segments_tld_t* tld);
diff --git a/src/heap.c b/src/heap.c
index dc21bd0a..2b7b7a99 100644
--- a/src/heap.c
+++ b/src/heap.c
@@ -147,6 +147,11 @@ static void mi_heap_collect_ex(mi_heap_t* heap, mi_collect_t collect)
   if (collect >= FORCE) {
     _mi_segment_thread_collect(&heap->tld->segments);
   }
+
+  // collect regions
+  if (collect >= FORCE && _mi_is_main_thread()) {
+    _mi_mem_collect(&heap->tld->stats);
+  }
 }
 
 void _mi_heap_collect_abandon(mi_heap_t* heap) {
diff --git a/src/memory.c b/src/memory.c
index a2cfe229..6a72e2e0 100644
--- a/src/memory.c
+++ b/src/memory.c
@@ -336,6 +336,31 @@ void _mi_mem_free(void* p, size_t size, size_t id, mi_stats_t* stats) {
   }
 }
 
+
+/* ----------------------------------------------------------------------------
+  collection
+-----------------------------------------------------------------------------*/
+void _mi_mem_collect(mi_stats_t* stats) {
+  // free every region that has no segments in use.
+  for (size_t i = 0; i < regions_count; i++) {
+    mem_region_t* region = &regions[i];
+    if (mi_atomic_read(&region->map) == 0 && region->start != NULL) {  
+      // if no segments used, try to claim the whole region
+      uintptr_t m;
+      do {
+        m = mi_atomic_read(&region->map);
+      } while(m == 0 && !mi_atomic_compare_exchange(&region->map, ~((uintptr_t)0), 0 ));
+      if (m == 0) {
+        // on success, free the whole region
+        if (region->start != NULL) _mi_os_free((void*)region->start, MI_REGION_SIZE, stats);
+        // and release 
+        region->start = 0;
+        mi_atomic_write(&region->map,0);
+      }
+    }
+  }
+}
+
 /* ----------------------------------------------------------------------------
   Other
 -----------------------------------------------------------------------------*/
diff --git a/test/test-stress.c b/test/test-stress.c
index 55d8dd9c..4e4d9c0d 100644
--- a/test/test-stress.c
+++ b/test/test-stress.c
@@ -36,6 +36,7 @@ static void* alloc_items(size_t items) {
   if ((rand()%100) == 0) items *= 100; // 1% huge objects;
   if (items==40) items++;              // pthreads uses that size for stack increases
   uintptr_t* p = mi_mallocn_tp(uintptr_t,items);
+  if(p == NULL) return NULL;
   for (uintptr_t i = 0; i < items; i++) p[i] = (items - i) ^ cookie;
   return p;
 }