#include "liballocations.h"

st_table *object_counts;

VALUE allocation_tracer;
VALUE free_tracer;

ID id_enabled;

/**
 * Called whenever a new Ruby object is allocated.
 */
void newobj_callback(VALUE tracepoint, void* data) {
    rb_trace_arg_t *trace_arg = rb_tracearg_from_tracepoint(tracepoint);

    st_data_t count = 0;

    VALUE obj   = rb_tracearg_object(trace_arg);
    VALUE klass = RBASIC_CLASS(obj);

    /* These aren't actually allocated so there's no point in tracking them. */
    if ( klass == Qtrue || klass == Qfalse || klass == Qnil ) {
        return;
    }

    st_lookup(object_counts, (st_data_t) klass, &count);
    st_insert(object_counts, (st_data_t) klass, count + 1);
}

/**
 * Called whenever a Ruby object is about to be released.
 *
 * Important: any Ruby allocations in this function will cause CRuby to
 * segfault.
 */
void freeobj_callback(VALUE tracepoint, void* data) {
    rb_trace_arg_t *trace_arg = rb_tracearg_from_tracepoint(tracepoint);

    st_data_t count;

    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);
        }
        /* Remove the entry if the count is now 0 */
        else {
            st_delete(object_counts, (st_data_t*) &klass, NULL);
        }
    }
}

/**
 * Copies every value in an st_table to a given Ruby Hash.
 */
static int each_count(st_data_t key, st_data_t value, st_data_t hash_ptr) {
    VALUE vkey = (VALUE) key;

    if ( rb_obj_is_kind_of(key, rb_cClass) == Qtrue ) {
        rb_hash_aset((VALUE) hash_ptr, vkey, INT2NUM(value));
    }

    return ST_CONTINUE;
}

/**
 * Returns a Hash containing the current allocation statistics.
 *
 * The returned Hash contains its own copy of the statistics, any further object
 * allocations/frees will not modify said Hash.
 *
 * call-seq:
 *     Allocations.to_hash -> Hash
 */
VALUE allocations_to_hash(VALUE self) {
    st_table *local_counts;
    VALUE hash;

    if ( !object_counts ) {
        return rb_hash_new();
    }

    local_counts = st_copy(object_counts);

    hash = rb_hash_new();

    st_foreach(local_counts, each_count, (st_data_t) hash);

    st_free_table(local_counts);

    return hash;
}

/**
 * Starts the counting of object allocations.
 *
 * call-seq:
 *     Allocations.start -> nil
 */
VALUE allocations_start(VALUE self) {
    if ( rb_ivar_get(self, id_enabled) == Qtrue ) {
        return Qnil;
    }

    object_counts = st_init_numtable();

    rb_ivar_set(self, id_enabled, Qtrue);

    rb_tracepoint_enable(allocation_tracer);
    rb_tracepoint_enable(free_tracer);

    return Qnil;
}

/**
 * Stops the counting of object allocations and clears the current statistics.
 *
 * call-seq:
 *     Allocations.stop -> nil
 */
VALUE allocations_stop(VALUE self) {
    if ( rb_ivar_get(self, id_enabled) != Qtrue ) {
        return Qnil;
    }

    rb_tracepoint_disable(allocation_tracer);
    rb_tracepoint_disable(free_tracer);

    if ( object_counts ) {
        st_free_table(object_counts);
    }

    object_counts = NULL;

    rb_ivar_set(self, id_enabled, Qfalse);

    return Qnil;
}

/**
 * Returns true if tracking allocations has been enabled, false otherwise.
 *
 * call-seq:
 *     Allocations.enabled? -> true/false
 */
VALUE allocations_enabled_p(VALUE self) {
    VALUE enabled = Qfalse;

    if ( rb_ivar_get(self, id_enabled) == Qtrue ) {
        enabled = Qtrue;
    }

    return enabled;
}

void Init_liballocations() {
    VALUE mAllocations = rb_define_module_under(rb_cObject, "Allocations");

    allocation_tracer = rb_tracepoint_new(Qnil, RUBY_INTERNAL_EVENT_NEWOBJ,
                                          newobj_callback, NULL);

    free_tracer = rb_tracepoint_new(Qnil, RUBY_INTERNAL_EVENT_FREEOBJ,
                                    freeobj_callback, NULL);

    id_enabled = rb_intern("enabled");

    rb_define_singleton_method(mAllocations, "to_hash", allocations_to_hash, 0);
    rb_define_singleton_method(mAllocations, "start", allocations_start, 0);
    rb_define_singleton_method(mAllocations, "stop", allocations_stop, 0);
    rb_define_singleton_method(mAllocations, "enabled?", allocations_enabled_p, 0);

    rb_define_const(mAllocations, "ALLOCATION_TRACER", allocation_tracer);
    rb_define_const(mAllocations, "FREE_TRACER", free_tracer);
}
