#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));
}
