From 25cc427ec5f41d92b5e77c3c02387326a11f0bd2 Mon Sep 17 00:00:00 2001
From: Andrew Noyes <andrew@weaselab.dev>
Date: Wed, 27 Mar 2024 15:34:24 -0700
Subject: [PATCH] Assert safe_free size is correct in debug builds

Closes #16
---
 Internal.h | 34 +++++++++++++++++++++++-----------
 1 file changed, 23 insertions(+), 11 deletions(-)

diff --git a/Internal.h b/Internal.h
index 45e455c..45ff778 100644
--- a/Internal.h
+++ b/Internal.h
@@ -60,6 +60,10 @@ inline int64_t peakMallocBytes = 0;
 
 inline thread_local int64_t mallocBytesDelta = 0;
 
+#ifndef NDEBUG
+constexpr auto kMallocHeaderSize = 16;
+#endif
+
 // malloc that aborts on OOM and thus always returns a non-null pointer. Must be
 // paired with `safe_free`.
 __attribute__((always_inline)) inline void *safe_malloc(size_t s) {
@@ -69,18 +73,20 @@ __attribute__((always_inline)) inline void *safe_malloc(size_t s) {
   if (mallocBytes > peakMallocBytes) {
     peakMallocBytes = mallocBytes;
   }
-  void *p = malloc(s);
-  if (p == nullptr) {
-    abort();
-  }
-  return p;
-#else
-  void *p = malloc(s);
-  if (p == nullptr) {
-    abort();
-  }
-  return p;
 #endif
+  void *p = malloc(s
+#ifndef NDEBUG
+                   + kMallocHeaderSize
+#endif
+  );
+  if (p == nullptr) {
+    abort();
+  }
+#ifndef NDEBUG
+  memcpy(p, &s, sizeof(s));
+  (char *&)p += kMallocHeaderSize;
+#endif
+  return p;
 }
 
 // Must be paired with `safe_malloc`.
@@ -93,6 +99,12 @@ __attribute__((always_inline)) inline void safe_free(void *p, size_t s) {
   mallocBytes -= s;
   free(p);
 #else
+#ifndef NDEBUG
+  (char *&)p -= kMallocHeaderSize;
+  size_t expected;
+  memcpy(&expected, p, sizeof(expected));
+  assert(s == expected);
+#endif
   free(p);
 #endif
 }