Add optimized_struct option for fast Struct serialization

gmalette · gmalette · commit 9312864ddf7a · 2025-12-10T15:47:17.000-05:00
Adds a C-level fast path for serializing and deserializing Ruby Struct
types, bypassing Ruby proc callbacks entirely.

Usage:
  factory.register_type(0x01, MyStruct, optimized_struct: true)

Benefits:
- Directly accesses struct fields via RSTRUCT_GET in C
- No Ruby method calls for field access
- No proc allocation or invocation overhead
- Supports both regular and keyword_init structs

The wire format is identical to the recursive proc-based approach,
ensuring compatibility with existing serialized data.

This provides significant performance improvements for applications
that serialize many Struct objects.
diff --git a/ext/msgpack/factory_class.c b/ext/msgpack/factory_class.c
@@ -250,6 +250,27 @@ static VALUE Factory_register_type_internal(VALUE self, VALUE rb_ext_type, VALUE
         if(RTEST(rb_hash_aref(options, ID2SYM(rb_intern("recursive"))))) {
             flags |= MSGPACK_EXT_RECURSIVE;
         }
+
+        /* optimized_struct: true enables C-level fast path for Struct types */
+        if (RTEST(rb_hash_aref(options, ID2SYM(rb_intern("optimized_struct"))))) {
+            /* Verify it's actually a Struct subclass */
+            if (!(rb_obj_is_kind_of(ext_module, rb_cClass) &&
+                  RTEST(rb_class_inherited_p(ext_module, rb_cStruct)))) {
+                rb_raise(rb_eArgError, "optimized_struct: true requires a Struct subclass");
+            }
+            /* Verify packer/unpacker procs weren't also provided */
+            if (RTEST(packer_proc) || RTEST(unpacker_proc)) {
+                rb_raise(rb_eArgError, "optimized_struct: true cannot be used with packer or unpacker options");
+            }
+            /* Verify recursive wasn't also provided - optimized_struct handles recursion automatically */
+            if (flags & MSGPACK_EXT_RECURSIVE) {
+                rb_raise(rb_eArgError, "optimized_struct: true cannot be used with recursive option");
+            }
+            flags |= MSGPACK_EXT_STRUCT_FAST_PATH;
+            /* Store the class itself - C code uses it directly */
+            packer_proc = ext_module;
+            unpacker_proc = ext_module;
+        }
     }
 
     msgpack_packer_ext_registry_put(self, &fc->pkrg, ext_module, ext_type, flags, packer_proc);
diff --git a/ext/msgpack/packer.c b/ext/msgpack/packer.c
@@ -91,6 +91,28 @@ void msgpack_packer_write_hash_value(msgpack_packer_t* pk, VALUE v)
     rb_hash_foreach(v, write_hash_foreach, (VALUE) pk);
 }
 
+/* Pack a Struct's fields directly using C API, bypassing Ruby callbacks */
+struct msgpack_packer_struct_args_t;
+typedef struct msgpack_packer_struct_args_t msgpack_packer_struct_args_t;
+struct msgpack_packer_struct_args_t {
+    msgpack_packer_t* pk;
+    VALUE v;
+};
+
+static VALUE msgpack_packer_write_struct_fields_protected(VALUE value)
+{
+    msgpack_packer_struct_args_t *args = (msgpack_packer_struct_args_t *)value;
+    msgpack_packer_t* pk = args->pk;
+    VALUE v = args->v;
+
+    long len = RSTRUCT_LEN(v);
+    for (int i = 0; i < len; i++) {
+        VALUE field = RSTRUCT_GET(v, i);
+        msgpack_packer_write_value(pk, field);
+    }
+    return Qnil;
+}
+
 struct msgpack_call_proc_args_t;
 typedef struct msgpack_call_proc_args_t msgpack_call_proc_args_t;
 struct msgpack_call_proc_args_t {
@@ -114,6 +136,33 @@ bool msgpack_packer_try_write_with_ext_type_lookup(msgpack_packer_t* pk, VALUE v
         return false;
     }
 
+    if(ext_flags & MSGPACK_EXT_STRUCT_FAST_PATH) {
+        /* Fast path for Struct: directly access fields in C, no Ruby callbacks */
+        VALUE held_buffer = MessagePack_Buffer_hold(&pk->buffer);
+
+        msgpack_buffer_t parent_buffer = pk->buffer;
+        msgpack_buffer_init(PACKER_BUFFER_(pk));
+
+        /* Write struct fields with exception handling */
+        int exception_occured = 0;
+        msgpack_packer_struct_args_t args = { pk, v };
+        rb_protect(msgpack_packer_write_struct_fields_protected, (VALUE)&args, &exception_occured);
+
+        if (exception_occured) {
+            msgpack_buffer_destroy(PACKER_BUFFER_(pk));
+            pk->buffer = parent_buffer;
+            rb_jump_tag(exception_occured); // re-raise the exception
+        } else {
+            VALUE payload = msgpack_buffer_all_as_string(PACKER_BUFFER_(pk));
+            msgpack_buffer_destroy(PACKER_BUFFER_(pk));
+            pk->buffer = parent_buffer;
+            msgpack_packer_write_ext(pk, ext_type, payload);
+        }
+
+        RB_GC_GUARD(held_buffer);
+        return true;
+    }
+
     if(ext_flags & MSGPACK_EXT_RECURSIVE) {
         VALUE held_buffer = MessagePack_Buffer_hold(&pk->buffer);
 
diff --git a/ext/msgpack/packer_ext_registry.h b/ext/msgpack/packer_ext_registry.h
@@ -22,6 +22,7 @@
 #include "ruby.h"
 
 #define MSGPACK_EXT_RECURSIVE 0b0001
+#define MSGPACK_EXT_STRUCT_FAST_PATH 0b0010
 
 struct msgpack_packer_ext_registry_t;
 typedef struct msgpack_packer_ext_registry_t msgpack_packer_ext_registry_t;
diff --git a/ext/msgpack/unpacker.c b/ext/msgpack/unpacker.c
@@ -231,6 +231,13 @@ static inline int object_complete_ext(msgpack_unpacker_t* uk, int ext_type, VALU
     VALUE proc = msgpack_unpacker_ext_registry_lookup(uk->ext_registry, ext_type, &ext_flags);
 
     if(proc != Qnil) {
+        /* Handle struct fast path for empty structs (0 fields) */
+        if (ext_flags & MSGPACK_EXT_STRUCT_FAST_PATH) {
+            VALUE struct_class = proc;
+            VALUE obj = rb_class_new_instance(0, NULL, struct_class);
+            return object_complete(uk, obj);
+        }
+
         VALUE obj;
         VALUE arg = (str == Qnil ? rb_str_buf_new(0) : str);
         int raised;
@@ -371,6 +378,70 @@ static inline int read_raw_body_begin(msgpack_unpacker_t* uk, int raw_type)
 
     if(!(raw_type == RAW_TYPE_STRING || raw_type == RAW_TYPE_BINARY)) {
         proc = msgpack_unpacker_ext_registry_lookup(uk->ext_registry, raw_type, &ext_flags);
+
+        if(proc != Qnil && ext_flags & MSGPACK_EXT_STRUCT_FAST_PATH) {
+            /* Fast path for Struct: proc is actually the Struct class
+             * Read fields directly and construct struct in C */
+            VALUE struct_class = proc;
+            uk->last_object = Qnil;
+            reset_head_byte(uk);
+            uk->reading_raw_remaining = 0;
+
+            /* Get struct members */
+            VALUE members = rb_struct_s_members(struct_class);
+            long num_fields = RARRAY_LEN(members);
+
+            /* Check if this is a keyword_init struct (Ruby 2.7+) */
+            VALUE keyword_init = Qfalse;
+            if (rb_respond_to(struct_class, rb_intern("keyword_init?"))) {
+                keyword_init = rb_funcall(struct_class, rb_intern("keyword_init?"), 0);
+            }
+
+            /* Push a recursive marker so nested reads don't prematurely return */
+            _msgpack_unpacker_stack_push(uk, STACK_TYPE_RECURSIVE, 1, Qnil);
+
+            VALUE obj;
+            if (num_fields == 0) {
+                /* Special case for empty structs */
+                obj = rb_class_new_instance(0, NULL, struct_class);
+            } else if (RTEST(keyword_init)) {
+                /* For keyword_init structs, build a hash with member names as keys */
+                VALUE kwargs = rb_hash_new();
+                for (long i = 0; i < num_fields; i++) {
+                    int ret = msgpack_unpacker_read(uk, 0);
+                    if (ret < 0) {
+                        msgpack_unpacker_stack_pop(uk);
+                        return ret;
+                    }
+                    VALUE key = rb_ary_entry(members, i);
+                    rb_hash_aset(kwargs, key, uk->last_object);
+                }
+                /* Call new with keyword arguments */
+                obj = rb_class_new_instance_kw(1, &kwargs, struct_class, RB_PASS_KEYWORDS);
+            } else {
+                /* For regular structs, use positional arguments
+                 * Use RB_ALLOCV to avoid stack overflow with large structs */
+                VALUE allocv_holder;
+                VALUE *values = RB_ALLOCV_N(VALUE, allocv_holder, num_fields);
+                for (int i = 0; i < num_fields; i++) {
+                    int ret = msgpack_unpacker_read(uk, 0);
+                    if (ret < 0) {
+                        msgpack_unpacker_stack_pop(uk);
+                        RB_ALLOCV_END(allocv_holder);
+                        return ret;
+                    }
+                    values[i] = uk->last_object;
+                }
+                obj = rb_class_new_instance((int)num_fields, values, struct_class);
+                RB_ALLOCV_END(allocv_holder);
+            }
+
+            RB_GC_GUARD(struct_class);
+            RB_GC_GUARD(members);
+            msgpack_unpacker_stack_pop(uk);
+            return object_complete(uk, obj);
+        }
+
         if(proc != Qnil && ext_flags & MSGPACK_EXT_RECURSIVE) {
             VALUE obj;
             uk->last_object = Qnil;
diff --git a/ext/msgpack/unpacker_ext_registry.h b/ext/msgpack/unpacker_ext_registry.h
@@ -22,6 +22,7 @@
 #include "ruby.h"
 
 #define MSGPACK_EXT_RECURSIVE 0b0001
+#define MSGPACK_EXT_STRUCT_FAST_PATH 0b0010
 
 struct msgpack_unpacker_ext_registry_t;
 typedef struct msgpack_unpacker_ext_registry_t msgpack_unpacker_ext_registry_t;
diff --git a/spec/optimized_struct_spec.rb b/spec/optimized_struct_spec.rb
