From 08e05bb10bb7f8b7807668ae7e27e50db28feee5 Mon Sep 17 00:00:00 2001
From: Elena Lepilkina <Elena.Lepilkina@jetbrains.com>
Date: Fri, 20 Dec 2019 16:25:31 +0300
Subject: [PATCH] Fixes in allocator

---
 runtime/src/mimalloc/c/include/mimalloc-atomic.h  | 13 +++++++++++++
 .../src/mimalloc/c/include/mimalloc-internal.h    | 15 ++++++++++++++-
 runtime/src/mimalloc/c/include/mimalloc-types.h   |  4 +++-
 3 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/runtime/src/mimalloc/c/include/mimalloc-atomic.h b/runtime/src/mimalloc/c/include/mimalloc-atomic.h
index 10368df348f..3a56986d014 100644
--- a/runtime/src/mimalloc/c/include/mimalloc-atomic.h
+++ b/runtime/src/mimalloc/c/include/mimalloc-atomic.h
@@ -214,9 +214,22 @@ static inline void mi_atomic_write(volatile _Atomic(uintptr_t)* p, uintptr_t x)
     asm volatile ("pause" ::: "memory");
   }
 #elif defined(__arm__) || defined(__aarch64__)
+  #if defined(KONAN_MI_MALLOC)
+    #if defined(__arm__)
+      #include <sched.h>
+      static inline void mi_atomic_yield(void) {
+        sched_yield();
+      }
+    #else
+      static inline void mi_atomic_yield(void) {
+        asm volatile("yield");
+      }
+    #endif
+  #else
   static inline void mi_atomic_yield(void) {
     asm volatile("yield");
   }
+  #endif
 #endif
 #elif defined(__wasi__)
   #include <sched.h>
diff --git a/runtime/src/mimalloc/c/include/mimalloc-internal.h b/runtime/src/mimalloc/c/include/mimalloc-internal.h
index 452f0b68ad5..f346f048033 100644
--- a/runtime/src/mimalloc/c/include/mimalloc-internal.h
+++ b/runtime/src/mimalloc/c/include/mimalloc-internal.h
@@ -457,6 +457,10 @@ static inline uintptr_t _mi_thread_id(void) mi_attr_noexcept {
 }
 #elif (defined(__GNUC__) || defined(__clang__)) && \
       (defined(__x86_64__) || defined(__i386__) || defined(__arm__) || defined(__aarch64__))
+#if defined(KONAN_MI_MALLOC)
+  #include <pthread.h>
+  pthread_t pthread_self(void);
+#endif
 // TLS register on x86 is in the FS or GS register
 // see: https://akkadia.org/drepper/tls.pdf
 static inline uintptr_t _mi_thread_id(void) mi_attr_noexcept {
@@ -464,7 +468,16 @@ static inline uintptr_t _mi_thread_id(void) mi_attr_noexcept {
   #if defined(__i386__)
   __asm__("movl %%gs:0, %0" : "=r" (tid) : : );  // 32-bit always uses GS
   #elif defined(__MACH__)
-  __asm__("movq %%gs:0, %0" : "=r" (tid) : : );  // x86_64 macOS uses GS
+  #if defined(KONAN_MI_MALLOC)
+  #include <TargetConditionals.h>
+    #if TARGET_OS_EMBEDDED // iOS/tvOS/watchOS devices.
+      tid = pthread_self();
+    #else
+       __asm__("movq %%gs:0, %0" : "=r" (tid) : : );  // x86_64 macOS uses GS
+    #endif
+  #else
+    __asm__("movq %%gs:0, %0" : "=r" (tid) : : );  // x86_64 macOS uses GS
+  #endif
   #elif defined(__x86_64__)
   __asm__("movq %%fs:0, %0" : "=r" (tid) : : );  // x86_64 Linux, BSD uses FS
   #elif defined(__arm__)
diff --git a/runtime/src/mimalloc/c/include/mimalloc-types.h b/runtime/src/mimalloc/c/include/mimalloc-types.h
index 96e1860f878..de197103f37 100644
--- a/runtime/src/mimalloc/c/include/mimalloc-types.h
+++ b/runtime/src/mimalloc/c/include/mimalloc-types.h
@@ -17,7 +17,9 @@ terms of the MIT license. A copy of the license can be found in the file
 // ------------------------------------------------------
 
 // Define NDEBUG in the release version to disable assertions.
-// #define NDEBUG
+#if !defined(KONAN_MI_MALLOC)
+  #define NDEBUG
+#endif
 
 // Define MI_STAT as 1 to maintain statistics; set it to 2 to have detailed statistics (but costs some performance).
 // #define MI_STAT 1