* st.c (st_update): new function to lookup the given key and

update the value.  [ruby-dev:44998]


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

ChangeLog

@@ -1,3 +1,8 @@
+Tue Dec 27 22:04:27 2011  Nobuyoshi Nakada  <nobu@ruby-lang.org>
+
+	* st.c (st_update): new function to lookup the given key and
+	  update the value.  [ruby-dev:44998]
+
 Tue Dec 27 21:17:33 2011  Nobuyoshi Nakada  <nobu@ruby-lang.org>
 
 	* node.h (rb_args_info): change pre_args_num and post_args_num as

ext/-test-/st/update/extconf.rb

@@ -0,0 +1 @@
+create_makefile("-test-/st/update")

ext/-test-/st/update/update.c

@@ -0,0 +1,34 @@
+#include <ruby.h>
+#include <ruby/st.h>
+
+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
+test_st_update(VALUE self, VALUE key)
+{
+    if (st_update(RHASH_TBL(self), (st_data_t)key, update_func, 0))
+	return Qtrue;
+    else
+	return Qfalse;
+}
+
+void
+Init_update(void)
+{
+    VALUE st = rb_define_class_under(rb_define_module("Bug"), "StTable", rb_cHash);
+    rb_define_method(st, "st_update", test_st_update, 1);
+}
+

include/ruby/st.h

@@ -113,6 +113,7 @@ 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 *);
 int st_get_key(st_table *, st_data_t, st_data_t *);
+int st_update(st_table *table, st_data_t key, int (*func)(st_data_t key, st_data_t *value, st_data_t arg), st_data_t arg);
 int st_foreach(st_table *, int (*)(ANYARGS), st_data_t);
 int st_reverse_foreach(st_table *, int (*)(ANYARGS), st_data_t);
 void st_add_direct(st_table *, st_data_t, st_data_t);

st.c

@@ -732,6 +732,62 @@ st_cleanup_safe(st_table *table, st_data_t never)
     }
 }
 
+int
+st_update(st_table *table, st_data_t key, int (*func)(st_data_t key, st_data_t *value, st_data_t arg), st_data_t arg)
+{
+    st_index_t hash_val, bin_pos;
+    register st_table_entry *ptr, **last, *tmp;
+    st_data_t value;
+
+    if (table->entries_packed) {
+	st_index_t i;
+	for (i = 0; i < table->num_entries; i++) {
+	    if ((st_data_t)table->bins[i*2] == key) {
+		value = (st_data_t)table->bins[i*2+1];
+		switch ((*func)(key, &value, arg)) {
+		  case ST_CONTINUE:
+		    table->bins[i*2+1] = (struct st_table_entry*)value;
+		    break;
+		  case ST_DELETE:
+		    table->num_entries--;
+		    memmove(&table->bins[i*2], &table->bins[(i+1)*2],
+			    sizeof(struct st_table_entry*) * 2 * (table->num_entries-i));
+		}
+		return 1;
+	    }
+	}
+	return 0;
+    }
+
+    hash_val = do_hash(key, table);
+    FIND_ENTRY(table, ptr, hash_val, bin_pos);
+
+    if (ptr == 0) {
+	return 0;
+    }
+    else {
+	value = ptr->record;
+	switch ((*func)(ptr->key, &value, arg)) {
+	  case ST_CONTINUE:
+	    ptr->record = value;
+	    break;
+	  case ST_DELETE:
+	    last = &table->bins[bin_pos];
+	    for (; (tmp = *last) != 0; last = &tmp->next) {
+		if (ptr == tmp) {
+		    tmp = ptr->fore;
+		    *last = ptr->next;
+		    REMOVE_ENTRY(table, ptr);
+		    free(ptr);
+		    break;
+		}
+	    }
+	    break;
+	}
+	return 1;
+    }
+}
+
 int
 st_foreach(st_table *table, int (*func)(ANYARGS), st_data_t arg)
 {

test/-ext-/st/test_update.rb

@@ -0,0 +1,39 @@
+require 'test/unit'
+require "-test-/st/update"
+
+class Bug::StTable
+  class Test_Update < Test::Unit::TestCase
+    def setup
+      @tbl = Bug::StTable.new
+      @tbl[:a] = 1
+      @tbl[:b] = 2
+    end
+
+    def test_notfound
+      called = false
+      assert_equal(false, @tbl.st_update(:c) {called = true})
+      assert_equal(false, called)
+    end
+
+    def test_continue
+      args = nil
+      assert_equal(true, @tbl.st_update(:a) {|*x| args = x; false})
+      assert_equal({a: 1, b: 2}, @tbl, :a)
+      assert_equal([:a, 1], args)
+    end
+
+    def test_delete
+      args = nil
+      assert_equal(true, @tbl.st_update(:a) {|*x| args = x; nil})
+      assert_equal({b: 2}, @tbl, :a)
+      assert_equal([:a, 1], args)
+    end
+
+    def test_update
+      args = nil
+      assert_equal(true, @tbl.st_update(:a) {|*x| args = x; 3})
+      assert_equal({a: 3, b: 2}, @tbl, :a)
+      assert_equal([:a, 1], args)
+    end
+  end
+end