4 files changed, 40 insertions, 5 deletions
diff --git a/libdimension/compiler.h b/libdimension/compiler.h
index cb77ca6..d6435d0 100644
--- a/libdimension/compiler.h
+++ b/libdimension/compiler.h
@@ -39,9 +39,13 @@
 #endif
 
 #ifdef __GNUC__
-  #define DMNSN_HOT      __attribute__((hot))
-  #define DMNSN_INTERNAL __attribute__((visibility("hidden")))
+  #define DMNSN_HOT             __attribute__((hot))
+  #define DMNSN_INTERNAL        __attribute__((visibility("hidden")))
+  #define DMNSN_DESTRUCTOR      __attribute__((destructor(102)))
+  #define DMNSN_LATE_DESTRUCTOR __attribute__((destructor(101)))
 #else
   #define DMNSN_HOT
   #define DMNSN_INTERNAL
+  #define DMNSN_DESTRUCTOR
+  #define DMNSN_LATE_DESTRUCTOR
 #endif
diff --git a/libdimension/malloc.c b/libdimension/malloc.c
index 91364d4..2f11281 100644
--- a/libdimension/malloc.c
+++ b/libdimension/malloc.c
@@ -23,10 +23,14 @@
  * Dynamic memory.
  */
 
-#include "dimension.h"
+#include "dimension-impl.h"
 #include <stdlib.h>
 #include <string.h>
 
+#ifndef NDEBUG
+static size_t dmnsn_allocs = 0;
+#endif
+
 void *
 dmnsn_malloc(size_t size)
 {
@@ -34,12 +38,23 @@ dmnsn_malloc(size_t size)
   if (!ptr) {
     dmnsn_error("Memory allocation failed.");
   }
+
+#ifndef NDEBUG
+  __sync_fetch_and_add(&dmnsn_allocs, 1);
+#endif
+
   return ptr;
 }
 
 void *
 dmnsn_realloc(void *ptr, size_t size)
 {
+#ifndef NDEBUG
+  if (!ptr) {
+    __sync_fetch_and_add(&dmnsn_allocs, 1);
+  }
+#endif
+
   ptr = realloc(ptr, size);
   if (!ptr) {
     dmnsn_error("Memory allocation failed.");
@@ -58,5 +73,21 @@ dmnsn_strdup(const char *s)
 void
 dmnsn_free(void *ptr)
 {
+#ifndef NDEBUG
+  if (ptr) {
+    __sync_fetch_and_sub(&dmnsn_allocs, 1);
+  }
+#endif
+
   free(ptr);
 }
+
+#ifndef NDEBUG
+DMNSN_LATE_DESTRUCTOR static void
+dmnsn_leak_check(void)
+{
+  if (dmnsn_allocs > 0) {
+    dmnsn_warning("Leaking memory.");
+  }
+}
+#endif
diff --git a/libdimension/profile.c b/libdimension/profile.c
index f9dd2d9..8fca36a 100644
--- a/libdimension/profile.c
+++ b/libdimension/profile.c
@@ -169,7 +169,7 @@ dmnsn_print_bad_prediction(void *ptr)
   dmnsn_free(branch->location);
 }
 
-static void __attribute__((destructor))
+DMNSN_DESTRUCTOR static void
 dmnsn_print_bad_predictions(void)
 {
   dmnsn_dictionary *thread_profile = dmnsn_get_thread_profile();
diff --git a/libdimension/prtree.c b/libdimension/prtree.c
index e73ea16..2d8acc4 100644
--- a/libdimension/prtree.c
+++ b/libdimension/prtree.c
@@ -564,7 +564,7 @@ dmnsn_get_prtree_caches(void)
 
 /** Needed because pthreads doesn't destroy data from the main thread unless
     it exits with pthread_exit(). */
-static void __attribute__((destructor))
+DMNSN_DESTRUCTOR static void
 dmnsn_delete_main_prtree_caches(void)
 {
   dmnsn_delete_array(dmnsn_get_prtree_caches());