public
Last active

Fix Ruby Hash#shift

  • Download Gist
00-test_hash_shift.rb
Ruby
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25
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
1 2 3 4 5 6 7
.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
1 2 3 4 5 6 7
.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
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98
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
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38
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);
}
}

Please sign in to comment on this gist.

Something went wrong with that request. Please try again.