msgpack 1.7.1 → 1.7.3

data/ext/msgpack/packer.c

@@ -17,16 +17,11 @@
  */
 
 #include "packer.h"
+#include "buffer_class.h"
 
-static ID s_call;
-
-void msgpack_packer_static_init(void)
-{
-    s_call = rb_intern("call");
-}
-
-void msgpack_packer_static_destroy(void)
-{ }
+#if !defined(HAVE_RB_PROC_CALL_WITH_BLOCK)
+#define rb_proc_call_with_block(recv, argc, argv, block) rb_funcallv(recv, rb_intern("call"), argc, argv)
+#endif
 
 void msgpack_packer_init(msgpack_packer_t* pk)
 {
@@ -100,14 +95,13 @@ struct msgpack_call_proc_args_t;
 typedef struct msgpack_call_proc_args_t msgpack_call_proc_args_t;
 struct msgpack_call_proc_args_t {
     VALUE proc;
-    VALUE arg;
-    VALUE packer;
+    VALUE args[2];
 };
 
 VALUE msgpack_packer_try_calling_proc(VALUE value)
 {
     msgpack_call_proc_args_t *args = (msgpack_call_proc_args_t *)value;
-    return rb_funcall(args->proc, s_call, 2, args->arg, args->packer);
+    return rb_proc_call_with_block(args->proc, 2, args->args, Qnil);
 }
 
 bool msgpack_packer_try_write_with_ext_type_lookup(msgpack_packer_t* pk, VALUE v)
@@ -121,11 +115,13 @@ bool msgpack_packer_try_write_with_ext_type_lookup(msgpack_packer_t* pk, VALUE v
     }
 
     if(ext_flags & MSGPACK_EXT_RECURSIVE) {
+        VALUE held_buffer = MessagePack_Buffer_hold(&pk->buffer);
+
         msgpack_buffer_t parent_buffer = pk->buffer;
         msgpack_buffer_init(PACKER_BUFFER_(pk));
 
         int exception_occured = 0;
-        msgpack_call_proc_args_t args = { proc, v, pk->to_msgpack_arg };
+        msgpack_call_proc_args_t args = { proc, { v, pk->to_msgpack_arg } };
         rb_protect(msgpack_packer_try_calling_proc, (VALUE)&args, &exception_occured);
 
         if (exception_occured) {
@@ -139,8 +135,10 @@ bool msgpack_packer_try_write_with_ext_type_lookup(msgpack_packer_t* pk, VALUE v
             pk->buffer = parent_buffer;
             msgpack_packer_write_ext(pk, ext_type, payload);
         }
+
+        RB_GC_GUARD(held_buffer);
     } else {
-        VALUE payload = rb_funcall(proc, s_call, 1, v);
+        VALUE payload = rb_proc_call_with_block(proc, 1, &v, Qnil);
         StringValue(payload);
         msgpack_packer_write_ext(pk, ext_type, payload);
     }

data/ext/msgpack/packer.h

@@ -47,10 +47,6 @@ struct msgpack_packer_t {
 
 #define PACKER_BUFFER_(pk) (&(pk)->buffer)
 
-void msgpack_packer_static_init(void);
-
-void msgpack_packer_static_destroy(void);
-
 void msgpack_packer_init(msgpack_packer_t* pk);
 
 void msgpack_packer_destroy(msgpack_packer_t* pk);

data/ext/msgpack/packer_class.c

@@ -106,8 +106,8 @@ VALUE MessagePack_Packer_initialize(int argc, VALUE* argv, VALUE self)
 
     msgpack_packer_t *pk = MessagePack_Packer_get(self);
 
-    msgpack_packer_ext_registry_init(&pk->ext_registry);
-    pk->buffer_ref = Qnil;
+    msgpack_packer_ext_registry_init(self, &pk->ext_registry);
+    pk->buffer_ref = MessagePack_Buffer_wrap(PACKER_BUFFER_(pk), self);
 
     MessagePack_Buffer_set_options(PACKER_BUFFER_(pk), io, options);
 
@@ -352,7 +352,7 @@ static VALUE Packer_registered_types_internal(VALUE self)
     return rb_hash_new();
 }
 
-static VALUE Packer_register_type(int argc, VALUE* argv, VALUE self)
+static VALUE Packer_register_type_internal(VALUE self, VALUE rb_ext_type, VALUE ext_module, VALUE proc)
 {
     if (OBJ_FROZEN(self)) {
         rb_raise(rb_eFrozenError, "can't modify frozen MessagePack::Packer");
@@ -360,38 +360,12 @@ static VALUE Packer_register_type(int argc, VALUE* argv, VALUE self)
 
     msgpack_packer_t *pk = MessagePack_Packer_get(self);
 
-    int ext_type;
-    VALUE ext_module;
-    VALUE proc;
-    VALUE arg;
-
-    switch (argc) {
-    case 2:
-        /* register_type(0x7f, Time) {|obj| block... } */
-        rb_need_block();
-        proc = rb_block_lambda();
-        arg = proc;
-        break;
-    case 3:
-        /* register_type(0x7f, Time, :to_msgpack_ext) */
-        arg = argv[2];
-        proc = rb_funcall(arg, rb_intern("to_proc"), 0);
-        break;
-    default:
-        rb_raise(rb_eArgError, "wrong number of arguments (%d for 2..3)", argc);
-    }
-
-    ext_type = NUM2INT(argv[0]);
+    int ext_type = NUM2INT(rb_ext_type);
     if(ext_type < -128 || ext_type > 127) {
         rb_raise(rb_eRangeError, "integer %d too big to convert to `signed char'", ext_type);
     }
 
-    ext_module = argv[1];
-    if(rb_type(ext_module) != T_MODULE && rb_type(ext_module) != T_CLASS) {
-        rb_raise(rb_eArgError, "expected Module/Class but found %s.", rb_obj_classname(ext_module));
-    }
-
-    msgpack_packer_ext_registry_put(&pk->ext_registry, ext_module, ext_type, 0, proc, arg);
+    msgpack_packer_ext_registry_put(self, &pk->ext_registry, ext_module, ext_type, 0, proc);
 
     if (ext_module == rb_cSymbol) {
         pk->has_symbol_ext_type = true;
@@ -424,10 +398,6 @@ void MessagePack_Packer_module_init(VALUE mMessagePack)
     s_write = rb_intern("write");
 
     sym_compatibility_mode = ID2SYM(rb_intern("compatibility_mode"));
-
-    msgpack_packer_static_init();
-    msgpack_packer_ext_registry_static_init();
-
     cMessagePack_Packer = rb_define_class_under(mMessagePack, "Packer", rb_cObject);
 
     rb_define_alloc_func(cMessagePack_Packer, MessagePack_Packer_alloc);
@@ -466,7 +436,7 @@ void MessagePack_Packer_module_init(VALUE mMessagePack)
     rb_define_method(cMessagePack_Packer, "to_a", Packer_to_a, 0);
 
     rb_define_private_method(cMessagePack_Packer, "registered_types_internal", Packer_registered_types_internal, 0);
-    rb_define_method(cMessagePack_Packer, "register_type", Packer_register_type, -1);
+    rb_define_method(cMessagePack_Packer, "register_type_internal", Packer_register_type_internal, 3);
 
     rb_define_method(cMessagePack_Packer, "full_pack", Packer_full_pack, 0);
 }

data/ext/msgpack/packer_ext_registry.c

@@ -18,20 +18,10 @@
 
 #include "packer_ext_registry.h"
 
-static ID s_call;
-
-void msgpack_packer_ext_registry_static_init(void)
-{
-    s_call = rb_intern("call");
-}
-
-void msgpack_packer_ext_registry_static_destroy(void)
-{ }
-
-void msgpack_packer_ext_registry_init(msgpack_packer_ext_registry_t* pkrg)
+void msgpack_packer_ext_registry_init(VALUE owner, msgpack_packer_ext_registry_t* pkrg)
 {
-    pkrg->hash = Qnil;
-    pkrg->cache = Qnil;
+    RB_OBJ_WRITE(owner, &pkrg->hash, Qnil);
+    RB_OBJ_WRITE(owner, &pkrg->cache, Qnil);
 }
 
 void msgpack_packer_ext_registry_mark(msgpack_packer_ext_registry_t* pkrg)
@@ -40,32 +30,45 @@ void msgpack_packer_ext_registry_mark(msgpack_packer_ext_registry_t* pkrg)
     rb_gc_mark(pkrg->cache);
 }
 
-void msgpack_packer_ext_registry_dup(msgpack_packer_ext_registry_t* src,
+void msgpack_packer_ext_registry_borrow(VALUE owner, msgpack_packer_ext_registry_t* src,
         msgpack_packer_ext_registry_t* dst)
 {
-    if(RTEST(src->hash) && !rb_obj_frozen_p(src->hash)) {
-        dst->hash = rb_hash_dup(src->hash);
-        dst->cache = RTEST(src->cache) ? rb_hash_dup(src->cache) : Qnil;
+    if(RTEST(src->hash)) {
+        if(rb_obj_frozen_p(src->hash)) {
+            // If the type registry is frozen we can safely share it, and share the cache as well.
+            RB_OBJ_WRITE(owner, &dst->hash, src->hash);
+            RB_OBJ_WRITE(owner, &dst->cache, src->cache);
+        } else {
+            RB_OBJ_WRITE(owner, &dst->hash, rb_hash_dup(src->hash));
+            RB_OBJ_WRITE(owner, &dst->cache, NIL_P(src->cache) ? Qnil : rb_hash_dup(src->cache));
+        }
     } else {
-        // If the type registry is frozen we can safely share it, and share the cache as well.
-        dst->hash = src->hash;
-        dst->cache = src->cache;
+        RB_OBJ_WRITE(owner, &dst->hash, Qnil);
+        RB_OBJ_WRITE(owner, &dst->cache, Qnil);
     }
 }
 
-VALUE msgpack_packer_ext_registry_put(msgpack_packer_ext_registry_t* pkrg,
-        VALUE ext_module, int ext_type, int flags, VALUE proc, VALUE arg)
+void msgpack_packer_ext_registry_dup(VALUE owner, msgpack_packer_ext_registry_t* src,
+        msgpack_packer_ext_registry_t* dst)
 {
-    if (!RTEST(pkrg->hash)) {
-        pkrg->hash = rb_hash_new();
+    RB_OBJ_WRITE(owner, &dst->hash, NIL_P(src->hash) ? Qnil : rb_hash_dup(src->hash));
+    RB_OBJ_WRITE(owner, &dst->cache, NIL_P(src->cache) ? Qnil : rb_hash_dup(src->cache));
+}
+
+void msgpack_packer_ext_registry_put(VALUE owner, msgpack_packer_ext_registry_t* pkrg,
+        VALUE ext_module, int ext_type, int flags, VALUE proc)
+{
+    if(NIL_P(pkrg->hash)) {
+        RB_OBJ_WRITE(owner, &pkrg->hash, rb_hash_new());
     }
 
-    if (RTEST(pkrg->cache)) {
+    if(NIL_P(pkrg->cache)) {
+        RB_OBJ_WRITE(owner, &pkrg->cache, rb_hash_new());
+    } else {
         /* clear lookup cache not to miss added type */
         rb_hash_clear(pkrg->cache);
     }
 
-    // TODO: Ruby embeded array limit is 3, merging `proc` and `arg` would be good.
-    VALUE entry = rb_ary_new3(4, INT2FIX(ext_type), proc, arg, INT2FIX(flags));
-    return rb_hash_aset(pkrg->hash, ext_module, entry);
+    VALUE entry = rb_ary_new3(3, INT2FIX(ext_type), proc, INT2FIX(flags));
+    rb_hash_aset(pkrg->hash, ext_module, entry);
 }

data/ext/msgpack/packer_ext_registry.h

@@ -31,22 +31,21 @@ struct msgpack_packer_ext_registry_t {
     VALUE cache; // lookup cache for ext types inherited from a super class
 };
 
-void msgpack_packer_ext_registry_static_init(void);
-
-void msgpack_packer_ext_registry_static_destroy(void);
-
-void msgpack_packer_ext_registry_init(msgpack_packer_ext_registry_t* pkrg);
+void msgpack_packer_ext_registry_init(VALUE owner, msgpack_packer_ext_registry_t* pkrg);
 
 static inline void msgpack_packer_ext_registry_destroy(msgpack_packer_ext_registry_t* pkrg)
 { }
 
 void msgpack_packer_ext_registry_mark(msgpack_packer_ext_registry_t* pkrg);
 
-void msgpack_packer_ext_registry_dup(msgpack_packer_ext_registry_t* src,
+void msgpack_packer_ext_registry_borrow(VALUE owner, msgpack_packer_ext_registry_t* src,
         msgpack_packer_ext_registry_t* dst);
 
-VALUE msgpack_packer_ext_registry_put(msgpack_packer_ext_registry_t* pkrg,
-        VALUE ext_module, int ext_type, int flags, VALUE proc, VALUE arg);
+void msgpack_packer_ext_registry_dup(VALUE owner, msgpack_packer_ext_registry_t* src,
+        msgpack_packer_ext_registry_t* dst);
+
+void msgpack_packer_ext_registry_put(VALUE owner, msgpack_packer_ext_registry_t* pkrg,
+        VALUE ext_module, int ext_type, int flags, VALUE proc);
 
 static int msgpack_packer_ext_find_superclass(VALUE key, VALUE value, VALUE arg)
 {
@@ -68,7 +67,7 @@ static inline VALUE msgpack_packer_ext_registry_fetch(msgpack_packer_ext_registr
     VALUE type = rb_hash_lookup(pkrg->hash, lookup_class);
     if(type != Qnil) {
         *ext_type_result = FIX2INT(rb_ary_entry(type, 0));
-        *ext_flags_result = FIX2INT(rb_ary_entry(type, 3));
+        *ext_flags_result = FIX2INT(rb_ary_entry(type, 2));
         return rb_ary_entry(type, 1);
     }
 
@@ -77,7 +76,7 @@ static inline VALUE msgpack_packer_ext_registry_fetch(msgpack_packer_ext_registr
         VALUE type_inht = rb_hash_lookup(pkrg->cache, lookup_class);
         if(type_inht != Qnil) {
             *ext_type_result = FIX2INT(rb_ary_entry(type_inht, 0));
-            *ext_flags_result = FIX2INT(rb_ary_entry(type_inht, 3));
+            *ext_flags_result = FIX2INT(rb_ary_entry(type_inht, 2));
             return rb_ary_entry(type_inht, 1);
         }
     }
@@ -129,12 +128,9 @@ static inline VALUE msgpack_packer_ext_registry_lookup(msgpack_packer_ext_regist
     VALUE superclass = args[1];
     if(superclass != Qnil) {
         VALUE superclass_type = rb_hash_lookup(pkrg->hash, superclass);
-        if (!RTEST(pkrg->cache)) {
-            pkrg->cache = rb_hash_new();
-        }
         rb_hash_aset(pkrg->cache, lookup_class, superclass_type);
         *ext_type_result = FIX2INT(rb_ary_entry(superclass_type, 0));
-        *ext_flags_result = FIX2INT(rb_ary_entry(superclass_type, 3));
+        *ext_flags_result = FIX2INT(rb_ary_entry(superclass_type, 2));
         return rb_ary_entry(superclass_type, 1);
     }
 

data/ext/msgpack/unpacker.c

@@ -20,11 +20,15 @@
 #include "rmem.h"
 #include "extension_value_class.h"
 #include <assert.h>
+#include <limits.h>
+
+#if !defined(HAVE_RB_PROC_CALL_WITH_BLOCK)
+#define rb_proc_call_with_block(recv, argc, argv, block) rb_funcallv(recv, rb_intern("call"), argc, argv)
+#endif
 
 static int RAW_TYPE_STRING = 256;
 static int RAW_TYPE_BINARY = 257;
-
-static ID s_call;
+static int16_t INITIAL_BUFFER_CAPACITY_MAX = SHRT_MAX;
 
 static msgpack_rmem_t s_stack_rmem;
 
@@ -35,13 +39,16 @@ static inline VALUE rb_hash_new_capa(long capa)
 }
 #endif
 
+static inline int16_t initial_buffer_size(long size)
+{
+    return (size > INITIAL_BUFFER_CAPACITY_MAX) ? INITIAL_BUFFER_CAPACITY_MAX : size;
+}
+
 void msgpack_unpacker_static_init(void)
 {
     assert(sizeof(msgpack_unpacker_stack_entry_t) * MSGPACK_UNPACKER_STACK_CAPACITY <= MSGPACK_RMEM_PAGE_SIZE);
 
     msgpack_rmem_init(&s_stack_rmem);
-
-    s_call = rb_intern("call");
 }
 
 void msgpack_unpacker_static_destroy(void)
@@ -178,7 +185,8 @@ static inline int object_complete_ext(msgpack_unpacker_t* uk, int ext_type, VALU
 
     if(proc != Qnil) {
         VALUE obj;
-        obj = rb_funcall(proc, s_call, 1, str == Qnil ? rb_str_buf_new(0) : str);
+        VALUE arg = (str == Qnil ? rb_str_buf_new(0) : str);
+        obj = rb_proc_call_with_block(proc, 1, &arg, Qnil);
         return object_complete(uk, obj);
     }
 
@@ -308,7 +316,7 @@ static inline int read_raw_body_begin(msgpack_unpacker_t* uk, int raw_type)
             child_stack->parent = uk->stack;
             uk->stack = child_stack;
 
-            obj = rb_funcall(proc, s_call, 1, uk->self);
+            obj = rb_proc_call_with_block(proc, 1, &uk->self, Qnil);
 
             uk->stack = child_stack->parent;
             _msgpack_unpacker_free_stack(child_stack);
@@ -374,14 +382,14 @@ static int read_primitive(msgpack_unpacker_t* uk)
         if(count == 0) {
             return object_complete(uk, rb_ary_new());
         }
-        return _msgpack_unpacker_stack_push(uk, STACK_TYPE_ARRAY, count, rb_ary_new2(count));
+        return _msgpack_unpacker_stack_push(uk, STACK_TYPE_ARRAY, count, rb_ary_new2(initial_buffer_size(count)));
 
     SWITCH_RANGE(b, 0x80, 0x8f)  // FixMap
         int count = b & 0x0f;
         if(count == 0) {
             return object_complete(uk, rb_hash_new());
         }
-        return _msgpack_unpacker_stack_push(uk, STACK_TYPE_MAP_KEY, count*2, rb_hash_new_capa(count));
+        return _msgpack_unpacker_stack_push(uk, STACK_TYPE_MAP_KEY, count*2, rb_hash_new_capa(initial_buffer_size(count)));
 
     SWITCH_RANGE(b, 0xc0, 0xdf)  // Variable
         switch(b) {
@@ -604,7 +612,7 @@ static int read_primitive(msgpack_unpacker_t* uk)
                 if(count == 0) {
                     return object_complete(uk, rb_ary_new());
                 }
-                return _msgpack_unpacker_stack_push(uk, STACK_TYPE_ARRAY, count, rb_ary_new2(count));
+                return _msgpack_unpacker_stack_push(uk, STACK_TYPE_ARRAY, count, rb_ary_new2(initial_buffer_size(count)));
             }
 
         case 0xdd:  // array 32
@@ -614,7 +622,7 @@ static int read_primitive(msgpack_unpacker_t* uk)
                 if(count == 0) {
                     return object_complete(uk, rb_ary_new());
                 }
-                return _msgpack_unpacker_stack_push(uk, STACK_TYPE_ARRAY, count, rb_ary_new2(count));
+                return _msgpack_unpacker_stack_push(uk, STACK_TYPE_ARRAY, count, rb_ary_new2(initial_buffer_size(count)));
             }
 
         case 0xde:  // map 16
@@ -624,7 +632,7 @@ static int read_primitive(msgpack_unpacker_t* uk)
                 if(count == 0) {
                     return object_complete(uk, rb_hash_new());
                 }
-                return _msgpack_unpacker_stack_push(uk, STACK_TYPE_MAP_KEY, count*2, rb_hash_new_capa(count));
+                return _msgpack_unpacker_stack_push(uk, STACK_TYPE_MAP_KEY, count*2, rb_hash_new_capa(initial_buffer_size(count)));
             }
 
         case 0xdf:  // map 32
@@ -634,7 +642,7 @@ static int read_primitive(msgpack_unpacker_t* uk)
                 if(count == 0) {
                     return object_complete(uk, rb_hash_new());
                 }
-                return _msgpack_unpacker_stack_push(uk, STACK_TYPE_MAP_KEY, count*2, rb_hash_new_capa(count));
+                return _msgpack_unpacker_stack_push(uk, STACK_TYPE_MAP_KEY, count*2, rb_hash_new_capa(initial_buffer_size(count)));
             }
 
         default:

data/ext/msgpack/unpacker_class.c

@@ -346,43 +346,19 @@ static VALUE Unpacker_registered_types_internal(VALUE self)
     return mapping;
 }
 
-static VALUE Unpacker_register_type(int argc, VALUE* argv, VALUE self)
+static VALUE Unpacker_register_type_internal(VALUE self, VALUE rb_ext_type, VALUE ext_module, VALUE proc)
 {
     if (OBJ_FROZEN(self)) {
         rb_raise(rb_eFrozenError, "can't modify frozen MessagePack::Unpacker");
     }
 
-    msgpack_unpacker_t *uk = MessagePack_Unpacker_get(self);
-
-    int ext_type;
-    VALUE proc;
-    VALUE arg;
-    VALUE ext_module;
-
-    switch (argc) {
-    case 1:
-        /* register_type(0x7f) {|data| block... } */
-        rb_need_block();
-        proc = rb_block_lambda();
-        arg = proc;
-        ext_module = Qnil;
-        break;
-    case 3:
-        /* register_type(0x7f, Time, :from_msgpack_ext) */
-        ext_module = argv[1];
-        arg = argv[2];
-        proc = rb_obj_method(ext_module, arg);
-        break;
-    default:
-        rb_raise(rb_eArgError, "wrong number of arguments (%d for 1 or 3)", argc);
-    }
-
-    ext_type = NUM2INT(argv[0]);
+    int ext_type = NUM2INT(rb_ext_type);
     if(ext_type < -128 || ext_type > 127) {
         rb_raise(rb_eRangeError, "integer %d too big to convert to `signed char'", ext_type);
     }
 
-    msgpack_unpacker_ext_registry_put(&uk->ext_registry, ext_module, ext_type, 0, proc, arg);
+    msgpack_unpacker_t *uk = MessagePack_Unpacker_get(self);
+    msgpack_unpacker_ext_registry_put(self, &uk->ext_registry, ext_module, ext_type, 0, proc);
 
     return Qnil;
 }
@@ -415,7 +391,6 @@ VALUE MessagePack_Unpacker_new(int argc, VALUE* argv)
 void MessagePack_Unpacker_module_init(VALUE mMessagePack)
 {
     msgpack_unpacker_static_init();
-    msgpack_unpacker_ext_registry_static_init();
 
     mTypeError = rb_define_module_under(mMessagePack, "TypeError");
 
@@ -456,7 +431,7 @@ void MessagePack_Unpacker_module_init(VALUE mMessagePack)
     rb_define_method(cMessagePack_Unpacker, "reset", Unpacker_reset, 0);
 
     rb_define_private_method(cMessagePack_Unpacker, "registered_types_internal", Unpacker_registered_types_internal, 0);
-    rb_define_method(cMessagePack_Unpacker, "register_type", Unpacker_register_type, -1);
+    rb_define_private_method(cMessagePack_Unpacker, "register_type_internal", Unpacker_register_type_internal, 3);
 
     rb_define_method(cMessagePack_Unpacker, "full_unpack", Unpacker_full_unpack, 0);
 }

data/ext/msgpack/unpacker_ext_registry.c

@@ -18,19 +18,6 @@
 
 #include "unpacker_ext_registry.h"
 
-static ID s_call;
-static ID s_dup;
-
-void msgpack_unpacker_ext_registry_static_init(void)
-{
-    s_call = rb_intern("call");
-    s_dup = rb_intern("dup");
-}
-
-
-void msgpack_unpacker_ext_registry_static_destroy(void)
-{ }
-
 void msgpack_unpacker_ext_registry_mark(msgpack_unpacker_ext_registry_t* ukrg)
 {
     if (ukrg) {
@@ -76,11 +63,12 @@ void msgpack_unpacker_ext_registry_release(msgpack_unpacker_ext_registry_t* ukrg
     }
 }
 
-void msgpack_unpacker_ext_registry_put(msgpack_unpacker_ext_registry_t** ukrg,
-        VALUE ext_module, int ext_type, int flags, VALUE proc, VALUE arg)
+void msgpack_unpacker_ext_registry_put(VALUE owner, msgpack_unpacker_ext_registry_t** ukrg,
+        VALUE ext_module, int ext_type, int flags, VALUE proc)
 {
     msgpack_unpacker_ext_registry_t* ext_registry = msgpack_unpacker_ext_registry_cow(*ukrg);
 
-    ext_registry->array[ext_type + 128] = rb_ary_new3(4, ext_module, proc, arg, INT2FIX(flags));
+    VALUE entry = rb_ary_new3(3, ext_module, proc, INT2FIX(flags));
+    RB_OBJ_WRITE(owner, &ext_registry->array[ext_type + 128], entry);
     *ukrg = ext_registry;
 }

data/ext/msgpack/unpacker_ext_registry.h

@@ -31,10 +31,6 @@ struct msgpack_unpacker_ext_registry_t {
     VALUE array[256];
 };
 
-void msgpack_unpacker_ext_registry_static_init(void);
-
-void msgpack_unpacker_ext_registry_static_destroy(void);
-
 void msgpack_unpacker_ext_registry_release(msgpack_unpacker_ext_registry_t* ukrg);
 
 static inline void msgpack_unpacker_ext_registry_borrow(msgpack_unpacker_ext_registry_t* src, msgpack_unpacker_ext_registry_t** dst)
@@ -47,8 +43,8 @@ static inline void msgpack_unpacker_ext_registry_borrow(msgpack_unpacker_ext_reg
 
 void msgpack_unpacker_ext_registry_mark(msgpack_unpacker_ext_registry_t* ukrg);
 
-void msgpack_unpacker_ext_registry_put(msgpack_unpacker_ext_registry_t** ukrg,
-        VALUE ext_module, int ext_type, int flags, VALUE proc, VALUE arg);
+void msgpack_unpacker_ext_registry_put(VALUE owner, msgpack_unpacker_ext_registry_t** ukrg,
+        VALUE ext_module, int ext_type, int flags, VALUE proc);
 
 static inline VALUE msgpack_unpacker_ext_registry_lookup(msgpack_unpacker_ext_registry_t* ukrg,
         int ext_type, int* ext_flags_result)
@@ -56,7 +52,7 @@ static inline VALUE msgpack_unpacker_ext_registry_lookup(msgpack_unpacker_ext_re
     if (ukrg) {
         VALUE entry = ukrg->array[ext_type + 128];
         if (entry != Qnil) {
-            *ext_flags_result = FIX2INT(rb_ary_entry(entry, 3));
+            *ext_flags_result = FIX2INT(rb_ary_entry(entry, 2));
             return rb_ary_entry(entry, 1);
         }
     }

data/lib/msgpack/factory.rb

@@ -2,11 +2,46 @@ module MessagePack
   class Factory
     # see ext for other methods
 
+    def register_type(type, klass, options = { packer: :to_msgpack_ext, unpacker: :from_msgpack_ext })
+      raise FrozenError, "can't modify frozen MessagePack::Factory" if frozen?
+
+      if options
+        options = options.dup
+        case packer = options[:packer]
+        when nil, Proc
+          # all good
+        when String, Symbol
+          options[:packer] = packer.to_sym.to_proc
+        when Method
+          options[:packer] = packer.to_proc
+        when packer.respond_to?(:call)
+          options[:packer] = packer.method(:call).to_proc
+        else
+          raise ::TypeError, "expected :packer argument to be a callable object, got: #{packer.inspect}"
+        end
+
+        case unpacker = options[:unpacker]
+        when nil, Proc
+          # all good
+        when String, Symbol
+          options[:unpacker] = klass.method(unpacker).to_proc
+        when Method
+          options[:unpacker] = unpacker.to_proc
+        when packer.respond_to?(:call)
+          options[:unpacker] = unpacker.method(:call).to_proc
+        else
+          raise ::TypeError, "expected :unpacker argument to be a callable object, got: #{unpacker.inspect}"
+        end
+      end
+
+      register_type_internal(type, klass, options)
+    end
+
     # [ {type: id, class: Class(or nil), packer: arg, unpacker: arg}, ... ]
     def registered_types(selector=:both)
       packer, unpacker = registered_types_internal
-      # packer: Class -> [tid, proc, arg]
-      # unpacker: tid -> [klass, proc, arg]
+      # packer: Class -> [tid, proc, _flags]
+      # unpacker: tid -> [klass, proc, _flags]
 
       list = []
 
@@ -14,27 +49,31 @@ module MessagePack
       when :both
         packer.each_pair do |klass, ary|
           type = ary[0]
-          packer_arg = ary[2]
-          unpacker_arg = nil
-          if unpacker.has_key?(type) && unpacker[type][0] == klass
-            unpacker_arg = unpacker.delete(type)[2]
+          packer_proc = ary[1]
+          unpacker_proc = nil
+          if unpacker.has_key?(type)
+            unpacker_proc = unpacker.delete(type)[1]
           end
-          list << {type: type, class: klass, packer: packer_arg, unpacker: unpacker_arg}
+          list << {type: type, class: klass, packer: packer_proc, unpacker: unpacker_proc}
         end
 
         # unpacker definition only
         unpacker.each_pair do |type, ary|
-          list << {type: type, class: ary[0], packer: nil, unpacker: ary[2]}
+          list << {type: type, class: ary[0], packer: nil, unpacker: ary[1]}
         end
 
       when :packer
         packer.each_pair do |klass, ary|
-          list << {type: ary[0], class: klass, packer: ary[2]}
+          if ary[1]
+            list << {type: ary[0], class: klass, packer: ary[1]}
+          end
         end
 
       when :unpacker
         unpacker.each_pair do |type, ary|
-          list << {type: type, class: ary[0], unpacker: ary[2]}
+          if ary[1]
+            list << {type: type, class: ary[0], unpacker: ary[1]}
+          end
         end
 
       else

data/lib/msgpack/packer.rb

@@ -6,11 +6,16 @@ module MessagePack
     undef_method :dup
     undef_method :clone
 
+    def register_type(type, klass, method_name = nil, &block)
+      raise ArgumentError, "expected Module/Class got: #{klass.inspect}" unless klass.is_a?(Module)
+      register_type_internal(type, klass, block || method_name.to_proc)
+    end
+
     def registered_types
       list = []
 
       registered_types_internal.each_pair do |klass, ary|
-        list << {type: ary[0], class: klass, packer: ary[2]}
+        list << {type: ary[0], class: klass, packer: ary[1]}
       end
 
       list.sort{|a, b| a[:type] <=> b[:type] }

data/lib/msgpack/unpacker.rb

@@ -6,11 +6,20 @@ module MessagePack
     undef_method :dup
     undef_method :clone
 
+    def register_type(type, klass = nil, method_name = nil, &block)
+      if klass && method_name
+        block = klass.method(method_name).to_proc
+      elsif !block_given?
+        raise ArgumentError, "register_type takes either 3 arguments or a block"
+      end
+      register_type_internal(type, klass, block)
+    end
+
     def registered_types
       list = []
 
       registered_types_internal.each_pair do |type, ary|
-        list << {type: type, class: ary[0], unpacker: ary[2]}
+        list << {type: type, class: ary[0], unpacker: ary[1]}
       end
 
       list.sort{|a, b| a[:type] <=> b[:type] }