rts: specify deallocator of allocating functions

This patch adds a new `STG_MALLOC1` macro (and its counterpart
`STG_MALLOC2` for completeness) which allows to specify the deallocation
function to be used with allocations of allocating functions, and
applies it to `stg*allocBytes`.

It also fixes a case where `free` was used to free up an
`stgMallocBytes` allocation, found by the above change.

See: https://gcc.gnu.org/onlinedocs/gcc/Common-Function-Attributes.html#index-malloc-function-attribute
See: https://gitlab.haskell.org/ghc/ghc/-/issues/22381

3 files changed, 50 insertions, 6 deletions

diff --git a/rts/Heap.c b/rts/Heap.c
index 491324e706..a3be4da749 100644
--- a/rts/Heap.c
+++ b/rts/Heap.c
@@ -278,7 +278,7 @@ StgMutArrPtrs *heap_view_closurePtrs(Capability *cap, StgClosure *closure) {
     for (StgWord i = 0; i<nptrs; i++) {
         arr->payload[i] = ptrs[i];
     }
-    free(ptrs);
+    stgFree(ptrs);
 
     return arr;
 }
diff --git a/rts/RtsUtils.h b/rts/RtsUtils.h
index d3618d6435..5282be2035 100644
--- a/rts/RtsUtils.h
+++ b/rts/RtsUtils.h
@@ -17,18 +17,25 @@
 void initAllocator(void);
 void shutdownAllocator(void);
 
+void stgFree(void* p);
+
 void *stgMallocBytes(size_t n, char *msg)
-    STG_MALLOC;
+    STG_MALLOC STG_MALLOC1(stgFree);
 
-void *stgReallocBytes(void *p, size_t n, char *msg);
+void *stgReallocBytes(void *p, size_t n, char *msg)
+    STG_MALLOC1(stgFree);
+/* Note: `stgRallocBytes` can *not* be tagged as `STG_MALLOC`
+ * since its return value *can* alias an existing pointer (i.e.,
+ * the given pointer `p`).
+ * See the documentation of the `malloc` attribute in the GCC manual
+ * for more information.
+ */
 
 void *stgCallocBytes(size_t count, size_t size, char *msg)
-    STG_MALLOC;
+    STG_MALLOC STG_MALLOC1(stgFree);
 
 char *stgStrndup(const char *s, size_t n);
 
-void stgFree(void* p);
-
 /* -----------------------------------------------------------------------------
  * Misc other utilities
  * -------------------------------------------------------------------------- */
diff --git a/rts/include/Stg.h b/rts/include/Stg.h
index ffecd70240..4f1f9d3849 100644
--- a/rts/include/Stg.h
+++ b/rts/include/Stg.h
@@ -246,6 +246,43 @@
 
 #define STG_MALLOC GNUC3_ATTRIBUTE(__malloc__)
 
+/* Instead of relying on GCC version checks to expand attributes,
+ * use `__has_attribute` which is supported by GCC >= 5 and Clang. Hence, the
+ * following macros won't expand on older compiler versions, but since they're
+ * purely for optimization or static analysis purposes, there's no harm done.
+ *
+ * See: https://gcc.gnu.org/onlinedocs/cpp/_005f_005fhas_005fattribute.html
+ * See: https://clang.llvm.org/docs/LanguageExtensions.html#has-attribute
+ */
+#ifdef __has_attribute
+# define stg__has_attribute(attr) __has_attribute(attr)
+#else
+# define stg__has_attribute(attr) (0)
+#endif
+
+#ifdef __GNUC__
+# define STG_GNUC_GUARD_VERSION(major, minor) \
+    ((__GNUC__ > (major)) || \
+      ((__GNUC__ == (major)) && (__GNUC_MINOR__ >= (minor))))
+#else
+# define STG_GNUC_GUARD_VERSION(major, minor) (0)
+#endif
+
+/*
+ * The versions of the `__malloc__` attribute which take arguments are only
+ * supported in GCC 11 and later.
+ *
+ * See: https://gcc.gnu.org/onlinedocs/gcc/Common-Function-Attributes.html#index-malloc-function-attribute
+ * See: https://developers.redhat.com/blog/2021/04/30/detecting-memory-management-bugs-with-gcc-11-part-1-understanding-dynamic-allocation#attribute_malloc
+ */
+#if stg__has_attribute(__malloc__) && STG_GNUC_GUARD_VERSION(11, 0)
+# define STG_MALLOC1(deallocator) __attribute__((__malloc__(deallocator)))
+# define STG_MALLOC2(deallocator, ptrIndex) __attribute__((__malloc__(deallocator, ptrIndex)))
+#else
+# define STG_MALLOC1(deallocator)
+# define STG_MALLOC2(deallocator, ptrIndex)
+#endif
+
 /* -----------------------------------------------------------------------------
    Global type definitions
    -------------------------------------------------------------------------- */