identity_map.cpp

// Hacky prototype weak identity map extension.
#include <iostream>
#include <string>
#include <map>

#include <ruby/ruby.h>

#define ID_CONST_GET rb_intern("const_get")
#define CONST_GET(scope, constant) (rb_funcall(scope, ID_CONST_GET, 1, rb_str_new2(constant)))

#define CSTRING(v) RSTRING_PTR(TYPE(v) == T_STRING ? v : rb_funcall(v, rb_intern("to_s"), 0))
#define INSPECT(v) RSTRING_PTR(rb_funcall(v, rb_intern("inspect"), 0))

static VALUE mOm;
static VALUE cIdentityMap;

typedef struct IdentityMap {
    std::map<ID, std::string> value_id_key;
    std::map<std::string, VALUE> key_value;
};

static void identity_map_deallocate(IdentityMap *idmap) {
    delete idmap;
}

static VALUE rb_identity_map_alloc(VALUE klass) {
    IdentityMap *idmap = new IdentityMap;
    return Data_Wrap_Struct(klass, NULL, identity_map_deallocate, idmap);
}

static VALUE rb_identity_map_delete_by_id(VALUE self, VALUE id) {
    IdentityMap *idmap;
    Data_Get_Struct(self, IdentityMap, idmap);

    std::string key = idmap->value_id_key[NUM2ULONG(id)];

    // std::count << "delete key:" << key << ", id: " << NUM2ULONG(id) << std::endl;

    idmap->key_value.erase(key);
    idmap->value_id_key.erase(NUM2ULONG(id));

    return Qtrue;
}

static VALUE rb_identity_map_set(VALUE self, VALUE key, VALUE value) {
    IdentityMap *idmap;
    Data_Get_Struct(self, IdentityMap, idmap);

    if (!FL_ABLE(value)) {
        // TODO: Barf because we can't add a finalizer.
    }

    idmap->value_id_key[NUM2ULONG(rb_obj_id(value))] = CSTRING(key);
    idmap->key_value[CSTRING(key)]                   = value;

    // std::cout << "add key:" << INSPECT(key) << ", id:" << NUM2ULONG(rb_obj_id(value)) << std::endl;

    VALUE callback = rb_obj_method(self, ID2SYM(rb_intern("delete_by_id")));
    rb_funcall(CONST_GET(rb_cObject, "ObjectSpace"), rb_intern("define_finalizer"), 2, value, callback);

    return value;
}

static VALUE rb_identity_map_get(VALUE self, VALUE key) {
    IdentityMap *idmap;
    Data_Get_Struct(self, IdentityMap, idmap);
    return idmap->key_value[CSTRING(key)];
}

extern "C" {
    void Init_om(void) {
        mOm          = rb_define_module("Om");
        cIdentityMap = rb_define_class_under(mOm, "IdentityMap", rb_cObject);
        rb_define_alloc_func(cIdentityMap, rb_identity_map_alloc);

        rb_define_method(cIdentityMap, "set", RUBY_METHOD_FUNC(rb_identity_map_set), 2);
        rb_define_method(cIdentityMap, "get", RUBY_METHOD_FUNC(rb_identity_map_get), 1);
        rb_define_private_method(cIdentityMap, "delete_by_id", RUBY_METHOD_FUNC(rb_identity_map_delete_by_id), 1);
    }
}

identity_map.rb

module Om
  # Om::IdentityMap
  #
  # Really a weak hash set implementation.
  #--
  # TODO: Is 'hash set' the real name for a hash where both the keys and values must be unique?
  class IdentityMap
    def initialize
      @cache, @reverse_cache, @finalize = {}, {}, method(:finalize)
    end

    def get key
      value_id = @cache[key]
      return ObjectSpace._id2ref(value_id) unless value_id.nil?
      nil
    end

    #--
    # TODO: Barf if the value.object_id already exists in the cache.
    def set key, value
      @reverse_cache[value.object_id] = key
      @cache[key]                     = value.object_id
      ObjectSpace.define_finalizer(value, @finalize)
    end

    private
      def finalize value_id
        @cache.delete @reverse_cache.delete value_id
      end
  end
end # Om