Extended type support

ext/msgpack/core_ext.c

@@ -20,29 +20,6 @@
 #include "packer.h"
 #include "packer_class.h"
 
-static inline VALUE delegete_to_pack(int argc, VALUE* argv, VALUE self)
-{
-    if(argc == 0) {
-        return MessagePack_pack(1, &self);
-    } else if(argc == 1) {
-        /* write to io */
-        VALUE argv2[2];
-        argv2[0] = self;
-        argv2[1] = argv[0];
-        return MessagePack_pack(2, argv2);
-    } else {
-        rb_raise(rb_eArgError, "wrong number of arguments (%d for 0..1)", argc);
-    }
-}
-
-#define ENSURE_PACKER(argc, argv, packer, pk) \
-    if(argc != 1 || rb_class_of(argv[0]) != cMessagePack_Packer) { \
-        return delegete_to_pack(argc, argv, self); \
-    } \
-    VALUE packer = argv[0]; \
-    msgpack_packer_t *pk; \
-    Data_Get_Struct(packer, msgpack_packer_t, pk);
-
 static VALUE NilClass_to_msgpack(int argc, VALUE* argv, VALUE self)
 {
     ENSURE_PACKER(argc, argv, packer, pk);

ext/msgpack/extended_class.c

@@ -0,0 +1,172 @@
+#include "compat.h"
+#include "ruby.h"
+#include "extended_class.h"
+#include "packer.h"
+#include "packer.h"
+#include "packer_class.h"
+
+VALUE cMessagePack_Extended;
+
+#ifdef COMPAT_HAVE_ENCODING  /* see compat.h*/
+static int s_enc_ascii_8bit;
+#endif
+
+#define EXTENDED(from, name) \
+    msgpack_extended_t *name = NULL; \
+    Data_Get_Struct(from, msgpack_extended_t, name); \
+    if(name == NULL) { \
+        rb_raise(rb_eArgError, "NULL found for " # name " when shouldn't be."); \
+    }
+
+static void assert_fixnum_type(VALUE type)
+{
+    Check_Type(type, T_FIXNUM);
+
+    if ((FIX2INT(type) < INT8_MIN) || (FIX2INT(type) > INT8_MAX)) {
+        rb_raise(rb_eRangeError, "type should be <= %d and >= %d", INT8_MIN, INT8_MAX);
+    }
+}
+
+static void Extended_free(void* data)
+{
+    if(data == NULL) {
+        return;
+    }
+    msgpack_extended_t* ext = (msgpack_extended_t*) data;
+    free(ext);
+}
+
+void msgpack_extended_mark(msgpack_extended_t* ext)
+{
+    rb_gc_mark(ext->type);
+    rb_gc_mark(ext->data);
+}
+
+
+void msgpack_extended_init(msgpack_extended_t* ext)
+{
+    memset(ext, 0, sizeof(msgpack_extended_t));
+
+    ext->data = Qnil;
+    ext->type = Qnil;
+}
+
+static VALUE Extended_alloc(VALUE klass)
+{
+    msgpack_extended_t* ext = ALLOC_N(msgpack_extended_t, 1);
+    msgpack_extended_init(ext);
+
+    return Data_Wrap_Struct(klass, msgpack_extended_mark, Extended_free, ext);
+}
+
+static VALUE Extended_initialize(VALUE self, VALUE type, VALUE data)
+{
+    EXTENDED(self, ext);
+
+    assert_fixnum_type(type);
+    Check_Type(data, T_STRING);
+
+#ifdef COMPAT_HAVE_ENCODING
+    // TODO ruby 2.0 has String#b method
+    ENCODING_SET(data, s_enc_ascii_8bit);
+#endif
+
+    ext->type = type;
+    ext->data = data;
+
+    return self;
+}
+
+static VALUE Extended_data(VALUE self)
+{
+    EXTENDED(self, ext);
+    return ext->data;
+}
+
+static VALUE Extended_type(VALUE self)
+{
+    EXTENDED(self, ext);
+    return ext->type;
+}
+
+static VALUE Extended_type_setter(VALUE self, VALUE type)
+{
+    EXTENDED(self, ext);
+    ext->type = type;
+    return ext->type;
+}
+
+static VALUE Extended_to_msgpack(int argc, VALUE* argv, VALUE self)
+{
+    ENSURE_PACKER(argc, argv, packer, pk);
+    msgpack_packer_write_ext_value(pk, self);
+    return packer;
+}
+
+static VALUE Extended_eql(VALUE self, VALUE other)
+{
+    if(rb_class_of(self) != rb_class_of(other)) {
+        return Qfalse;
+    }
+
+    EXTENDED(self, ext);
+    EXTENDED(other, other_ext);
+
+    if(ext->type != other_ext->type) {
+        return Qfalse;
+    }
+
+    unsigned long self_len  = RSTRING_LEN((ext->data));
+    unsigned long other_len = RSTRING_LEN((other_ext->data));
+
+    if(self_len != other_len) {
+        return Qfalse;
+    }
+
+    if(memcmp(StringValuePtr(ext->data), StringValuePtr(other_ext->data), self_len) != 0) {
+        return Qfalse;
+    }
+
+    return Qtrue;
+}
+
+static VALUE MessagePack_register_extension_module_method(VALUE self, VALUE type, VALUE klass)
+{
+    assert_fixnum_type(type);
+    Check_Type(klass, T_CLASS);
+
+    const int itype = FIX2INT(type);
+    const char *klass_name = rb_class2name(klass);
+
+    st_add_direct(msgpack_extension_mappings, itype, (st_data_t)klass_name);
+
+    return Qnil;
+}
+
+void MessagePack_Extended_module_init(VALUE mMessagePack)
+{
+    cMessagePack_Extended = rb_define_class_under(mMessagePack, "ExtensionValue", rb_cObject);
+
+    rb_define_alloc_func(cMessagePack_Extended, Extended_alloc);
+
+    rb_define_method(cMessagePack_Extended, "initialize", Extended_initialize, 2);
+    rb_define_method(cMessagePack_Extended, "type", Extended_type, 0);
+    rb_define_method(cMessagePack_Extended, "type=", Extended_type_setter, 1);
+    rb_define_method(cMessagePack_Extended, "data", Extended_data, 0);
+    rb_define_method(cMessagePack_Extended, "to_msgpack", Extended_to_msgpack, -1);
+    rb_define_method(cMessagePack_Extended, "==", Extended_eql, 1);
+    rb_define_method(cMessagePack_Extended, "eql?", Extended_eql, 1);
+
+    msgpack_extension_mappings = st_init_numtable();
+
+    rb_define_module_function(mMessagePack,
+            "register_extension",
+            MessagePack_register_extension_module_method,
+            2);
+
+#ifdef COMPAT_HAVE_ENCODING
+    s_enc_ascii_8bit = rb_ascii8bit_encindex();
+#endif
+
+}
+

ext/msgpack/extended_class.h

@@ -0,0 +1,22 @@
+#ifndef MSGPACK_RUBY_EXTENDED_CLASS_H__
+#define MSGPACK_RUBY_EXTENDED_CLASS_H__
+
+#include "compat.h"
+
+struct msgpack_extended_t;
+typedef struct msgpack_extended_t msgpack_extended_t;
+
+st_table *msgpack_extension_mappings;
+
+struct msgpack_extended_t {
+    VALUE type;
+    VALUE data;
+};
+
+extern VALUE cMessagePack_Extended;
+
+void MessagePack_Extended_module_init(VALUE mMessagePack);
+
+void MessagePack_Extended_initialize(VALUE self, VALUE type, VALUE data);
+
+#endif

ext/msgpack/packer.h

@@ -316,6 +316,39 @@ static inline void msgpack_packer_write_map_header(msgpack_packer_t* pk, unsigne
     }
 }
 
+static inline void msgpack_packer_write_ext_header(msgpack_packer_t* pk, unsigned int n, int8_t type)
+{
+    if(n == 1) {
+        msgpack_buffer_ensure_writable(PACKER_BUFFER_(pk), 2);
+        msgpack_buffer_write_2(PACKER_BUFFER_(pk), 0xd4, type);
+    } else if(n == 2) {
+        msgpack_buffer_ensure_writable(PACKER_BUFFER_(pk), 2);
+        msgpack_buffer_write_2(PACKER_BUFFER_(pk), 0xd5, type);
+    } else if(n == 4) {
+        msgpack_buffer_ensure_writable(PACKER_BUFFER_(pk), 2);
+        msgpack_buffer_write_2(PACKER_BUFFER_(pk), 0xd6, type);
+    } else if(n == 8) {
+        msgpack_buffer_ensure_writable(PACKER_BUFFER_(pk), 2);
+        msgpack_buffer_write_2(PACKER_BUFFER_(pk), 0xd7, type);
+    } else if(n == 16) {
+        msgpack_buffer_ensure_writable(PACKER_BUFFER_(pk), 2);
+        msgpack_buffer_write_2(PACKER_BUFFER_(pk), 0xd8, type);
+    } else if(n < 256) {
+        msgpack_buffer_ensure_writable(PACKER_BUFFER_(pk), 3);
+        msgpack_buffer_write_2(PACKER_BUFFER_(pk), 0xc7, (uint8_t)n);
+        msgpack_buffer_write_1(PACKER_BUFFER_(pk), type);
+    } else if (n < 65536) {
+        msgpack_buffer_ensure_writable(PACKER_BUFFER_(pk), 4);
+        uint16_t be = _msgpack_be16(n);
+        msgpack_buffer_write_byte_and_data(PACKER_BUFFER_(pk), 0xc8, (const void*)&be, 2);
+        msgpack_buffer_write_1(PACKER_BUFFER_(pk), type);
+    } else {
+        msgpack_buffer_ensure_writable(PACKER_BUFFER_(pk), 6);
+        uint32_t be = _msgpack_be32(n);
+        msgpack_buffer_write_byte_and_data(PACKER_BUFFER_(pk), 0xc9, (const void*)&be, 4);
+        msgpack_buffer_write_1(PACKER_BUFFER_(pk), type);
+    }
+}
 
 void _msgpack_packer_write_string_to_io(msgpack_packer_t* pk, VALUE string);
 
@@ -350,6 +383,23 @@ static inline void msgpack_packer_write_symbol_value(msgpack_packer_t* pk, VALUE
 #endif
 }
 
+static inline void msgpack_packer_write_ext_value(msgpack_packer_t* pk, VALUE v)
+{
+    int8_t type = NUM2INT(rb_funcall(v, rb_intern("type"), 0));
+    VALUE data  = rb_funcall(v, rb_intern("data"), 0);
+
+    /* actual return type of RSTRING_LEN is long */
+    unsigned long len = RSTRING_LEN(data);
+    if(len > 0xffffffffUL) {
+        // TODO rb_eArgError?
+        rb_raise(rb_eArgError, "size of string is too long to pack: %lu bytes should be <= %lu", len, 0xffffffffUL);
+    }
+
+    msgpack_packer_write_ext_header(pk, (unsigned int)len, type);
+    msgpack_buffer_append_string(PACKER_BUFFER_(pk), data);
+}
+
+
 static inline void msgpack_packer_write_fixnum_value(msgpack_packer_t* pk, VALUE v)
 {
 #ifdef JRUBY

ext/msgpack/packer_class.c

@@ -128,6 +128,13 @@ static VALUE Packer_write_map_header(VALUE self, VALUE n)
     return self;
 }
 
+static VALUE Packer_write_ext_header(VALUE self, VALUE n, VALUE type)
+{
+    PACKER(self, pk);
+    msgpack_packer_write_ext_header(pk, NUM2UINT(n), NUM2INT(type));
+    return self;
+}
+
 static VALUE Packer_flush(VALUE self)
 {
     PACKER(self, pk);
@@ -278,6 +285,7 @@ void MessagePack_Packer_module_init(VALUE mMessagePack)
     rb_define_method(cMessagePack_Packer, "write_nil", Packer_write_nil, 0);
     rb_define_method(cMessagePack_Packer, "write_array_header", Packer_write_array_header, 1);
     rb_define_method(cMessagePack_Packer, "write_map_header", Packer_write_map_header, 1);
+    rb_define_method(cMessagePack_Packer, "write_ext_header", Packer_write_ext_header, 2);
     rb_define_method(cMessagePack_Packer, "flush", Packer_flush, 0);
 
     /* delegation methods */

ext/msgpack/packer_class.h

@@ -26,5 +26,28 @@ void MessagePack_Packer_module_init(VALUE mMessagePack);
 
 VALUE MessagePack_pack(int argc, VALUE* argv);
 
-#endif
+static inline VALUE delegete_to_pack(int argc, VALUE* argv, VALUE self)
+{
+    if(argc == 0) {
+        return MessagePack_pack(1, &self);
+    } else if(argc == 1) {
+        /* write to io */
+        VALUE argv2[2];
+        argv2[0] = self;
+        argv2[1] = argv[0];
+        return MessagePack_pack(2, argv2);
+    } else {
+        rb_raise(rb_eArgError, "wrong number of arguments (%d for 0..1)", argc);
+    }
+}
+
+#define ENSURE_PACKER(argc, argv, packer, pk) \
+    if(argc != 1 || rb_class_of(argv[0]) != cMessagePack_Packer) { \
+        return delegete_to_pack(argc, argv, self); \
+    } \
+    VALUE packer = argv[0]; \
+    msgpack_packer_t *pk; \
+    Data_Get_Struct(packer, msgpack_packer_t, pk);
 
+
+#endif

ext/msgpack/rbinit.c

@@ -17,6 +17,7 @@
  */
 
 #include "buffer_class.h"
+#include "extended_class.h"
 #include "packer_class.h"
 #include "unpacker_class.h"
 #include "core_ext.h"
@@ -26,6 +27,7 @@ void Init_msgpack(void)
     VALUE mMessagePack = rb_define_module("MessagePack");
 
     MessagePack_Buffer_module_init(mMessagePack);
+    MessagePack_Extended_module_init(mMessagePack);
     MessagePack_Packer_module_init(mMessagePack);
     MessagePack_Unpacker_module_init(mMessagePack);
     MessagePack_core_ext_module_init();