2 files changed, 49 insertions, 16 deletions

diff --git a/src/nvim/memory.c b/src/nvim/memory.c
index 1884d55999..3c0d001848 100644
--- a/src/nvim/memory.c
+++ b/src/nvim/memory.c
@@ -17,10 +17,31 @@
 // Force je_ prefix on jemalloc functions.
 # define JEMALLOC_NO_DEMANGLE
 # include <jemalloc/jemalloc.h>
-# define malloc(size) je_malloc(size)
-# define calloc(count, size) je_calloc(count, size)
-# define realloc(ptr, size) je_realloc(ptr, size)
-# define free(ptr) je_free(ptr)
+#endif
+
+#ifdef UNIT_TESTING
+# define malloc(size) mem_malloc(size)
+# define calloc(count, size) mem_calloc(count, size)
+# define realloc(ptr, size) mem_realloc(ptr, size)
+# define free(ptr) mem_free(ptr)
+# ifdef HAVE_JEMALLOC
+MemMalloc mem_malloc = &je_malloc;
+MemFree mem_free = &je_free;
+MemCalloc mem_calloc = &je_calloc;
+MemRealloc mem_realloc = &je_realloc;
+# else
+MemMalloc mem_malloc = &malloc;
+MemFree mem_free = &free;
+MemCalloc mem_calloc = &calloc;
+MemRealloc mem_realloc = &realloc;
+# endif
+#else
+# ifdef HAVE_JEMALLOC
+#  define malloc(size) je_malloc(size)
+#  define calloc(count, size) je_calloc(count, size)
+#  define realloc(ptr, size) je_realloc(ptr, size)
+#  define free(ptr) je_free(ptr)
+# endif
 #endif
 
 #ifdef INCLUDE_GENERATED_DECLARATIONS
@@ -353,15 +374,15 @@ char *xstpncpy(char *restrict dst, const char *restrict src, size_t maxlen)
 size_t xstrlcpy(char *restrict dst, const char *restrict src, size_t size)
   FUNC_ATTR_NONNULL_ALL
 {
-    size_t ret = strlen(src);
+  size_t ret = strlen(src);
 
-    if (size) {
-        size_t len = (ret >= size) ? size - 1 : ret;
-        memcpy(dst, src, len);
-        dst[len] = '\0';
-    }
+  if (size) {
+    size_t len = (ret >= size) ? size - 1 : ret;
+    memcpy(dst, src, len);
+    dst[len] = '\0';
+  }
 
-    return ret;
+  return ret;
 }
 
 /// strdup() wrapper
@@ -371,6 +392,7 @@ size_t xstrlcpy(char *restrict dst, const char *restrict src, size_t size)
 /// @return pointer to a copy of the string
 char *xstrdup(const char *str)
   FUNC_ATTR_MALLOC FUNC_ATTR_WARN_UNUSED_RESULT FUNC_ATTR_NONNULL_RET
+  FUNC_ATTR_NONNULL_ALL
 {
   return xmemdupz(str, strlen(str));
 }
@@ -401,6 +423,7 @@ void *xmemrchr(const void *src, uint8_t c, size_t len)
 /// @return pointer to a copy of the string
 char *xstrndup(const char *str, size_t len)
   FUNC_ATTR_MALLOC FUNC_ATTR_WARN_UNUSED_RESULT FUNC_ATTR_NONNULL_RET
+  FUNC_ATTR_NONNULL_ALL
 {
   char *p = memchr(str, '\0', len);
   return xmemdupz(str, p ? (size_t)(p - str) : len);
@@ -488,13 +511,13 @@ void time_to_bytes(time_t time_, uint8_t buf[8])
 void free_all_mem(void)
 {
   buf_T       *buf, *nextbuf;
+  static bool entered = false;
 
-  // When we cause a crash here it is caught and Vim tries to exit cleanly.
-  // Don't try freeing everything again.
-  if (entered_free_all_mem) {
+  /* When we cause a crash here it is caught and Vim tries to exit cleanly.
+   * Don't try freeing everything again. */
+  if (entered)
     return;
-  }
-  entered_free_all_mem = true;
+  entered = true;
 
   // Don't want to trigger autocommands from here on.
   block_autocmds();
diff --git a/src/nvim/memory.h b/src/nvim/memory.h
index 62cc78360c..0b422957d5 100644
--- a/src/nvim/memory.h
+++ b/src/nvim/memory.h
@@ -5,6 +5,16 @@
 #include <stddef.h>  // for size_t
 #include <time.h>    // for time_t
 
+typedef void *(*MemMalloc)(size_t);
+typedef void (*MemFree)(void *);
+typedef void *(*MemCalloc)(size_t, size_t);
+typedef void *(*MemRealloc)(void *, size_t);
+
+extern MemMalloc mem_malloc;
+extern MemFree mem_free;
+extern MemCalloc mem_calloc;
+extern MemRealloc mem_realloc;
+
 #ifdef INCLUDE_GENERATED_DECLARATIONS
 # include "memory.h.generated.h"
 #endif