Let execution context local storage be an ID table

6 files changed, 42 insertions, 22 deletions

diff --git a/benchmark/fiber_locals.yml b/benchmark/fiber_locals.yml
new file mode 100644
index 0000000000..8588686477
--- /dev/null
+++ b/benchmark/fiber_locals.yml
@@ -0,0 +1,8 @@
+prelude: |
+  th = Thread.current
+  th[:key] = :val
+benchmark:
+  key?: th.key?(:key)
+  []: th[:key]
+  keys: th.keys
+loop_count: 1_000_000
diff --git a/common.mk b/common.mk
index 6341b430ce..22e7d99237 100644
--- a/common.mk
+++ b/common.mk
@@ -1919,6 +1919,7 @@ cont.$(OBJEXT): {$(VPATH)}defines.h
 cont.$(OBJEXT): {$(VPATH)}eval_intern.h
 cont.$(OBJEXT): {$(VPATH)}gc.h
 cont.$(OBJEXT): {$(VPATH)}id.h
+cont.$(OBJEXT): {$(VPATH)}id_table.h
 cont.$(OBJEXT): {$(VPATH)}intern.h
 cont.$(OBJEXT): {$(VPATH)}internal.h
 cont.$(OBJEXT): {$(VPATH)}method.h
diff --git a/cont.c b/cont.c
index d4a2bf9353..df197a7363 100644
--- a/cont.c
+++ b/cont.c
@@ -27,6 +27,7 @@
 #include "internal/warnings.h"
 #include "mjit.h"
 #include "vm_core.h"
+#include "id_table.h"
 
 static const int DEBUG = 0;
 
@@ -1018,7 +1019,7 @@ fiber_free(void *ptr)
     //if (DEBUG) fprintf(stderr, "fiber_free: %p[%p]\n", fiber, fiber->stack.base);
 
     if (fiber->cont.saved_ec.local_storage) {
-        st_free_table(fiber->cont.saved_ec.local_storage);
+        rb_id_table_free(fiber->cont.saved_ec.local_storage);
     }
 
     cont_free(&fiber->cont);
@@ -1037,7 +1038,7 @@ fiber_memsize(const void *ptr)
      * vm.c::thread_memsize already counts th->ec->local_storage
      */
     if (saved_ec->local_storage && fiber != th->root_fiber) {
-        size += st_memsize(saved_ec->local_storage);
+        size += rb_id_table_memsize(saved_ec->local_storage);
     }
     size += cont_memsize(&fiber->cont);
     return size;
diff --git a/thread.c b/thread.c
index 4faf04c1ce..5b23eaa92e 100644
--- a/thread.c
+++ b/thread.c
@@ -3219,11 +3219,11 @@ threadptr_local_aref(rb_thread_t *th, ID id)
 	return th->ec->local_storage_recursive_hash;
     }
     else {
-	st_data_t val;
-	st_table *local_storage = th->ec->local_storage;
+	VALUE val;
+	struct rb_id_table *local_storage = th->ec->local_storage;
 
-	if (local_storage != NULL && st_lookup(local_storage, id, &val)) {
-	    return (VALUE)val;
+	if (local_storage != NULL && rb_id_table_lookup(local_storage, id, &val)) {
+	    return val;
 	}
 	else {
 	    return Qnil;
@@ -3340,7 +3340,7 @@ rb_thread_fetch(int argc, VALUE *argv, VALUE self)
 	return target_th->ec->local_storage_recursive_hash;
     }
     else if (id && target_th->ec->local_storage &&
-	     st_lookup(target_th->ec->local_storage, id, &val)) {
+	     rb_id_table_lookup(target_th->ec->local_storage, id, &val)) {
 	return val;
     }
     else if (block_given) {
@@ -3362,18 +3362,18 @@ threadptr_local_aset(rb_thread_t *th, ID id, VALUE val)
 	return val;
     }
     else {
-	st_table *local_storage = th->ec->local_storage;
+	struct rb_id_table *local_storage = th->ec->local_storage;
 
 	if (NIL_P(val)) {
 	    if (!local_storage) return Qnil;
-	    st_delete_wrap(local_storage, id);
+	    rb_id_table_delete(local_storage, id);
 	    return Qnil;
 	}
 	else {
 	    if (local_storage == NULL) {
-		th->ec->local_storage = local_storage = st_init_numtable();
+		th->ec->local_storage = local_storage = rb_id_table_create(0);
 	    }
-	    st_insert(local_storage, id, val);
+	    rb_id_table_insert(local_storage, id, val);
 	    return val;
 	}
     }
@@ -3486,13 +3486,14 @@ rb_thread_variable_set(VALUE thread, VALUE id, VALUE val)
 static VALUE
 rb_thread_key_p(VALUE self, VALUE key)
 {
+    VALUE val;
     ID id = rb_check_id(&key);
-    st_table *local_storage = rb_thread_ptr(self)->ec->local_storage;
+    struct rb_id_table *local_storage = rb_thread_ptr(self)->ec->local_storage;
 
     if (!id || local_storage == NULL) {
 	return Qfalse;
     }
-    else if (st_is_member(local_storage, id)) {
+    else if (rb_id_table_lookup(local_storage, id, &val)) {
 	return Qtrue;
     }
     else {
@@ -3500,11 +3501,11 @@ rb_thread_key_p(VALUE self, VALUE key)
     }
 }
 
-static int
-thread_keys_i(ID key, VALUE value, VALUE ary)
+static enum rb_id_table_iterator_result
+thread_keys_i(ID key, VALUE value, void *ary)
 {
-    rb_ary_push(ary, ID2SYM(key));
-    return ST_CONTINUE;
+    rb_ary_push((VALUE)ary, ID2SYM(key));
+    return ID_TABLE_CONTINUE;
 }
 
 int
@@ -3530,11 +3531,11 @@ rb_thread_alone(void)
 static VALUE
 rb_thread_keys(VALUE self)
 {
-    st_table *local_storage = rb_thread_ptr(self)->ec->local_storage;
+    struct rb_id_table *local_storage = rb_thread_ptr(self)->ec->local_storage;
     VALUE ary = rb_ary_new();
 
     if (local_storage) {
-	st_foreach(local_storage, thread_keys_i, ary);
+	rb_id_table_foreach(local_storage, thread_keys_i, (void *)ary);
     }
     return ary;
 }
diff --git a/vm.c b/vm.c
index 5dd78a5e83..e772d3b112 100644
--- a/vm.c
+++ b/vm.c
@@ -2501,6 +2501,13 @@ rb_execution_context_update(const rb_execution_context_t *ec)
     }
 }
 
+static enum rb_id_table_iterator_result
+mark_local_storage_i(VALUE local, void *data)
+{
+    rb_gc_mark(local);
+    return ID_TABLE_CONTINUE;
+}
+
 void
 rb_execution_context_mark(const rb_execution_context_t *ec)
 {
@@ -2544,7 +2551,9 @@ rb_execution_context_mark(const rb_execution_context_t *ec)
 
     RUBY_MARK_UNLESS_NULL(ec->errinfo);
     RUBY_MARK_UNLESS_NULL(ec->root_svar);
-    rb_mark_tbl(ec->local_storage);
+    if (ec->local_storage) {
+        rb_id_table_foreach_values(ec->local_storage, mark_local_storage_i, NULL);
+    }
     RUBY_MARK_UNLESS_NULL(ec->local_storage_recursive_hash);
     RUBY_MARK_UNLESS_NULL(ec->local_storage_recursive_hash_for_trace);
     RUBY_MARK_UNLESS_NULL(ec->private_const_reference);
@@ -2639,7 +2648,7 @@ thread_memsize(const void *ptr)
 	size += th->ec->vm_stack_size * sizeof(VALUE);
     }
     if (th->ec->local_storage) {
-	size += st_memsize(th->ec->local_storage);
+	size += rb_id_table_memsize(th->ec->local_storage);
     }
     return size;
 }
diff --git a/vm_core.h b/vm_core.h
index 70575fa66b..a2f6af3c30 100644
--- a/vm_core.h
+++ b/vm_core.h
@@ -865,7 +865,7 @@ typedef struct rb_execution_context_struct {
     struct rb_thread_struct *thread_ptr;
 
     /* storage (ec (fiber) local) */
-    st_table *local_storage;
+    struct rb_id_table *local_storage;
     VALUE local_storage_recursive_hash;
     VALUE local_storage_recursive_hash_for_trace;