Use batch APIs to create arrays and hashes

Basically a port of ruby/json#678

Rather than to allocate the container and push elements one by one,
we accumulate them on a stack and then use the faster batch APIs
to directly create the final container.

Author: byroot

ext/msgpack/unpacker.c

@@ -28,7 +28,6 @@
 
 static int RAW_TYPE_STRING = 256;
 static int RAW_TYPE_BINARY = 257;
-static int16_t INITIAL_BUFFER_CAPACITY_MAX = SHRT_MAX;
 
 static msgpack_rmem_t s_stack_rmem;
 
@@ -39,9 +38,98 @@ static inline VALUE rb_hash_new_capa(long capa)
 }
 #endif
 
-static inline int16_t initial_buffer_size(long size)
+/* rvalue_stack functions */
+
+static inline void _msgpack_unpacker_rvalue_stack_free(msgpack_rvalue_stack_t* value_stack) {
+    switch (value_stack->type) {
+        case STACK_TYPE_UNALLOCATED:
+            break;
+        case STACK_TYPE_RMEM:
+            if (!msgpack_rmem_free(&s_stack_rmem, value_stack->data)) {
+                rb_bug("Failed to free an rmem pointer, memory leak?");
+            }
+            break;
+        case STACK_TYPE_HEAP:
+            xfree(value_stack->data);
+            break;
+    }
+    memset(value_stack, 0, sizeof(msgpack_rvalue_stack_t));
+}
+
+static void _msgpack_unpacker_rvalue_stack_grow(msgpack_rvalue_stack_t* value_stack) {
+    switch (value_stack->type) {
+        case STACK_TYPE_UNALLOCATED: {
+            value_stack->data = msgpack_rmem_alloc(&s_stack_rmem);
+            value_stack->capacity = MSGPACK_RMEM_PAGE_SIZE / sizeof(VALUE);
+            value_stack->type = STACK_TYPE_RMEM;
+            break;
+        }
+        case STACK_TYPE_RMEM: {
+            size_t new_capacity = value_stack->capacity * 2;
+            VALUE *new_ptr = ALLOC_N(VALUE, new_capacity);
+            MEMCPY(new_ptr, value_stack->data, VALUE, value_stack->depth);
+            if (!msgpack_rmem_free(&s_stack_rmem, value_stack->data)) {
+                rb_bug("Failed to free an rmem pointer, memory leak?");
+            }
+            value_stack->type = STACK_TYPE_HEAP;
+            value_stack->data = new_ptr;
+            value_stack->capacity = new_capacity;
+            break;
+        }
+        case STACK_TYPE_HEAP: {
+            size_t new_capacity = value_stack->capacity * 2;
+            REALLOC_N(value_stack->data, VALUE, new_capacity);
+            value_stack->capacity = new_capacity;
+            break;
+        }
+    }
+}
+
+static inline void _msgpack_unpacker_rvalue_stack_push(msgpack_unpacker_t* uk, VALUE value) {
+    size_t free_slots = uk->value_stack.capacity - uk->value_stack.depth;
+
+    if (RB_UNLIKELY(free_slots == 0)) {
+        _msgpack_unpacker_rvalue_stack_grow(&uk->value_stack);
+        free_slots = uk->value_stack.capacity - uk->value_stack.depth;
+    }
+
+    RB_OBJ_WRITE(uk->self, &uk->value_stack.data[uk->value_stack.depth++], value);
+}
+
+// static inline VALUE _msgpack_unpacker_rvalue_stack_top(msgpack_unpacker_t* uk)
+// {
+//     return uk->value_stack.data[uk->value_stack.depth - 1];
+// }
+
+static inline VALUE _msgpack_unpacker_rvalue_stack_create_array(msgpack_unpacker_t* uk, size_t length) {
+    VALUE *data = &uk->value_stack.data[uk->value_stack.depth - length];
+
+    VALUE array = rb_ary_new_from_values(length, data);
+
+    RB_OBJ_WRITE(uk->self, data, array);
+    uk->value_stack.depth -= (length - 1);
+
+    return RB_GC_GUARD(array);
+}
+
+static inline VALUE _msgpack_unpacker_rvalue_stack_create_hash(msgpack_unpacker_t* uk, size_t items_count) {
+    size_t length = items_count / 2;
+    VALUE *data = &uk->value_stack.data[uk->value_stack.depth - items_count];
+
+    VALUE hash = rb_hash_new_capa(length);
+    rb_hash_bulk_insert(items_count, data, hash);
+
+    RB_OBJ_WRITE(uk->self, data, hash);
+    uk->value_stack.depth -= (items_count - 1);
+
+    return RB_GC_GUARD(hash);
+}
+
+void msgpack_unpacker_mark_rvalue_stack(msgpack_rvalue_stack_t* value_stack)
 {
-    return (size > INITIAL_BUFFER_CAPACITY_MAX) ? INITIAL_BUFFER_CAPACITY_MAX : size;
+    if (value_stack->data) {
+        rb_gc_mark_locations(value_stack->data, value_stack->data + value_stack->depth);
+    }
 }
 
 void msgpack_unpacker_static_init(void)
@@ -88,27 +176,15 @@ static inline void _msgpack_unpacker_free_stack(msgpack_unpacker_stack_t* stack)
 void _msgpack_unpacker_destroy(msgpack_unpacker_t* uk)
 {
     _msgpack_unpacker_free_stack(uk->stack);
+    _msgpack_unpacker_rvalue_stack_free(&uk->value_stack);
     msgpack_buffer_destroy(UNPACKER_BUFFER_(uk));
 }
 
-void msgpack_unpacker_mark_stack(msgpack_unpacker_stack_t* stack)
-{
-    while (stack) {
-        msgpack_unpacker_stack_entry_t* s = stack->data;
-        msgpack_unpacker_stack_entry_t* send = stack->data + stack->depth;
-        for(; s < send; s++) {
-            rb_gc_mark(s->object);
-            rb_gc_mark(s->key);
-        }
-        stack = stack->parent;
-    }
-}
-
 void msgpack_unpacker_mark(msgpack_unpacker_t* uk)
 {
     rb_gc_mark(uk->last_object);
     rb_gc_mark(uk->reading_raw);
-    msgpack_unpacker_mark_stack(uk->stack);
+    msgpack_unpacker_mark_rvalue_stack(&uk->value_stack);
     /* See MessagePack_Buffer_wrap */
     /* msgpack_buffer_mark(UNPACKER_BUFFER_(uk)); */
     rb_gc_mark(uk->buffer_ref);
@@ -160,13 +236,15 @@ static inline int object_complete(msgpack_unpacker_t* uk, VALUE object)
     }
 
     uk->last_object = object;
+    _msgpack_unpacker_rvalue_stack_push(uk, object);
     reset_head_byte(uk);
     return PRIMITIVE_OBJECT_COMPLETE;
 }
 
 static inline int object_complete_symbol(msgpack_unpacker_t* uk, VALUE object)
 {
     uk->last_object = object;
+    _msgpack_unpacker_rvalue_stack_push(uk, object);
     reset_head_byte(uk);
     return PRIMITIVE_OBJECT_COMPLETE;
 }
@@ -204,7 +282,7 @@ static inline msgpack_unpacker_stack_entry_t* _msgpack_unpacker_stack_entry_top(
     return &uk->stack->data[uk->stack->depth-1];
 }
 
-static inline int _msgpack_unpacker_stack_push(msgpack_unpacker_t* uk, enum stack_type_t type, size_t count, VALUE object)
+static inline int _msgpack_unpacker_stack_push(msgpack_unpacker_t* uk, enum stack_type_t type, size_t count)
 {
     reset_head_byte(uk);
 
@@ -214,22 +292,20 @@ static inline int _msgpack_unpacker_stack_push(msgpack_unpacker_t* uk, enum stac
 
     msgpack_unpacker_stack_entry_t* next = &uk->stack->data[uk->stack->depth];
     next->count = count;
+    next->size = count;
     next->type = type;
-    next->object = object;
-    next->key = Qnil;
-
     uk->stack->depth++;
     return PRIMITIVE_CONTAINER_START;
 }
 
-static inline VALUE msgpack_unpacker_stack_pop(msgpack_unpacker_t* uk)
+static inline size_t msgpack_unpacker_stack_pop(msgpack_unpacker_t* uk)
 {
     return --uk->stack->depth;
 }
 
 static inline bool msgpack_unpacker_stack_is_empty(msgpack_unpacker_t* uk)
 {
-    return uk->stack->depth == 0;
+    return uk->stack->depth == 0 || _msgpack_unpacker_stack_entry_top(uk)->type == STACK_TYPE_RECURSIVE;
 }
 
 #ifdef USE_CASE_RANGE
@@ -259,9 +335,7 @@ static inline bool is_reading_map_key(msgpack_unpacker_t* uk)
 {
     if(uk->stack->depth > 0) {
         msgpack_unpacker_stack_entry_t* top = _msgpack_unpacker_stack_entry_top(uk);
-        if(top->type == STACK_TYPE_MAP_KEY) {
-            return true;
-        }
+        return top->type == STACK_TYPE_MAP && (top->count % 2 == 0);
     }
     return false;
 }
@@ -312,14 +386,9 @@ static inline int read_raw_body_begin(msgpack_unpacker_t* uk, int raw_type)
             reset_head_byte(uk);
             uk->reading_raw_remaining = 0;
 
-            msgpack_unpacker_stack_t* child_stack = _msgpack_unpacker_new_stack();
-            child_stack->parent = uk->stack;
-            uk->stack = child_stack;
-
+            _msgpack_unpacker_stack_push(uk, STACK_TYPE_RECURSIVE, 1);
             obj = rb_proc_call_with_block(proc, 1, &uk->self, Qnil);
-
-            uk->stack = child_stack->parent;
-            _msgpack_unpacker_free_stack(child_stack);
+            msgpack_unpacker_stack_pop(uk);
 
             return object_complete(uk, obj);
         }
@@ -382,14 +451,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(initial_buffer_size(count)));
+        return _msgpack_unpacker_stack_push(uk, STACK_TYPE_ARRAY, 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(initial_buffer_size(count)));
+        return _msgpack_unpacker_stack_push(uk, STACK_TYPE_MAP, count*2);
 
     SWITCH_RANGE(b, 0xc0, 0xdf)  // Variable
         switch(b) {
@@ -612,7 +681,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(initial_buffer_size(count)));
+                return _msgpack_unpacker_stack_push(uk, STACK_TYPE_ARRAY, count);
             }
 
         case 0xdd:  // array 32
@@ -622,7 +691,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(initial_buffer_size(count)));
+                return _msgpack_unpacker_stack_push(uk, STACK_TYPE_ARRAY, count);
             }
 
         case 0xde:  // map 16
@@ -632,7 +701,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(initial_buffer_size(count)));
+                return _msgpack_unpacker_stack_push(uk, STACK_TYPE_MAP, count*2);
             }
 
         case 0xdf:  // map 32
@@ -642,7 +711,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(initial_buffer_size(count)));
+                return _msgpack_unpacker_stack_push(uk, STACK_TYPE_MAP, count*2);
             }
 
         default:
@@ -730,28 +799,23 @@ int msgpack_unpacker_read(msgpack_unpacker_t* uk, size_t target_stack_depth)
         container_completed:
         {
             msgpack_unpacker_stack_entry_t* top = _msgpack_unpacker_stack_entry_top(uk);
-            switch(top->type) {
-            case STACK_TYPE_ARRAY:
-                rb_ary_push(top->object, uk->last_object);
-                break;
-            case STACK_TYPE_MAP_KEY:
-                top->key = uk->last_object;
-                top->type = STACK_TYPE_MAP_VALUE;
-                break;
-            case STACK_TYPE_MAP_VALUE:
-                if(uk->symbolize_keys && rb_type(top->key) == T_STRING) {
-                    /* here uses rb_str_intern instead of rb_intern so that Ruby VM can GC unused symbols */
-                    rb_hash_aset(top->object, rb_str_intern(top->key), uk->last_object);
-                } else {
-                    rb_hash_aset(top->object, top->key, uk->last_object);
-                }
-                top->type = STACK_TYPE_MAP_KEY;
-                break;
-            }
             size_t count = --top->count;
 
             if(count == 0) {
-                object_complete(uk, top->object);
+                switch(top->type) {
+                case STACK_TYPE_ARRAY:
+                    _msgpack_unpacker_rvalue_stack_create_array(uk, top->size);
+                    break;
+                case STACK_TYPE_MAP:
+                    _msgpack_unpacker_rvalue_stack_create_hash(uk, top->size);
+                    break;
+                case STACK_TYPE_RECURSIVE:
+                    object_complete(uk, _msgpack_unpacker_rvalue_stack_pop(uk));
+                    return PRIMITIVE_OBJECT_COMPLETE;
+                    break;
+                }
+
+                object_complete(uk, _msgpack_unpacker_rvalue_stack_pop(uk));
                 if(msgpack_unpacker_stack_pop(uk) <= target_stack_depth) {
                     return PRIMITIVE_OBJECT_COMPLETE;
                 }

ext/msgpack/unpacker.h

@@ -26,29 +26,42 @@
 struct msgpack_unpacker_t;
 typedef struct msgpack_unpacker_t msgpack_unpacker_t;
 typedef struct msgpack_unpacker_stack_t msgpack_unpacker_stack_t;
+typedef struct msgpack_rvalue_stack_t msgpack_rvalue_stack_t;
+
+enum rvalue_stack_type_t {
+    STACK_TYPE_UNALLOCATED = 0,
+    STACK_TYPE_RMEM = 1,
+    STACK_TYPE_HEAP = 2,
+};
+
+struct msgpack_rvalue_stack_t {
+    enum rvalue_stack_type_t type;
+    size_t depth;
+    size_t capacity;
+    VALUE *data;
+};
 
 enum stack_type_t {
     STACK_TYPE_ARRAY,
-    STACK_TYPE_MAP_KEY,
-    STACK_TYPE_MAP_VALUE,
+    STACK_TYPE_MAP,
+    STACK_TYPE_RECURSIVE,
 };
 
 typedef struct {
     size_t count;
+    size_t size;
     enum stack_type_t type;
-    VALUE object;
-    VALUE key;
 } msgpack_unpacker_stack_entry_t;
 
 struct msgpack_unpacker_stack_t {
     size_t depth;
     size_t capacity;
     msgpack_unpacker_stack_entry_t *data;
-    msgpack_unpacker_stack_t *parent;
 };
 
 struct msgpack_unpacker_t {
     msgpack_buffer_t buffer;
+    msgpack_rvalue_stack_t value_stack;
     msgpack_unpacker_stack_t *stack;
     unsigned int head_byte;
 
@@ -124,9 +137,13 @@ int msgpack_unpacker_read(msgpack_unpacker_t* uk, size_t target_stack_depth);
 
 int msgpack_unpacker_skip(msgpack_unpacker_t* uk, size_t target_stack_depth);
 
+static inline VALUE _msgpack_unpacker_rvalue_stack_pop(msgpack_unpacker_t* uk) {
+    return uk->value_stack.data[--uk->value_stack.depth];
+}
+
 static inline VALUE msgpack_unpacker_get_last_object(msgpack_unpacker_t* uk)
 {
-    return uk->last_object;
+    return _msgpack_unpacker_rvalue_stack_pop(uk);
 }