From 6b6a9bace941fc0ced9daf124e34c327003f3e12 Mon Sep 17 00:00:00 2001
From: Andrew Noyes <andrew@weaselab.dev>
Date: Wed, 13 Mar 2024 16:57:38 -0700
Subject: [PATCH] Fix SHOW_MEMORY for gcc and glibc on linux

---
 ConflictSet.cpp |  4 +++-
 Internal.h      | 17 +++++++++++++----
 2 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/ConflictSet.cpp b/ConflictSet.cpp
index 88e791b..8efa9eb 100644
--- a/ConflictSet.cpp
+++ b/ConflictSet.cpp
@@ -47,7 +47,9 @@ limitations under the License.
 #if __has_builtin(__builtin_assume)
 #define assume(e) __builtin_assume(e)
 #else
-#define assume(e) __attribute__((assume(e)))
+#define assume(e)                                                              \
+  if (!(e))                                                                    \
+  __builtin_unreachable()
 #endif
 #else
 #define assume assert
diff --git a/Internal.h b/Internal.h
index eb59dc1..30ae021 100644
--- a/Internal.h
+++ b/Internal.h
@@ -43,6 +43,8 @@ operator<=>(const std::span<const uint8_t> &lhs,
 #if SHOW_MEMORY
 #ifdef __APPLE__
 #include <malloc/malloc.h>
+#else
+#include <malloc.h>
 #endif
 inline int64_t mallocBytes = 0;
 inline int64_t peakMallocBytes = 0;
@@ -50,16 +52,21 @@ inline int64_t peakMallocBytes = 0;
 
 // malloc that aborts on OOM and thus always returns a non-null pointer
 __attribute__((always_inline)) inline void *safe_malloc(size_t s) {
+  void *p = malloc(s);
+  if (p == nullptr) {
+    abort();
+  }
 #if SHOW_MEMORY
+#ifdef __APPLE__
   mallocBytes += s;
+#else
+  mallocBytes += malloc_usable_size(p);
+#endif
   if (mallocBytes > peakMallocBytes) {
     peakMallocBytes = mallocBytes;
   }
 #endif
-  if (void *p = malloc(s)) {
-    return p;
-  }
-  abort();
+  return p;
 }
 
 // There's nothing safer about this than free. Only called safe_free for
@@ -68,6 +75,8 @@ __attribute__((always_inline)) inline void safe_free(void *p) {
 #if SHOW_MEMORY
 #ifdef __APPLE__
   mallocBytes -= malloc_size(p);
+#else
+  mallocBytes -= malloc_usable_size(p);
 #endif
 #endif
   free(p);