From 3b567b1054efa256d32e90ecf5ab2807ab43e6d0 Mon Sep 17 00:00:00 2001
From: Daan <daanl@outlook.com>
Date: Wed, 21 May 2025 09:20:29 -0700
Subject: [PATCH] ensure page stats are done on the subproc stats in case the
 tld is NULL

---
 include/mimalloc/internal.h |  1 +
 src/arena.c                 | 13 +++++++++----
 src/heap.c                  |  2 +-
 src/stats.c                 |  8 +++++++-
 4 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/include/mimalloc/internal.h b/include/mimalloc/internal.h
index d5c6ca2a..4dade9dd 100644
--- a/include/mimalloc/internal.h
+++ b/include/mimalloc/internal.h
@@ -234,6 +234,7 @@ void          _mi_heap_page_reclaim(mi_heap_t* heap, mi_page_t* page);
 
 // "stats.c"
 void          _mi_stats_done(mi_stats_t* stats);
+void          _mi_stats_merge_thread(mi_tld_t* tld);
 void          _mi_stats_merge_from(mi_stats_t* to, mi_stats_t* from);
 mi_msecs_t    _mi_clock_now(void);
 mi_msecs_t    _mi_clock_end(mi_msecs_t start);
diff --git a/src/arena.c b/src/arena.c
index ca9f988e..06fb9cb8 100644
--- a/src/arena.c
+++ b/src/arena.c
@@ -799,7 +799,7 @@ mi_page_t* _mi_arenas_page_alloc(mi_heap_t* heap, size_t block_size, size_t bloc
   return page;
 }
 
-void _mi_arenas_page_free(mi_page_t* page, mi_tld_t* tld /* can be NULL */) {
+void _mi_arenas_page_free(mi_page_t* page, mi_tld_t* stats_tld /* can be NULL */) {
   mi_assert_internal(_mi_is_aligned(page, MI_PAGE_ALIGN));
   mi_assert_internal(_mi_ptr_page(page)==page);
   mi_assert_internal(mi_page_is_owned(page));
@@ -807,9 +807,14 @@ void _mi_arenas_page_free(mi_page_t* page, mi_tld_t* tld /* can be NULL */) {
   mi_assert_internal(mi_page_is_abandoned(page));
   mi_assert_internal(page->next==NULL && page->prev==NULL);
 
-  if (tld==NULL) { tld = _mi_thread_tld(); }
-  mi_tld_stat_decrease(tld, page_bins[_mi_page_bin(page)], 1);
-  mi_tld_stat_decrease(tld, pages, 1);
+  if (stats_tld != NULL) { 
+    mi_tld_stat_decrease(stats_tld, page_bins[_mi_page_bin(page)], 1);
+    mi_tld_stat_decrease(stats_tld, pages, 1);
+  }
+  else {
+    mi_os_stat_decrease(page_bins[_mi_page_bin(page)], 1);
+    mi_os_stat_decrease(pages, 1);
+  }
 
   #if MI_DEBUG>1
   if (page->memid.memkind==MI_MEM_ARENA && !mi_page_is_full(page)) {
diff --git a/src/heap.c b/src/heap.c
index 861af1b9..85285bdb 100644
--- a/src/heap.c
+++ b/src/heap.c
@@ -130,7 +130,7 @@ static void mi_heap_collect_ex(mi_heap_t* heap, mi_collect_t collect)
   _mi_arenas_collect(collect == MI_FORCE /* force purge? */, collect >= MI_FORCE /* visit all? */, heap->tld);
 
   // merge statistics
-  _mi_stats_merge_from(&_mi_subproc()->stats, &heap->tld->stats);
+  _mi_stats_merge_thread(heap->tld);
 }
 
 void _mi_heap_collect_abandon(mi_heap_t* heap) {
diff --git a/src/stats.c b/src/stats.c
index 42107187..70020e40 100644
--- a/src/stats.c
+++ b/src/stats.c
@@ -391,6 +391,7 @@ void mi_stats_reset(void) mi_attr_noexcept {
 }
 
 void _mi_stats_merge_from(mi_stats_t* to, mi_stats_t* from) {
+  mi_assert_internal(to != NULL && from != NULL);
   if (to != from) {
     mi_stats_add(to, from);
     _mi_memzero(from, sizeof(mi_stats_t));
@@ -401,8 +402,13 @@ void _mi_stats_done(mi_stats_t* stats) {  // called from `mi_thread_done`
   _mi_stats_merge_from(&_mi_subproc()->stats, stats);
 }
 
+void _mi_stats_merge_thread(mi_tld_t* tld) {
+  mi_assert_internal(tld != NULL && tld->subproc != NULL);
+  _mi_stats_merge_from( &tld->subproc->stats, &tld->stats );
+}
+
 void mi_stats_merge(void) mi_attr_noexcept {
-  _mi_stats_done( mi_get_tld_stats() );
+  _mi_stats_merge_thread( _mi_thread_tld() );
 }
 
 void mi_stats_print_out(mi_output_fun* out, void* arg) mi_attr_noexcept {