From f37aff6ee273cb149f3d103598c1c38ab673268d Mon Sep 17 00:00:00 2001
From: Daan <daanl@outlook.com>
Date: Mon, 9 Dec 2024 22:27:40 -0800
Subject: [PATCH] fix for macOS 14 and earlier

---
 include/mimalloc/internal.h |  2 +-
 src/heap.c                  |  6 +++---
 src/init.c                  | 17 +++++++++++++----
 3 files changed, 17 insertions(+), 8 deletions(-)

diff --git a/include/mimalloc/internal.h b/include/mimalloc/internal.h
index c189a082..8a61a58e 100644
--- a/include/mimalloc/internal.h
+++ b/include/mimalloc/internal.h
@@ -186,7 +186,7 @@ size_t      _mi_bin_size(uint8_t bin);           // for stats
 uint8_t     _mi_bin(size_t size);                // for stats
 
 // "heap.c"
-void        _mi_heap_init(mi_heap_t* heap, mi_arena_id_t arena_id, bool noreclaim, uint8_t tag);
+void        _mi_heap_init(mi_heap_t* heap, mi_arena_id_t arena_id, bool noreclaim, uint8_t tag, mi_tld_t* tld);
 void        _mi_heap_destroy_pages(mi_heap_t* heap);
 void        _mi_heap_collect_abandon(mi_heap_t* heap);
 void        _mi_heap_set_default_direct(mi_heap_t* heap);
diff --git a/src/heap.c b/src/heap.c
index d2914361..ee0a8ce9 100644
--- a/src/heap.c
+++ b/src/heap.c
@@ -182,9 +182,9 @@ mi_heap_t* mi_heap_get_backing(void) {
   return bheap;
 }
 
-void _mi_heap_init(mi_heap_t* heap, mi_arena_id_t arena_id, bool noreclaim, uint8_t tag) {
+void _mi_heap_init(mi_heap_t* heap, mi_arena_id_t arena_id, bool noreclaim, uint8_t tag, mi_tld_t* tld) {
   _mi_memcpy_aligned(heap, &_mi_heap_empty, sizeof(mi_heap_t));
-  heap->tld = _mi_tld();
+  heap->tld = (tld == NULL ? _mi_tld() : tld);  // avoid reading the thread-local tld during initialization
   heap->thread_id  = _mi_thread_id();
   heap->arena_id   = arena_id;
   heap->allow_page_reclaim = !noreclaim;
@@ -216,7 +216,7 @@ mi_decl_nodiscard mi_heap_t* mi_heap_new_ex(int heap_tag, bool allow_destroy, mi
   mi_heap_t* heap = mi_heap_malloc_tp(bheap, mi_heap_t);  // todo: OS allocate in secure mode?
   if (heap == NULL) return NULL;
   mi_assert(heap_tag >= 0 && heap_tag < 256);
-  _mi_heap_init(heap, arena_id, allow_destroy /* no reclaim? */, (uint8_t)heap_tag /* heap tag */);
+  _mi_heap_init(heap, arena_id, allow_destroy /* no reclaim? */, (uint8_t)heap_tag /* heap tag */, NULL);
   return heap;
 }
 
diff --git a/src/init.c b/src/init.c
index 2070405d..19e111d3 100644
--- a/src/init.c
+++ b/src/init.c
@@ -214,7 +214,7 @@ static void mi_heap_main_init(void) {
   if (_mi_heap_main.cookie == 0) {
     _mi_heap_main.thread_id = _mi_thread_id();
     _mi_heap_main.cookie = 1;
-    #if defined(_WIN32) && !defined(MI_SHARED_LIB)
+    #if defined(__APPLE__) || defined(_WIN32) && !defined(MI_SHARED_LIB)
       _mi_random_init_weak(&_mi_heap_main.random);    // prevent allocation failure during bcrypt dll initialization with static linking
     #else
       _mi_random_init(&_mi_heap_main.random);
@@ -344,8 +344,12 @@ static bool _mi_thread_heap_init(void) {
     //mi_assert_internal(_mi_heap_default->tld->heap_backing == mi_prim_get_default_heap());
   }
   else {
-    // allocate heap and thread local data
-    mi_tld_t* tld = _mi_tld();  // allocates & initializes tld if needed
+    // allocates tld data 
+    // note: we cannot access thread-locals yet as that can cause (recursive) allocation (on macOS <= 14 for 
+    // example where the loader allocates thread-local data on demand).
+    mi_tld_t* tld = mi_tld_alloc();  
+    
+    // allocate and initialize the heap
     mi_memid_t memid;
     mi_heap_t* heap = (tld == NULL ? NULL : (mi_heap_t*)_mi_meta_zalloc(sizeof(mi_heap_t), &memid));
     if (heap==NULL || tld==NULL) {
@@ -353,8 +357,13 @@ static bool _mi_thread_heap_init(void) {
       return false;
     }
     heap->memid = memid;
-    _mi_heap_init(heap, _mi_arena_id_none(), false /* can reclaim */, 0 /* default tag */);
+    _mi_heap_init(heap, _mi_arena_id_none(), false /* can reclaim */, 0 /* default tag */, tld);
+
+    // associate the heap with this thread
+    // (this is safe, on macOS for example, the heap is set in a dedicated TLS slot and thus does not cause recursive allocation)
     _mi_heap_set_default_direct(heap);
+    // now that the heap is set for this thread, we can set the thread-local tld.
+    mi_tld = tld; 
   }
   return false;
 }