Add xfree_sized and xrealloc_sized

Doxyfile

@@ -24,6 +24,7 @@ OUTPUT_DIRECTORY       = doc
 JAVADOC_AUTOBRIEF      = YES
 OPTIMIZE_OUTPUT_FOR_C  = YES
 INPUT                  = src src/util include include/prism include/prism/util
+EXCLUDE                = include/prism/debug_allocator.h
 HTML_OUTPUT            = c
 SORT_MEMBER_DOCS       = NO
 GENERATE_LATEX         = NO

ext/prism/extension.c

@@ -413,7 +413,7 @@ dump(int argc, VALUE *argv, VALUE self) {
     if (options.freeze) rb_obj_freeze(value);
 
 #ifdef PRISM_BUILD_DEBUG
-    xfree(dup);
+    xfree_sized(dup, length);
 #endif
 
     pm_string_free(&input);
@@ -929,7 +929,7 @@ parse(int argc, VALUE *argv, VALUE self) {
     VALUE value = parse_input(&input, &options);
 
 #ifdef PRISM_BUILD_DEBUG
-    xfree(dup);
+    xfree_sized(dup, length);
 #endif
 
     pm_string_free(&input);

include/prism/debug_allocator.h

@@ -0,0 +1,99 @@
+/**
+ * @file debug_allocator.h
+ *
+ * Decorate allocation function to ensure sizes are correct.
+ */
+#ifndef PRISM_DEBUG_ALLOCATOR_H
+#define PRISM_DEBUG_ALLOCATOR_H
+
+#include <string.h>
+#include <stdio.h>
+#include <stdlib.h>
+
+static inline void *
+pm_debug_malloc(size_t size)
+{
+    size_t *memory = xmalloc(size + sizeof(size_t));
+    memory[0] = size;
+    return memory + 1;
+}
+
+static inline void *
+pm_debug_calloc(size_t nmemb, size_t size)
+{
+    size_t total_size = nmemb * size;
+    void *ptr = pm_debug_malloc(total_size);
+    memset(ptr, 0, total_size);
+    return ptr;
+}
+
+static inline void *
+pm_debug_realloc(void *ptr, size_t size)
+{
+    if (ptr == NULL) {
+        return pm_debug_malloc(size);
+    }
+
+    size_t *memory = (size_t *)ptr;
+    void *raw_memory = memory - 1;
+    memory = (size_t *)xrealloc(raw_memory, size + sizeof(size_t));
+    memory[0] = size;
+    return memory + 1;
+}
+
+static inline void
+pm_debug_free(void *ptr)
+{
+    if (ptr != NULL) {
+        size_t *memory = (size_t *)ptr;
+        xfree(memory - 1);
+    }
+}
+
+static inline void
+pm_debug_free_sized(void *ptr, size_t old_size)
+{
+    if (ptr != NULL) {
+        size_t *memory = (size_t *)ptr;
+        if (old_size != memory[-1]) {
+            fprintf(stderr, "[BUG] buffer %p was allocated with size %lu but freed with size %lu\n", ptr, memory[-1], old_size);
+            abort();
+        }
+        xfree_sized(memory - 1, old_size + sizeof(size_t));
+    }
+}
+
+static inline void *
+pm_debug_realloc_sized(void *ptr, size_t size, size_t old_size)
+{
+    if (ptr == NULL) {
+        if (old_size != 0) {
+            fprintf(stderr, "[BUG] realloc_sized called with NULL pointer and old size %lu\n", old_size);
+            abort();
+        }
+        return pm_debug_malloc(size);
+    }
+
+    size_t *memory = (size_t *)ptr;
+    if (old_size != memory[-1]) {
+        fprintf(stderr, "[BUG] buffer %p was allocated with size %lu but realloced with size %lu\n", ptr, memory[-1], old_size);
+        abort();
+    }
+    return pm_debug_realloc(ptr, size);
+}
+
+#undef xmalloc
+#undef xrealloc
+#undef xcalloc
+#undef xfree
+#undef xrealloc_sized
+#undef xfree_sized
+
+#define xmalloc          pm_debug_malloc
+#define xrealloc         pm_debug_realloc
+#define xcalloc          pm_debug_calloc
+#define xfree            pm_debug_free
+#define xrealloc_sized   pm_debug_realloc_sized
+#define xfree_sized      pm_debug_free_sized
+
+#endif

include/prism/defines.h

@@ -161,10 +161,12 @@
  * ```
  * #ifndef PRISM_XALLOCATOR_H
  * #define PRISM_XALLOCATOR_H
- * #define xmalloc      my_malloc
- * #define xrealloc     my_realloc
- * #define xcalloc      my_calloc
- * #define xfree        my_free
+ * #define xmalloc          my_malloc
+ * #define xrealloc         my_realloc
+ * #define xcalloc          my_calloc
+ * #define xfree            my_free
+ * #define xrealloc_sized   my_realloc_sized // (optional)
+ * #define xfree_sized      my_free_sized    // (optional)
  * #endif
  * ```
  */
@@ -204,6 +206,28 @@
     #endif
 #endif
 
+#ifndef xfree_sized
+/**
+ * The free_sized function that should be used. This can be overridden with the
+ * PRISM_XALLOCATOR define.
+ * If not defined, defaults to calling xfree.
+ */
+    #define xfree_sized(p, s) xfree(((void)(s), (p)))
+#endif
+
+#ifndef xrealloc_sized
+/**
+ * The xrealloc_sized function that should be used. This can be overridden with the
+ * PRISM_XALLOCATOR define.
+ * If not defined, defaults to calling xrealloc.
+ */
+    #define xrealloc_sized(p, ns, os) xrealloc((p), ((void)(os), (ns)))
+#endif
+
+#ifdef PRISM_BUILD_DEBUG
+    #include "prism/debug_allocator.h"
+#endif
+
 /**
  * If PRISM_BUILD_MINIMAL is defined, then we're going to define every possible
  * switch that will turn off certain features of prism.

prism.gemspec

@@ -48,6 +48,7 @@ Gem::Specification.new do |spec|
     "ext/prism/extension.h",
     "include/prism.h",
     "include/prism/ast.h",
+    "include/prism/debug_allocator.h",
     "include/prism/defines.h",
     "include/prism/diagnostic.h",
     "include/prism/encoding.h",

src/options.c

@@ -226,10 +226,10 @@ pm_options_free(pm_options_t *options) {
             pm_string_free(&scope->locals[local_index]);
         }
 
-        xfree(scope->locals);
+        xfree_sized(scope->locals, scope->locals_count * sizeof(pm_string_t));
     }
 
-    xfree(options->scopes);
+    xfree_sized(options->scopes, options->scopes_count * sizeof(pm_options_scope_t));
 }
 
 /**