* array.c (rb_ary_uniq_bang, rb_ary_uniq): unique by the result of

given block.  [ruby-dev:37998]

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

ChangeLog

@@ -1,3 +1,8 @@
+Sun Feb 15 00:45:41 2009  Nobuyoshi Nakada  <nobu@ruby-lang.org>
+
+	* array.c (rb_ary_uniq_bang, rb_ary_uniq): unique by the result of
+	  given block.  [ruby-dev:37998]
+
 Sun Feb 15 00:39:44 2009  Nobuyoshi Nakada  <nobu@ruby-lang.org>
 
 	* array.c (ary_resize_capa): should not overwrite outside embedded

array.c

@@ -2887,15 +2887,43 @@ ary_add_hash(VALUE hash, VALUE ary)
     return hash;
 }
 
-static VALUE
+static inline VALUE
-ary_make_hash(VALUE ary)
+ary_tmp_hash_new(void)
 {
     VALUE hash = rb_hash_new();
 
     RBASIC(hash)->klass = 0;
+    return hash;
+}
+
+static VALUE
+ary_make_hash(VALUE ary)
+{
+    VALUE hash = ary_tmp_hash_new();
     return ary_add_hash(hash, ary);
 }
 
+static VALUE
+ary_add_hash_by(VALUE hash, VALUE ary)
+{
+    long i;
+
+    for (i = 0; i < RARRAY_LEN(ary); ++i) {
+	VALUE v = rb_ary_elt(ary, i), k = rb_yield(v);
+	if (rb_hash_lookup2(hash, k, Qundef) == Qundef) {
+	    rb_hash_aset(hash, k, v);
+	}
+    }
+    return hash;
+}
+
+static VALUE
+ary_make_hash_by(VALUE ary)
+{
+    VALUE hash = ary_tmp_hash_new();
+    return ary_add_hash_by(hash, ary);
+}
+
 static inline void
 ary_recycle_hash(VALUE hash)
 {
@@ -3010,6 +3038,13 @@ rb_ary_or(VALUE ary1, VALUE ary2)
     return ary3;
 }
 
+static int
+push_value(st_data_t key, st_data_t val, st_data_t ary)
+{
+    rb_ary_push((VALUE)ary, (VALUE)val);
+    return ST_CONTINUE;
+}
+
 /*
  *  call-seq:
  *     array.uniq! -> array or nil
@@ -3022,6 +3057,8 @@ rb_ary_or(VALUE ary1, VALUE ary2)
  *     a.uniq!   #=> ["a", "b", "c"]
  *     b = [ "a", "b", "c" ]
  *     b.uniq!   #=> nil
+ *     c = [ "a:def", "a:xyz", "b:abc", "b:xyz", "c:jkl" ]
+ *     c.uniq! {|s| s[/^\w+/]}  #=> [ "a:def", "b:abc", "c:jkl" ] 
  */
 
 static VALUE
@@ -3030,8 +3067,17 @@ rb_ary_uniq_bang(VALUE ary)
     VALUE hash, v;
     long i, j;
 
+    if (rb_block_given_p()) {
+	hash = ary_make_hash_by(ary);
+	if (RARRAY_LEN(ary) == (i = RHASH_SIZE(hash))) {
+	    return Qnil;
+	}
+	ary_resize_capa(ary, i);
+	ARY_SET_LEN(ary, 0);
+	st_foreach(RHASH_TBL(hash), push_value, ary);
+    }
+    else {
 	hash = ary_make_hash(ary);
-
 	if (RARRAY_LEN(ary) == RHASH_SIZE(hash)) {
 	    return Qnil;
 	}
@@ -3042,6 +3088,7 @@ rb_ary_uniq_bang(VALUE ary)
 	    }
 	}
 	ARY_SET_LEN(ary, j);
+    }
     ary_recycle_hash(hash);
 
     return ary;
@@ -3055,21 +3102,31 @@ rb_ary_uniq_bang(VALUE ary)
  *     
  *     a = [ "a", "a", "b", "b", "c" ]
  *     a.uniq   #=> ["a", "b", "c"]
+ *     c = [ "a:def", "a:xyz", "b:abc", "b:xyz", "c:jkl" ]
+ *     c.uniq {|s| s[/^\w+/]}  #=> [ "a:def", "b:abc", "c:jkl" ] 
  */
 
 static VALUE
 rb_ary_uniq(VALUE ary)
 {
-    VALUE hash = ary_make_hash(ary), v;
+    VALUE hash, uniq, v;
-    VALUE uniq = ary_new(rb_obj_class(ary), RHASH_SIZE(hash));
     long i;
 
+    if (rb_block_given_p()) {
+	hash = ary_make_hash_by(ary);
+	uniq = ary_new(rb_obj_class(ary), RHASH_SIZE(hash));
+	st_foreach(RHASH_TBL(hash), push_value, uniq);
+    }
+    else {
+	hash = ary_make_hash(ary);
+	uniq = ary_new(rb_obj_class(ary), RHASH_SIZE(hash));
 	for (i=0; i<RARRAY_LEN(ary); i++) {
 	    st_data_t vv = (st_data_t)(v = rb_ary_elt(ary, i));
 	    if (st_delete(RHASH_TBL(hash), &vv, 0)) {
 		rb_ary_push(uniq, v);
 	    }
 	}
+    }
     ary_recycle_hash(hash);
 
     return uniq;

test/ruby/test_array.rb

@@ -1240,6 +1240,11 @@ class TestArray < Test::Unit::TestCase
     assert_equal(@cls[1, 2, 3, 4, nil], a.uniq)
     assert_equal(b, a)
 
+    c = @cls["a:def", "a:xyz", "b:abc", "b:xyz", "c:jkl"]
+    d = c.dup
+    assert_equal(@cls[ "a:def", "b:abc", "c:jkl" ], c.uniq {|s| s[/^\w+/]})
+    assert_equal(d, c)
+
     assert_equal(@cls[1, 2, 3], @cls[1, 2, 3].uniq)
   end
 
@@ -1248,6 +1253,10 @@ class TestArray < Test::Unit::TestCase
     assert_equal(@cls[1, 2, 3, 4, nil], a.uniq!)
     assert_equal(@cls[1, 2, 3, 4, nil], a)
 
+    c = @cls["a:def", "a:xyz", "b:abc", "b:xyz", "c:jkl"]
+    assert_equal(@cls[ "a:def", "b:abc", "c:jkl" ], c.uniq! {|s| s[/^\w+/]})
+    assert_equal(@cls[ "a:def", "b:abc", "c:jkl" ], c)
+
     assert_nil(@cls[1, 2, 3].uniq!)
   end