* benchmark/bm_hash_shift.rb: add benchmark for Hash#shift

* hash.c (rb_hash_shift): use st_shift if hash is not being iterated to
  delete element without iterating the whole hash.

* hash.c (shift_i): remove function

* include/ruby/st.h (st_shift): add st_shift function

* st.c (st_shift): ditto

[Bug #8312] [ruby-core:54524] Patch by funny-falcon

git-svn-id: svn+ssh://ci.ruby-lang.org/ruby/trunk@40457 b2dd03c8-39d4-4d8f-98ff-823fe69b080e

5 files changed, 64 insertions, 17 deletions

diff --git a/ChangeLog b/ChangeLog
index e83ad892bb..9c995af4eb 100644
--- a/ChangeLog
+++ b/ChangeLog
@@ -1,3 +1,18 @@
+Wed Apr 25 14:26:00 2013  Charlie Somerville  <charlie@charliesomerville.com>
+
+	* benchmark/bm_hash_shift.rb: add benchmark for Hash#shift
+
+	* hash.c (rb_hash_shift): use st_shift if hash is not being iterated to
+	  delete element without iterating the whole hash.
+
+	* hash.c (shift_i): remove function
+
+	* include/ruby/st.h (st_shift): add st_shift function
+
+	* st.c (st_shift): ditto
+
+	[Bug #8312] [ruby-core:54524] Patch by funny-falcon
+
 Thu Apr 25 12:03:38 2013  Tanaka Akira  <akr@fsij.org>
 
 	* ext/socket/extconf.rb: Extract C programs as toplevel constants.
diff --git a/benchmark/bm_hash_shift.rb b/benchmark/bm_hash_shift.rb
new file mode 100644
index 0000000000..a645671a5b
--- /dev/null
+++ b/benchmark/bm_hash_shift.rb
@@ -0,0 +1,10 @@
+h = {}
+
+10000.times do |i|
+  h[i] = nil
+end
+
+50000.times do
+  k, v = h.shift
+  h[k] = v
+end
diff --git a/hash.c b/hash.c
index a06e7a352b..43b6757f14 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 ef838b7aeb..88135f6a38 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 9151267263..5b59538966 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)
 {