Support prefixed and non-prefixed JeMalloc functions

Note that linking a JeMalloc library without prefixed functions will replace all memory allocations, including any call to "new", not just the ones specifically called via Star::malloc etc.

3 files changed, 41 insertions, 1 deletions

diff --git a/cmake/FindJeMalloc.cmake b/cmake/FindJeMalloc.cmake
index 54abd29..2c35ca9 100644
--- a/cmake/FindJeMalloc.cmake
+++ b/cmake/FindJeMalloc.cmake
@@ -37,8 +37,23 @@ find_package_handle_standard_args(JeMalloc DEFAULT_MSG
     JEMALLOC_INCLUDE_DIR
 )
 
+# Check if the JeMalloc library was compiled with the "je_" prefix.
+include(CheckSymbolExists)
+set(CMAKE_REQUIRED_INCLUDES ${JEMALLOC_INCLUDE_DIR})
+set(CMAKE_REQUIRED_LIBRARIES ${JEMALLOC_LIBRARY})
+check_symbol_exists("je_malloc" "jemalloc/jemalloc.h" _jemalloc_is_prefixed)
+unset(CMAKE_REQUIRED_INCLUDES)
+unset(CMAKE_REQUIRED_LIBRARIES)
+
+if(_jemalloc_is_prefixed)
+    set(JEMALLOC_IS_PREFIXED TRUE)
+endif()
+
+unset(_jemalloc_is_prefixed)
+
 mark_as_advanced(
     JEMALLOC_ROOT_DIR
     JEMALLOC_LIBRARY
     JEMALLOC_INCLUDE_DIR
+    JEMALLOC_IS_PREFIXED
 )
diff --git a/source/core/CMakeLists.txt b/source/core/CMakeLists.txt
index bbad6f4..ce35500 100644
--- a/source/core/CMakeLists.txt
+++ b/source/core/CMakeLists.txt
@@ -214,4 +214,10 @@ ELSEIF (STAR_SYSTEM_FAMILY_WINDOWS)
 ENDIF ()
 
 ADD_LIBRARY (star_core OBJECT ${star_core_SOURCES} ${star_core_HEADERS})
-TARGET_PRECOMPILE_HEADERS (star_core PUBLIC StarPch.hpp)
\ No newline at end of file
+TARGET_PRECOMPILE_HEADERS (star_core PUBLIC StarPch.hpp)
+
+IF(STAR_USE_JEMALLOC AND JEMALLOC_IS_PREFIXED)
+  SET_SOURCE_FILES_PROPERTIES(StarMemory.cpp PROPERTIES
+          COMPILE_DEFINITIONS STAR_JEMALLOC_IS_PREFIXED
+          )
+ENDIF()
\ No newline at end of file
diff --git a/source/core/StarMemory.cpp b/source/core/StarMemory.cpp
index 3855283..9e73eb6 100644
--- a/source/core/StarMemory.cpp
+++ b/source/core/StarMemory.cpp
@@ -7,6 +7,7 @@
 namespace Star {
 
 #ifdef STAR_USE_JEMALLOC
+#ifdef STAR_JEMALLOC_IS_PREFIXED
   void* malloc(size_t size) {
     return je_malloc(size);
   }
@@ -33,6 +34,24 @@ namespace Star {
   }
 
   void free(void* ptr) {
+    ::free(ptr);
+  }
+
+  void free(void* ptr, size_t size) {
+    if (ptr)
+      ::sdallocx(ptr, size, 0);
+  }
+#endif
+#else
+  void* malloc(size_t size) {
+    return ::malloc(size);
+  }
+
+  void* realloc(void* ptr, size_t size) {
+    return ::realloc(ptr, size);
+  }
+
+  void free(void* ptr) {
     return ::free(ptr);
   }