update for new dynamic direction on windows 64-bit

1 files changed, 55 insertions, 11 deletions

diff --git a/src/init.c b/src/init.c
index f55b731..06aa28c 100644
--- a/src/init.c
+++ b/src/init.c
@@ -373,6 +373,47 @@ void mi_thread_done(void) mi_attr_noexcept {
 // --------------------------------------------------------
 static void mi_process_done(void);
 
+static bool os_preloading = true;    // true until this module is initialized
+static bool mi_redirected = false;   // true if malloc redirects to mi_malloc
+
+// Returns true if this module has not been initialized; Don't use C runtime routines until it returns false.
+bool _mi_preloading() {
+  return os_preloading;
+}
+
+// Communicate with the redirection module on Windows
+#if defined(_WIN32) && defined(MI_SHARED_LIB) 
+mi_decl_export void _mi_redirect_init() {
+  // called on redirection
+  mi_redirected = true;
+}
+__declspec(dllimport) bool mi_allocator_init(const char** message);
+__declspec(dllimport) void mi_allocator_done();
+#else
+static bool mi_allocator_init(const char** message) {
+  if (message != NULL) *message = NULL;
+  return true;
+}
+static void mi_allocator_done() {
+  // nothing to do
+}
+#endif
+
+// Called once by the process loader
+static void mi_process_load(void) {
+  os_preloading = false;
+  atexit(&mi_process_done);
+  mi_process_init();
+  //mi_stats_reset();
+  if (mi_redirected) _mi_verbose_message("malloc is redirected.\n");
+
+  // show message from the redirector (if present)
+  const char* msg = NULL;
+  mi_allocator_init(&msg);
+  if (msg != NULL) _mi_verbose_message(msg);  
+}
+
+// Initialize the process; called by thread_init or the process loader
 void mi_process_init(void) mi_attr_noexcept {
   // ensure we are called once
   if (_mi_process_is_initialized) return;
@@ -381,7 +422,7 @@ void mi_process_init(void) mi_attr_noexcept {
   // when using dynamic linking with interpose.
   mi_heap_t* h = _mi_heap_default;
   _mi_process_is_initialized = true;
-
+  
   _mi_heap_main.thread_id = _mi_thread_id();
   _mi_verbose_message("process init: 0x%zx\n", _mi_heap_main.thread_id);
   uintptr_t random = _mi_random_init(_mi_heap_main.thread_id)  ^ (uintptr_t)h;
@@ -389,15 +430,16 @@ void mi_process_init(void) mi_attr_noexcept {
   _mi_heap_main.cookie = (uintptr_t)&_mi_heap_main ^ random;
   #endif
   _mi_heap_main.random = _mi_random_shuffle(random);
+  mi_process_setup_auto_thread_done();
+  _mi_os_init();  
   #if (MI_DEBUG)
   _mi_verbose_message("debug level : %d\n", MI_DEBUG);
   #endif
-  atexit(&mi_process_done);
-  mi_process_setup_auto_thread_done();
-  mi_stats_reset();
-  _mi_os_init();
+  mi_thread_init();
+  mi_stats_reset();  // only call stat reset *after* thread init (or the heap tld == NULL)
 }
 
+// Called when the process is done (through `at_exit`)
 static void mi_process_done(void) {
   // only shutdown if we were initialized
   if (!_mi_process_is_initialized) return;
@@ -413,7 +455,9 @@ static void mi_process_done(void) {
       mi_option_is_enabled(mi_option_verbose)) {
     mi_stats_print(NULL);
   }
+  mi_allocator_done();
   _mi_verbose_message("process done: 0x%zx\n", _mi_heap_main.thread_id);
+  os_preloading = true; // don't call the C runtime anymore
 }
 
 
@@ -425,8 +469,8 @@ static void mi_process_done(void) {
   __declspec(dllexport) BOOL WINAPI DllMain(HINSTANCE inst, DWORD reason, LPVOID reserved) {
     UNUSED(reserved);
     UNUSED(inst);
-    if (reason==DLL_PROCESS_ATTACH) {
-      mi_process_init();
+    if (reason==DLL_PROCESS_ATTACH) {     
+      mi_process_load();
     }
     else if (reason==DLL_THREAD_DETACH) {
       mi_thread_done();
@@ -437,7 +481,7 @@ static void mi_process_done(void) {
 #elif defined(__cplusplus)
   // C++: use static initialization to detect process start
   static bool _mi_process_init(void) {
-    mi_process_init();
+    mi_process_load();
     return (_mi_heap_main.thread_id != 0);
   }
   static bool mi_initialized = _mi_process_init();
@@ -445,14 +489,14 @@ static void mi_process_done(void) {
 #elif defined(__GNUC__) || defined(__clang__)
   // GCC,Clang: use the constructor attribute
   static void __attribute__((constructor)) _mi_process_init(void) {
-    mi_process_init();
+    mi_process_load();
   }
 
 #elif defined(_MSC_VER)
   // MSVC: use data section magic for static libraries
   // See <https://www.codeguru.com/cpp/misc/misc/applicationcontrol/article.php/c6945/Running-Code-Before-and-After-Main.htm>
   static int _mi_process_init(void) {
-    mi_process_init();
+    mi_process_load();
     return 0;
   }
   typedef int(*_crt_cb)(void);
@@ -467,5 +511,5 @@ static void mi_process_done(void) {
   #pragma data_seg()
 
 #else
-#pragma message("define a way to call mi_process_init/done on your platform")
+#pragma message("define a way to call mi_process_load on your platform")
 #endif