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

@@ -418,7 +418,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);
@@ -930,7 +930,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,47 @@
     #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(p)
+#endif
+
+/**
+ * Convenience macro to call xfree_sized with a pointer to a fixed size object.
+ * Example:
+ *      pm_call_node_t *target = xmalloc(sizeof(pm_call_node_t));
+ *      ...
+ *      PM_FREE_SIZED(target);
+ */
+#define PM_FREE_SIZED(p) xfree_sized((p), (sizeof(*(p))))
+
+/**
+ * Convenience macro to call xfree_sized with a pointer to an array of fixed size objects.
+ * Example:
+ *      size_t capacity = 12;
+ *      pm_call_node_t *targets = xcalloc(capacity, sizeof(pm_call_node_t));
+ *      ...
+ *      PM_FREE_SIZED_N(target, capacity);
+ */
+#define PM_FREE_SIZED_N(p, n) xfree_sized((p), (sizeof(*(p))) * n)
+
+#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, 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

@@ -47,6 +47,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/prism.c

@@ -286,7 +286,7 @@ lex_mode_pop(pm_parser_t *parser) {
     } else {
         parser->lex_modes.index--;
         pm_lex_mode_t *prev = parser->lex_modes.current->prev;
-        xfree(parser->lex_modes.current);
+        PM_FREE_SIZED(parser->lex_modes.current);
         parser->lex_modes.current = prev;
     }
 }
@@ -777,7 +777,7 @@ pm_parser_scope_shareable_constant_set(pm_parser_t *parser, pm_shareable_constan
 static void
 pm_locals_free(pm_locals_t *locals) {
     if (locals->capacity > 0) {
-        xfree(locals->locals);
+        PM_FREE_SIZED_N(locals->locals, locals->capacity);
     }
 }
 
@@ -2883,7 +2883,7 @@ pm_call_and_write_node_create(pm_parser_t *parser, pm_call_node_t *target, const
     // Here we're going to free the target, since it is no longer necessary.
     // However, we don't want to call `pm_node_destroy` because we want to keep
     // around all of its children since we just reused them.
-    xfree(target);
+    PM_FREE_SIZED(target);
 
     return node;
 }
@@ -2937,7 +2937,7 @@ pm_index_and_write_node_create(pm_parser_t *parser, pm_call_node_t *target, cons
     // Here we're going to free the target, since it is no longer necessary.
     // However, we don't want to call `pm_node_destroy` because we want to keep
     // around all of its children since we just reused them.
-    xfree(target);
+    PM_FREE_SIZED(target);
 
     return node;
 }