Merge pull request #2078 from suzukaze/add-block-with-array.uniq_bang

Add block argument with Array#uniq! and Array#uniq

2 files changed, 40 insertions, 8 deletions

diff --git a/mrbgems/mruby-array-ext/mrblib/array.rb b/mrbgems/mruby-array-ext/mrblib/array.rb
index 7f48811f1..6c47235fe 100644
--- a/mrbgems/mruby-array-ext/mrblib/array.rb
+++ b/mrbgems/mruby-array-ext/mrblib/array.rb
@@ -1,7 +1,8 @@
 class Array
   ##
   # call-seq:
-  #    ary.uniq! -> ary or nil
+  #    ary.uniq!                -> ary or nil
+  #    ary.uniq! { |item| ... } -> ary or nil
   #
   # Removes duplicate elements from +self+.
   # Returns <code>nil</code> if no changes are made (that is, no
@@ -11,13 +12,27 @@ class Array
   #    a.uniq!   #=> ["a", "b", "c"]
   #    b = [ "a", "b", "c" ]
   #    b.uniq!   #=> nil
+  #    c = [["student","sam"], ["student","george"], ["teacher","matz"]]
+  #    c.uniq! { |s| s.first } # => [["student", "sam"], ["teacher", "matz"]]
   #
-  def uniq!
+  def uniq!(&block)
     ary = self.dup
     result = []
-    while ary.size > 0
-      result << ary.shift
-      ary.delete(result.last)
+    if block
+      hash = {}
+      while ary.size > 0
+        val = ary.shift
+        key = block.call(val)
+        hash[key] = val unless hash.has_key?(key)
+      end
+      hash.each_value do |value|
+        result << value
+      end
+    else
+      while ary.size > 0
+        result << ary.shift
+        ary.delete(result.last)
+      end
     end
     if result.size == self.size
       nil
@@ -28,16 +43,24 @@ class Array
 
   ##
   # call-seq:
-  #    ary.uniq   -> new_ary
+  #    ary.uniq                -> new_ary
+  #    ary.uniq { |item| ... } -> new_ary
   #
   # Returns a new array by removing duplicate values in +self+.
   #
   #    a = [ "a", "a", "b", "b", "c" ]
   #    a.uniq   #=> ["a", "b", "c"]
   #
-  def uniq
+  #    b = [["student","sam"], ["student","george"], ["teacher","matz"]]
+  #    b.uniq { |s| s.first } # => [["student", "sam"], ["teacher", "matz"]]
+  #
+  def uniq(&block)
     ary = self.dup
-    ary.uniq!
+    if block
+      ary.uniq!(&block)
+    else 
+      ary.uniq!
+    end
     ary
   end
 
diff --git a/mrbgems/mruby-array-ext/test/array.rb b/mrbgems/mruby-array-ext/test/array.rb
index 9a0f25f9a..d157a5b4d 100644
--- a/mrbgems/mruby-array-ext/test/array.rb
+++ b/mrbgems/mruby-array-ext/test/array.rb
@@ -36,12 +36,21 @@ assert("Array#uniq!") do
 
   b = [ "a", "b", "c" ]
   assert_nil b.uniq!
+
+  c = [["student","sam"], ["student","george"], ["teacher","matz"]]
+  assert_equal [["student", "sam"], ["teacher", "matz"]], c.uniq! { |s| s.first }
+
+  d = [["student","sam"], ["teacher","matz"]]
+  assert_nil d.uniq! { |s| s.first }
 end
 
 assert("Array#uniq") do
   a = [1, 2, 3, 1]
   assert_equal [1, 2, 3], a.uniq
   assert_equal [1, 2, 3, 1], a
+
+  b = [["student","sam"], ["student","george"], ["teacher","matz"]]
+  assert_equal [["student", "sam"], ["teacher", "matz"]], b.uniq { |s| s.first } 
 end
 
 assert("Array#-") do