Use a state struct for keeping track of counts

This allows us to define a custom mark/free function using
Data_Wrap_Struct. This in turn makes the code a bit more robust by
marking objects that are counted so they're not garbage collected by
accident.

ext/liballocations/liballocations.c

 #include "liballocations.h"
-st_table *object_counts;
-
 VALUE allocation_tracer;
 VALUE free_tracer;
+VALUE state_const;
 ID id_enabled;
@@ -15,6 +14,8 @@ void newobj_callback(VALUE tracepoint, void* data) {
     st_data_t count = 0;
+    AllocationState *state = allocation_state_get_struct(state_const);
+
     VALUE obj   = rb_tracearg_object(trace_arg);
     VALUE klass = RBASIC_CLASS(obj);
@@ -23,8 +24,8 @@ void newobj_callback(VALUE tracepoint, void* data) {
         return;
     }
-    st_lookup(object_counts, (st_data_t) klass, &count);
-    st_insert(object_counts, (st_data_t) klass, count + 1);
+    st_lookup(state->object_counts, (st_data_t) klass, &count);
+    st_insert(state->object_counts, (st_data_t) klass, count + 1);
 }
 /**
@@ -34,20 +35,21 @@ void newobj_callback(VALUE tracepoint, void* data) {
  * segfault.
  */
 void freeobj_callback(VALUE tracepoint, void* data) {
-    rb_trace_arg_t *trace_arg = rb_tracearg_from_tracepoint(tracepoint);
-
     st_data_t count;
+    rb_trace_arg_t *trace_arg = rb_tracearg_from_tracepoint(tracepoint);
+    AllocationState *state = allocation_state_get_struct(state_const);
+
     VALUE obj   = rb_tracearg_object(trace_arg);
     VALUE klass = RBASIC_CLASS(obj);
-    if ( st_lookup(object_counts, (st_data_t) klass, &count) ) {
-        if ( count > 0 && (count - 1) > 0) {
-            st_insert(object_counts, (st_data_t) klass, count - 1);
+    if ( st_lookup(state->object_counts, (st_data_t) klass, &count) ) {
+        if ( count > 0 && (count - 1) > 0 ) {
+            st_insert(state->object_counts, (st_data_t) klass, count - 1);
         }
         /* Remove the entry if the count is now 0 */
         else {
-            st_delete(object_counts, (st_data_t*) &klass, NULL);
+            st_delete(state->object_counts, (st_data_t*) &klass, NULL);
         }
     }
 }
@@ -75,15 +77,16 @@ static int each_count(st_data_t key, st_data_t value, st_data_t hash_ptr) {
  *     Allocations.to_hash -> Hash
  */
 VALUE allocations_to_hash(VALUE self) {
+    AllocationState *state = allocation_state_get_struct(state_const);
+
     st_table *local_counts;
     VALUE hash;
-    if ( !object_counts ) {
+    if ( !state->object_counts ) {
         return rb_hash_new();
     }
-    local_counts = st_copy(object_counts);
-
+    local_counts = allocation_state_copy_table(state);
     hash = rb_hash_new();
     st_foreach(local_counts, each_count, (st_data_t) hash);
@@ -100,11 +103,13 @@ VALUE allocations_to_hash(VALUE self) {
  *     Allocations.start -> nil
  */
 VALUE allocations_start(VALUE self) {
+    AllocationState *state = allocation_state_get_struct(state_const);
+
     if ( rb_ivar_get(self, id_enabled) == Qtrue ) {
         return Qnil;
     }
-    object_counts = st_init_numtable();
+    allocation_state_allocate_counts(state);
     rb_ivar_set(self, id_enabled, Qtrue);
@@ -121,6 +126,8 @@ VALUE allocations_start(VALUE self) {
  *     Allocations.stop -> nil
  */
 VALUE allocations_stop(VALUE self) {
+    AllocationState *state = allocation_state_get_struct(state_const);
+
     if ( rb_ivar_get(self, id_enabled) != Qtrue ) {
         return Qnil;
     }
@@ -128,11 +135,7 @@ VALUE allocations_stop(VALUE self) {
     rb_tracepoint_disable(allocation_tracer);
     rb_tracepoint_disable(free_tracer);
-    if ( object_counts ) {
-        st_free_table(object_counts);
-    }
-
-    object_counts = NULL;
+    allocation_state_reset_counts(state);
     rb_ivar_set(self, id_enabled, Qfalse);
@@ -173,4 +176,8 @@ void Init_liballocations() {
     rb_define_const(mAllocations, "ALLOCATION_TRACER", allocation_tracer);
     rb_define_const(mAllocations, "FREE_TRACER", free_tracer);
+
+    Init_allocations_state();
+
+    state_const = rb_const_get(mAllocations, rb_intern("STATE"));
 }

ext/liballocations/liballocations.h

@@ -5,6 +5,8 @@
 #include <ruby/debug.h>
 #include <ruby/st.h>
+#include "state.h"
+
 void Init_liballocations();
 #endif

ext/liballocations/state.c

+#include "state.h"
+
+static int mark_each_entry(st_data_t key, st_data_t value, st_data_t data) {
+    rb_gc_mark((VALUE) key);
+
+    return ST_CONTINUE;
+}
+
+// Initializes the table used for storing object counts.
+void allocation_state_allocate_counts(AllocationState *state) {
+    state->object_counts = st_init_numtable();
+}
+
+// Resets the table used for storing object counts.
+void allocation_state_reset_counts(AllocationState *state) {
+    if ( state->object_counts ) {
+        st_free_table(state->object_counts);
+    }
+
+    state->object_counts = NULL;
+}
+
+// Returns the AllocationState wrapped by `object`
+AllocationState *allocation_state_get_struct(VALUE object) {
+    AllocationState *state;
+
+    Data_Get_Struct(object, AllocationState, state);
+
+    return state;
+}
+
+st_table *allocation_state_copy_table(AllocationState *state) {
+    return st_copy(state->object_counts);
+}
+
+void allocation_state_mark(AllocationState *state) {
+    if ( state->object_counts ) {
+        st_foreach(state->object_counts, mark_each_entry, (st_data_t) NULL);
+    }
+}
+
+void allocation_state_free(AllocationState *state) {
+    allocation_state_reset_counts(state);
+
+    free(state);
+}
+
+VALUE allocation_state_allocate(VALUE klass) {
+    AllocationState *state = ALLOC(AllocationState);
+
+    state->object_counts = NULL;
+
+    return Data_Wrap_Struct(klass, allocation_state_mark, allocation_state_free,
+                            state);
+}
+
+void Init_allocations_state() {
+    VALUE mAllocations = rb_const_get(rb_cObject, rb_intern("Allocations"));
+    VALUE cState = rb_define_class_under(mAllocations, "State", rb_cObject);
+
+    rb_define_alloc_func(cState, allocation_state_allocate);
+
+    rb_define_const(mAllocations, "STATE",
+                    rb_funcall(cState, rb_intern("new"), 0));
+}

ext/liballocations/state.h

+#ifndef ALLOCATIONS_STATE_H
+#define ALLOCATIONS_STATE_H
+
+#include "liballocations.h"
+
+typedef struct {
+    st_table *object_counts
+} AllocationState;
+
+extern void allocation_state_reset_counts(AllocationState*);
+extern void allocation_state_allocate_counts(AllocationState*);
+extern AllocationState *allocation_state_get_struct(VALUE object);
+extern st_table *allocation_state_copy_table(AllocationState*);
+
+extern void Init_allocations_state();
+
+#endif