msgpack-ruby-frozen-str-cache.diff

diff --git a/ext/msgpack/buffer.h b/ext/msgpack/buffer.h
index 9387254..5c8af67 100644
--- a/ext/msgpack/buffer.h
+++ b/ext/msgpack/buffer.h
@@ -20,6 +20,7 @@
 
 #include "compat.h"
 #include "sysdep.h"
+#include "frozen_str_cache.h"
 
 #ifndef MSGPACK_BUFFER_STRING_WRITE_REFERENCE_DEFAULT
 #define MSGPACK_BUFFER_STRING_WRITE_REFERENCE_DEFAULT (512*1024)
@@ -428,7 +429,7 @@ static inline VALUE msgpack_buffer_read_top_as_string(msgpack_buffer_t* b, size_
 {
 #ifndef DISABLE_BUFFER_READ_REFERENCE_OPTIMIZE
     /* optimize */
-    if(!will_be_frozen &&
+    if(!will_be_frozen &&  /* skips if will_be_frozen because freezing strings causes copying and spoils CoW */
             b->head->mapped_string != NO_MAPPED_STRING &&
             length >= b->read_reference_threshold) {
         VALUE result = _msgpack_buffer_refer_head_mapped_string(b, length);
@@ -442,6 +443,16 @@ static inline VALUE msgpack_buffer_read_top_as_string(msgpack_buffer_t* b, size_
     return result;
 }
 
+#ifndef DISABLE_MAP_KEY_CACHE
+static inline VALUE msgpack_buffer_read_top_with_fstring_cache(msgpack_buffer_t* b,
+        msgpack_frozen_str_cache_t* fstring_cache, size_t length)
+{
+    VALUE result = msgpack_frozen_str_cache_get(fstring_cache, b->read_buffer, length);
+    _msgpack_buffer_consumed(b, length);
+    return result;
+}
+#endif
+
 
 #endif
 
diff --git a/ext/msgpack/extconf.rb b/ext/msgpack/extconf.rb
index 4ca8cda..e2f9795 100644
--- a/ext/msgpack/extconf.rb
+++ b/ext/msgpack/extconf.rb
@@ -17,6 +17,7 @@ end
 #$CFLAGS << %[ -DDISABLE_RMEM_REUSE_INTERNAL_FRAGMENT]
 #$CFLAGS << %[ -DDISABLE_BUFFER_READ_REFERENCE_OPTIMIZE]
 #$CFLAGS << %[ -DDISABLE_BUFFER_READ_TO_S_OPTIMIZE]
+#$CFLAGS << %[ -DDISABLE_MAP_KEY_CACHE]
 
 if defined?(RUBY_ENGINE) && RUBY_ENGINE == 'rbx'
   # msgpack-ruby doesn't modify data came from RSTRING_PTR(str)
diff --git a/ext/msgpack/frozen_str_cache.c b/ext/msgpack/frozen_str_cache.c
new file mode 100644
index 0000000..7ed3182
--- /dev/null
+++ b/ext/msgpack/frozen_str_cache.c
@@ -0,0 +1,154 @@
+#include "frozen_str_cache.h"
+#include "buffer.h"
+
+#ifndef DISABLE_MAP_KEY_CACHE
+
+void msgpack_frozen_str_cache_init(msgpack_frozen_str_cache_t* cache)
+{
+    memset(cache, 0, sizeof(msgpack_frozen_str_cache_t));
+    for (int i = 0; i < 256; i++) {
+        rb_gc_register_address(&cache->values[i]);
+    }
+}
+
+void msgpack_frozen_str_cache_destroy(msgpack_frozen_str_cache_t* cache)
+{
+    for (int i = 0; i < 256; i++) {
+        rb_gc_unregister_address(&cache->values[i]);
+    }
+}
+
+void msgpack_frozen_str_cache_mark(msgpack_frozen_str_cache_t* cache)
+{
+    for (int i = 0; i < 256; i++) {
+        rb_gc_mark(cache->values[i]);
+    }
+}
+
+#define HASH16_1(hash16) ((hash16) & 0xff)
+#define HASH16_2(hash16) ((hash16 >> 8))
+
+/**
+ * If cache hits, set used flag for the index of the value.
+ */
+#define SET_RECENTLY_USED(cache, idx)   (((cache)->recently_used_bitvec)[(idx) >> 3] |= (1 << ((idx) & 0x1f)))
+#define GET_RECENTLY_USED(cache, idx)   (((cache)->recently_used_bitvec)[(idx) >> 3] &  (1 << ((idx) & 0x1f)))
+#define RESET_RECENTLY_USED(cache, idx) (((cache)->recently_used_bitvec)[(idx) >> 3] & ~(1 << ((idx) & 0x1f)))
+
+static inline unsigned int hash_fnv_16(const char* buffer, int length)
+{
+    // TODO review this hash function
+    uint32_t h32 = 0x811c9dc5;
+    while (length >= 4) {
+        uint32_t u32;
+        memcpy(&u32, buffer, 4);
+        h32 ^= u32;
+        h32 *= 16777619;
+        length -= 4;
+        buffer += 4;
+    }
+    if (length & 0x1) {  /* length is 1 or 3 */
+        h32 ^= buffer[0];
+        h32 ^= (19 ^ buffer[0]) * 47;
+        length--;
+        buffer++;
+    }
+    if (length) {  /* length is 2 */
+        uint32_t u16;
+        memcpy(&u16, buffer, 4);
+        h32 ^= (u16 & 0xff);
+        h32 *= 16777619;
+        h32 ^= (u16 >> 8);
+        h32 *= 16777619;
+    }
+    return (((h32 & 0xffff) * 47) ^ (h32 >> 16)) & 0xffff;
+}
+
+static inline bool string_equals(VALUE v, const char* buffer, size_t length)
+{
+    return v != 0 &&
+        RSTRING_LEN(v) == length &&
+        memcmp(RSTRING_PTR(v), buffer, length) == 0;
+}
+
+static inline VALUE frozen_str_new(const char* buffer, size_t length)
+{
+    VALUE str = rb_str_new(buffer, length);
+#ifdef COMPAT_HAVE_ENCODING
+    ENCODING_SET(str, msgpack_rb_encindex_utf8);
+#endif
+    rb_obj_freeze(str);
+    return str;
+}
+
+VALUE msgpack_frozen_str_cache_get(msgpack_frozen_str_cache_t* cache,
+        const char* buffer, size_t length)
+{
+    if (length == 0) {
+        return frozen_str_new(buffer, length);   // TODO cache the empty string
+    }
+    unsigned int hash16 = hash_fnv_16(buffer, length);
+    //printf("hash16: %u\n", hash16);
+
+    /* Lookup table1 */
+    uint8_t idx1 = cache->table1[HASH16_1(hash16)];
+    VALUE value1 = cache->values[idx1];
+    if (string_equals(value1, buffer, length)) {
+        //printf("hit table1\n");
+        SET_RECENTLY_USED(cache, idx1);
+        return value1;
+    }
+
+    /* Lookup table2 */
+    uint8_t idx2 = cache->table2[HASH16_2(hash16)];
+    VALUE value2 = cache->values[idx2];
+    if (string_equals(value2, buffer, length)) {
+        //printf("hit table2\n");
+        SET_RECENTLY_USED(cache, idx2);
+        return value2;
+    }
+
+    /* Cache miss. Create a new string */
+    VALUE value = frozen_str_new(buffer, length);
+    //printf("miss table1=%s table2=%s\n", (value1 == 0 ? "free" : "used"), (value2 == 0 ? "free" : "used"));
+
+    /* Try to cache the new string using the clock position at most twice */
+    unsigned int clock_idx;
+    for (int i = 0; i < 2; i++) {
+        clock_idx = (++cache->clock) & 0xff;
+        if (GET_RECENTLY_USED(cache, clock_idx)) {
+            /* This clock_idx is recently used. Drop the flag and go next */
+            RESET_RECENTLY_USED(cache, clock_idx);
+        } else {
+            /* This clock_idx is not used recently. Try to replace it */
+            if (value1 == 0) {
+                goto replace_table1;  /* table1 is free */
+
+            } else if (value2 == 0) {
+                goto replace_table2;  /* table2 is free */
+
+            } else if (GET_RECENTLY_USED(cache, idx2)) {
+                goto replace_table2;  /* table2 is full but not recently used */
+
+            } else if (GET_RECENTLY_USED(cache, idx1)) {
+                goto replace_table1;  /* table1 is full but not recently used */
+            }
+        }
+    }
+
+    /* Tried to replace two clock positions but failed. Giveup. */
+    return value;
+
+replace_table1:
+    cache->values[clock_idx] = value;
+    cache->table1[HASH16_1(hash16)] = clock_idx;
+    return value;
+
+replace_table2:
+    cache->values[clock_idx] = value;
+    cache->table2[HASH16_2(hash16)] = clock_idx;
+    return value;
+}
+
+#endif
+
diff --git a/ext/msgpack/frozen_str_cache.h b/ext/msgpack/frozen_str_cache.h
new file mode 100644
index 0000000..0aebe38
--- /dev/null
+++ b/ext/msgpack/frozen_str_cache.h
@@ -0,0 +1,49 @@
+/*
+ * MessagePack for Ruby
+ *
+ * Copyright (C) 2008-2013 Sadayuki Furuhashi
+ *
+ *    Licensed under the Apache License, Version 2.0 (the "License");
+ *    you may not use this file except in compliance with the License.
+ *    You may obtain a copy of the License at
+ *
+ *        http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *    Unless required by applicable law or agreed to in writing, software
+ *    distributed under the License is distributed on an "AS IS" BASIS,
+ *    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *    See the License for the specific language governing permissions and
+ *    limitations under the License.
+ */
+#ifndef MSGPACK_RUBY_FROZEN_STR_H__
+#define MSGPACK_RUBY_FROZEN_STR_H__
+
+#include "compat.h"
+
+#ifndef DISABLE_MAP_KEY_CACHE
+
+struct msgpack_frozen_str_cache_t {
+    uint32_t recently_used_bitvec[8];
+    uint8_t table1[256];
+    uint8_t table2[256];
+    VALUE values[256];
+    unsigned int clock;
+};
+
+typedef struct msgpack_frozen_str_cache_t msgpack_frozen_str_cache_t;
+
+
+void msgpack_frozen_str_cache_init(msgpack_frozen_str_cache_t* cache);
+
+void msgpack_frozen_str_cache_destroy(msgpack_frozen_str_cache_t* cache);
+
+void msgpack_frozen_str_cache_mark(msgpack_frozen_str_cache_t* cache);
+
+VALUE msgpack_frozen_str_cache_get(msgpack_frozen_str_cache_t* cache,
+        const char* buffer, size_t length);
+
+#endif
+
+
+#endif
+
diff --git a/ext/msgpack/unpacker.c b/ext/msgpack/unpacker.c
index 8b1f699..44b2f71 100644
--- a/ext/msgpack/unpacker.c
+++ b/ext/msgpack/unpacker.c
@@ -18,6 +18,7 @@
 
 #include "unpacker.h"
 #include "rmem.h"
+#include "frozen_str_cache.h"
 #include "extension_value_class.h"
 
 #if !defined(DISABLE_RMEM) && !defined(DISABLE_UNPACKER_STACK_RMEM) && \
@@ -34,12 +35,20 @@ static ID s_call;
 static msgpack_rmem_t s_stack_rmem;
 #endif
 
+#ifndef DISABLE_MAP_KEY_CACHE
+static msgpack_frozen_str_cache_t s_frozen_str_cache;
+#endif
+
 void msgpack_unpacker_static_init()
 {
 #ifdef UNPACKER_STACK_RMEM
     msgpack_rmem_init(&s_stack_rmem);
 #endif
 
+#ifndef DISABLE_MAP_KEY_CACHE
+    msgpack_frozen_str_cache_init(&s_frozen_str_cache);
+#endif
+
     s_call = rb_intern("call");
 }
 
@@ -48,6 +57,10 @@ void msgpack_unpacker_static_destroy()
 #ifdef UNPACKER_STACK_RMEM
     msgpack_rmem_destroy(&s_stack_rmem);
 #endif
+
+#ifndef DISABLE_MAP_KEY_CACHE
+    msgpack_frozen_str_cache_destroy(&s_frozen_str_cache);
+#endif
 }
 
 #define HEAD_BYTE_REQUIRED 0xc1
@@ -290,8 +303,19 @@ static inline int read_raw_body_begin(msgpack_unpacker_t* uk, int raw_type)
     if(length <= msgpack_buffer_top_readable_size(UNPACKER_BUFFER_(uk))) {
         /* don't use zerocopy for hash keys but get a frozen string directly
          * because rb_hash_aset freezes keys and it causes copying */
-        bool will_freeze = is_reading_map_key(uk);
-        VALUE string = msgpack_buffer_read_top_as_string(UNPACKER_BUFFER_(uk), length, will_freeze);
+        bool is_map_key = is_reading_map_key(uk);
+#ifndef DISABLE_MAP_KEY_CACHE
+        /* optimize */
+        if(is_map_key && raw_type == RAW_TYPE_STRING) {
+            /* Deserializing a string key of a Hash. Using map key cache optimization.
+             * Bypassing object_complete_string and calling object_complete directly. */
+            VALUE string = msgpack_buffer_read_top_with_fstring_cache(UNPACKER_BUFFER_(uk), &s_frozen_str_cache, length);
+            int ret = object_complete(uk, string);
+            uk->reading_raw_remaining = 0;
+            return ret;
+        }
+#endif
+        VALUE string = msgpack_buffer_read_top_as_string(UNPACKER_BUFFER_(uk), length, is_map_key);
         int ret;
         if(raw_type == RAW_TYPE_STRING) {
             ret = object_complete_string(uk, string);
@@ -300,7 +324,7 @@ static inline int read_raw_body_begin(msgpack_unpacker_t* uk, int raw_type)
         } else {
             ret = object_complete_ext(uk, raw_type, string);
         }
-        if(will_freeze) {
+        if(is_map_key) {
             rb_obj_freeze(string);
         }
         uk->reading_raw_remaining = 0;