* st.c (st_update): table can be unpacked in the callback.

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

4 files changed, 48 insertions, 3 deletions

diff --git a/ChangeLog b/ChangeLog
index 05c900f970..1830e804b3 100644
--- a/ChangeLog
+++ b/ChangeLog
@@ -1,4 +1,6 @@
-Tue Feb  7 14:29:16 2012  Nobuyoshi Nakada  <nobu@ruby-lang.org>
+Tue Feb  7 14:52:10 2012  Nobuyoshi Nakada  <nobu@ruby-lang.org>
+
+	* st.c (st_update): table can be unpacked in the callback.
 
 	* st.c (st_foreach): should not yield same pair when checking
 	  after unpacking.
diff --git a/ext/-test-/st/numhash/numhash.c b/ext/-test-/st/numhash/numhash.c
index e186cd43d9..adf3d258b9 100644
--- a/ext/-test-/st/numhash/numhash.c
+++ b/ext/-test-/st/numhash/numhash.c
@@ -57,6 +57,30 @@ numhash_each(VALUE self)
     return st_foreach((st_table *)DATA_PTR(self), numhash_i, self) ? Qtrue : Qfalse;
 }
 
+static int
+update_func(st_data_t key, st_data_t *value, st_data_t arg)
+{
+    VALUE ret = rb_yield_values(2, (VALUE)key, (VALUE)*value);
+    switch (ret) {
+      case Qfalse:
+	return ST_STOP;
+      case Qnil:
+	return ST_DELETE;
+      default:
+	*value = ret;
+	return ST_CONTINUE;
+    }
+}
+
+static VALUE
+numhash_update(VALUE self, VALUE key)
+{
+    if (st_update((st_table *)DATA_PTR(self), (st_data_t)key, update_func, 0))
+	return Qtrue;
+    else
+	return Qfalse;
+}
+
 void
 Init_numhash(void)
 {
@@ -66,4 +90,6 @@ Init_numhash(void)
     rb_define_method(st, "[]", numhash_aref, 1);
     rb_define_method(st, "[]=", numhash_aset, 2);
     rb_define_method(st, "each", numhash_each, 0);
+    rb_define_method(st, "update", numhash_update, 1);
 }
+
diff --git a/st.c b/st.c
index 215e27b298..c45338abfd 100644
--- a/st.c
+++ b/st.c
@@ -761,12 +761,21 @@ st_update(st_table *table, st_data_t key, int (*func)(st_data_t key, st_data_t *
     st_index_t hash_val, bin_pos;
     register st_table_entry *ptr, **last, *tmp;
     st_data_t value;
+    int retval;
 
     if (table->entries_packed) {
 	st_index_t i = find_packed_index(table, key);
 	if (i < table->num_entries) {
 	    value = PVAL(table, i);
-	    switch ((*func)(key, &value, arg)) {
+	    retval = (*func)(key, &value, arg);
+	    if (!table->entries_packed) {
+		hash_val = do_hash(key, table);
+		bin_pos = hash_val % table->num_bins;
+		ptr = find_entry(table, key, hash_val, bin_pos);
+		if (ptr == 0) return 0;
+		goto unpacked;
+	    }
+	    switch (retval) {
 	      case ST_CONTINUE:
 		PVAL_SET(table, i, value);
 		break;
@@ -787,7 +796,9 @@ st_update(st_table *table, st_data_t key, int (*func)(st_data_t key, st_data_t *
     }
     else {
 	value = ptr->record;
-	switch ((*func)(ptr->key, &value, arg)) {
+	retval = (*func)(ptr->key, &value, arg);
+      unpacked:
+	switch (retval) {
 	  case ST_CONTINUE:
 	    ptr->record = value;
 	    break;
diff --git a/test/-ext-/st/test_numhash.rb b/test/-ext-/st/test_numhash.rb
index 9ec072e1a5..c092049bf9 100644
--- a/test/-ext-/st/test_numhash.rb
+++ b/test/-ext-/st/test_numhash.rb
@@ -17,5 +17,11 @@ class Bug::StNumHash
       end
       assert_equal([*0..5], keys)
     end
+
+    def test_update
+      assert_equal(true, @tbl.update(0) {@tbl[5] = :x})
+      assert_equal(:x, @tbl[0])
+      assert_equal(:x, @tbl[5])
+    end
   end
 end