Fix Ruby Hash#shift

00-test_hash_shift.rb

require 'benchmark'

N = 50000
K = 10000

a = {}
K.times{|i| a[i] = nil}

class Hash
  def custom_shift
    ar = first
    delete ar[0]
    ar
  end
end

Benchmark.bm do |x|
  x.report("Hash#shift") do
    N.times{ k,v = a.shift; a[k] = v }
  end

  x.report("Hash#custom_shift") do
    N.times{ k,v = a.custom_shift; a[k] = v }
  end
end

01-before_fix

.btr~$ ./ruby -I .ext/x86_64-linux -I . -I ../lib -v
ruby 2.1.0dev (2013-04-23 trunk 40423) [x86_64-linux]

.btr~$ ./ruby -I .ext/x86_64-linux -I . -I ../lib ../test_hash_shift.rb 
      user     system      total        real
Hash#shift  2.810000   0.000000   2.810000 (  2.812111)
Hash#custom_shift  0.030000   0.000000   0.030000 (  0.034924)

02-after_fix

.btr~$ ./ruby -I .ext/x86_64-linux -I . -I ../lib -v
ruby 2.1.0dev (2013-04-23 trunk 40423) [x86_64-linux]

.btr~$ ./ruby -I .ext/x86_64-linux -I . -I ../lib ../test_hash_shift.rb 
       user     system      total        real
Hash#shift  0.020000   0.000000   0.020000 (  0.017113)
Hash#custom_shift  0.030000   0.000000   0.030000 (  0.034295)

03-fix_hash_shift.diff

diff --git a/hash.c b/hash.c
index a06e7a3..43b6757 100644
--- a/hash.c
+++ b/hash.c
@@ -875,17 +875,6 @@ struct shift_var {
 };
 
 static int
-shift_i(VALUE key, VALUE value, VALUE arg)
-{
-    struct shift_var *var = (struct shift_var *)arg;
-
-    if (var->key != Qundef) return ST_STOP;
-    var->key = key;
-    var->val = value;
-    return ST_DELETE;
-}
-
-static int
 shift_i_safe(VALUE key, VALUE value, VALUE arg)
 {
     struct shift_var *var = (struct shift_var *)arg;
@@ -916,14 +905,17 @@ rb_hash_shift(VALUE hash)
     rb_hash_modify_check(hash);
     if (RHASH(hash)->ntbl) {
 	var.key = Qundef;
-	rb_hash_foreach(hash, RHASH_ITER_LEV(hash) > 0 ? shift_i_safe : shift_i,
-			(VALUE)&var);
-
-	if (var.key != Qundef) {
-	    if (RHASH_ITER_LEV(hash) > 0) {
+	if (RHASH_ITER_LEV(hash) == 0) {
+	    if (st_shift(RHASH(hash)->ntbl, &var.key, &var.val)) {
+		return rb_assoc_new(var.key, var.val);
+	    }
+	}
+	else {
+	    rb_hash_foreach(hash, shift_i_safe, (VALUE)&var);
+	    if (var.key != Qundef) {
 		rb_hash_delete_key(hash, var.key);
+		return rb_assoc_new(var.key, var.val);
 	    }
-	    return rb_assoc_new(var.key, var.val);
 	}
     }
     return hash_default_value(hash, Qnil);
diff --git a/include/ruby/st.h b/include/ruby/st.h
index ef838b7..88135f6 100644
--- a/include/ruby/st.h
+++ b/include/ruby/st.h
@@ -102,6 +102,7 @@ st_table *st_init_strcasetable(void);
 st_table *st_init_strcasetable_with_size(st_index_t);
 int st_delete(st_table *, st_data_t *, st_data_t *); /* returns 0:notfound 1:deleted */
 int st_delete_safe(st_table *, st_data_t *, st_data_t *, st_data_t);
+int st_shift(st_table *, st_data_t *, st_data_t *); /* returns 0:notfound 1:deleted */
 int st_insert(st_table *, st_data_t, st_data_t);
 int st_insert2(st_table *, st_data_t, st_data_t, st_data_t (*)(st_data_t));
 int st_lookup(st_table *, st_data_t, st_data_t *);
diff --git a/st.c b/st.c
index 9151267..5b59538 100644
--- a/st.c
+++ b/st.c
@@ -793,6 +793,35 @@ st_delete_safe(register st_table *table, register st_data_t *key, st_data_t *val
     return 0;
 }
 
+int
+st_shift(register st_table *table, register st_data_t *key, st_data_t *value)
+{
+    st_index_t hash_val;
+    st_table_entry **prev;
+    register st_table_entry *ptr;
+
+    if (table->num_entries == 0) {
+        if (value != 0) *value = 0;
+        return 0;
+    }
+
+    if (table->entries_packed) {
+        if (value != 0) *value = PVAL(table, 0);
+        *key = PKEY(table, 0);
+        remove_packed_entry(table, 0);
+        return 1;
+    }
+
+    prev = &table->bins[table->head->hash % table->num_bins];
+    while ((ptr = *prev) != table->head) prev = &ptr->next;
+    *prev = ptr->next;
+    if (value != 0) *value = ptr->record;
+    *key = ptr->key;
+    remove_entry(table, ptr);
+    st_free_entry(ptr);
+    return 1;
+}
+
 void
 st_cleanup_safe(st_table *table, st_data_t never)
 {

04-simpler_fix_hash_shift.diff

diff --git a/hash.c b/hash.c
index a06e7a3..36ac232 100644
--- a/hash.c
+++ b/hash.c
@@ -875,17 +875,6 @@ struct shift_var {
 };
 
 static int
-shift_i(VALUE key, VALUE value, VALUE arg)
-{
-    struct shift_var *var = (struct shift_var *)arg;
-
-    if (var->key != Qundef) return ST_STOP;
-    var->key = key;
-    var->val = value;
-    return ST_DELETE;
-}
-
-static int
 shift_i_safe(VALUE key, VALUE value, VALUE arg)
 {
     struct shift_var *var = (struct shift_var *)arg;
@@ -916,13 +905,10 @@ rb_hash_shift(VALUE hash)
     rb_hash_modify_check(hash);
     if (RHASH(hash)->ntbl) {
 	var.key = Qundef;
-	rb_hash_foreach(hash, RHASH_ITER_LEV(hash) > 0 ? shift_i_safe : shift_i,
-			(VALUE)&var);
+	rb_hash_foreach(hash, shift_i_safe, (VALUE)&var);
 
 	if (var.key != Qundef) {
-	    if (RHASH_ITER_LEV(hash) > 0) {
-		rb_hash_delete_key(hash, var.key);
-	    }
+	    rb_hash_delete_key(hash, var.key);
 	    return rb_assoc_new(var.key, var.val);
 	}
     }