RFC: Move allocated, freed tracking into alloc, free functions

include/mimalloc/prim/prim-tls.h

@@ -9,6 +9,7 @@ terms of the MIT license. A copy of the license can be found in the file
 #define MI_INTERNAL_PRIM_TLS_H
 
 #include "../types.h"
+#include "../alloc/page.h"
 #include "../heap/heap.h"
 #include "../heap/theap.h"
 #include "../util/threadlocal.h"
@@ -442,4 +443,13 @@ static inline mi_theap_t* _mi_page_associated_theap_peek(const mi_page_t* page)
   return theap;
 }
 
+// Gets the theap belonging to a page.
+static inline mi_theap_t* mi_theap_of_page(mi_page_t* page) {
+  mi_theap_t* theap = page->theap;
+  if mi_unlikely(mi_page_thread_id(page) != _mi_prim_thread_id()) {
+    theap = _mi_page_associated_theap_peek(page);
+  }
+  return theap;
+}
+
 #endif  // MI_PRIM_TLS_H

src/alloc/alloc.c

@@ -143,6 +143,20 @@ static mi_decl_forceinline mi_decl_restrict void* mi_page_malloc_zero(mi_theap_t
 
   // track per-block statistics
   mi_assert_internal(mi_page_alloc_count(page) + mi_page_last_used(page) >= mi_page_used(page));
+
+  #if MI_STATS>=2
+  const size_t bin = _mi_bin(bsize);
+  mi_theap_stat_counter_increase(theap, malloc_count, 1);
+  mi_theap_stat_increase(theap, allocated, bsize);
+  // deprecated
+  if (bin!=MI_BIN_HUGE) {
+    mi_theap_stat_counter_increase(theap, malloc_normal_count, 1);
+    mi_theap_stat_increase(theap, malloc_normal, bsize);
+  } else {
+    mi_theap_stat_counter_increase(theap, malloc_huge_count, 1);
+    mi_theap_stat_increase(theap,malloc_huge,bsize);
+  }
+  #endif
 
   // in debug mode initialize with 0xD0
   #if (MI_DEBUG>0) && !MI_TRACK_ENABLED && !MI_TSAN

src/alloc/free.c

@@ -54,6 +54,21 @@ static inline void mi_free_block_local(mi_page_t* page, mi_block_t* block, bool
   page->local_free = block;
   const bool is_empty = (mi_xused_used_count(xused) == 0);
   mi_assert_internal(mi_page_alloc_count(page) + mi_page_last_used(page) >= mi_page_used(page));
+
+  #if MI_STATS>=2
+  mi_theap_t* theap = mi_theap_of_page(page);
+  const size_t bsize = mi_page_usable_block_size(page);
+  const size_t bin = _mi_bin(bsize);
+  mi_theap_stat_counter_increase(theap, free_count, 1);
+  mi_theap_stat_counter_increase(theap, freed, bsize);
+  mi_theap_stat_decrease(theap, allocated, bsize);
+  if (bin != MI_BIN_HUGE) {
+    mi_theap_stat_decrease(theap, malloc_normal, bsize);
+  } else {
+    mi_theap_stat_decrease(theap, malloc_huge, bsize);
+  }
+  #endif
+
   if mi_unlikely(is_empty) {      // is used count zero ?
     if (page->retire_expire==0) { // no need to re-retire retired pages (happens when we alloc/free one block repeatedly in an empty page)
       _mi_page_retire(page); 

src/alloc/page.c

@@ -138,17 +138,6 @@ bool _mi_page_is_valid(mi_page_t* page) {
 ----------------------------------------------------------- */
 
 
-#if MI_STATS || (MI_SAMPLE==1)
-// Gets the theap belonging to a page.
-static mi_theap_t* mi_theap_of_page(mi_page_t* page) {
-  mi_theap_t* theap = page->theap;
-  if mi_unlikely(mi_page_thread_id(page) != _mi_prim_thread_id()) { 
-    theap = _mi_page_associated_theap_peek(page);
-  }
-  return theap;
-}
-#endif
-
 #if MI_SAMPLE==1 /* for ==2, the countdown is already done at every `alloc/alloc.c:mi_page_alloc_zero` */
 
 // Adjust the theap->sample_countdown based on the current page allocation count
@@ -198,14 +187,13 @@ static void mi_theap_page_merge_stats(mi_theap_t* theapx, const mi_page_t* page,
   // adjust stats
   const size_t bin = _mi_bin(bsize);
   if (alloc_count > 0) {
+    mi_theapx_stat_increase(heap, theapx, malloc_bins[bin], alloc_count);
+    #if MI_STATS==1
     mi_theapx_stat_counter_increase(heap, theapx, malloc_count, alloc_count);
     mi_theapx_stat_increase(heap, theapx, allocated, allocated);
-    mi_theapx_stat_increase(heap, theapx, malloc_bins[bin], alloc_count);
     if (bin==MI_BIN_HUGE) { mi_theapx_stat_increase(heap,theapx,malloc_huge,allocated); } // separate so we have a count and total
-    #if MI_STATS==1
     // use coarse total requested bytes
     mi_theapx_stat_counter_increase(heap, theapx, malloc_requested, requested);      
-    #endif
     // deprecated
     if (bin!=MI_BIN_HUGE) { 
       mi_theapx_stat_counter_increase(heap, theapx, malloc_normal_count, alloc_count); 
@@ -214,15 +202,18 @@ static void mi_theap_page_merge_stats(mi_theap_t* theapx, const mi_page_t* page,
     else {
       mi_theapx_stat_counter_increase(heap, theapx, malloc_huge_count, alloc_count);
     }
+    #endif
   }
   if (free_count > 0) {
+    mi_theapx_stat_decrease(heap, theapx, malloc_bins[bin], free_count);
+    #if MI_STATS==1
     mi_theapx_stat_counter_increase(heap, theapx, free_count, free_count);
     mi_theapx_stat_counter_increase(heap, theapx, freed, freed);
     mi_theapx_stat_decrease(heap, theapx, allocated, freed);
-    mi_theapx_stat_decrease(heap, theapx, malloc_bins[bin], free_count);      
     if (bin==MI_BIN_HUGE) { mi_theapx_stat_decrease(heap, theapx, malloc_huge, freed); }
     // deprecated
     if (bin!=MI_BIN_HUGE) { mi_theapx_stat_decrease(heap, theapx, malloc_normal, freed); }      
+    #endif
   }
 
   // merge the theap stats into the heap once N bytes were allocated or freed since the last merge