add Enumerable#group_by

2 files changed, 26 insertions, 0 deletions

diff --git a/mrbgems/mruby-enum-ext/mrblib/enum.rb b/mrbgems/mruby-enum-ext/mrblib/enum.rb
index b63de0d31..f250d39f1 100644
--- a/mrbgems/mruby-enum-ext/mrblib/enum.rb
+++ b/mrbgems/mruby-enum-ext/mrblib/enum.rb
@@ -142,4 +142,23 @@ module Enumerable
     block.call(ary) unless ary.empty?
   end
 
+  ##
+  # call-seq:
+  #    enum.group_by {| obj | block }  -> a_hash
+  #
+  # Returns a hash, which keys are evaluated result from the
+  # block, and values are arrays of elements in <i>enum</i>
+  # corresponding to the key.
+  #
+  #    (1..6).group_by {|i| i%3}   #=> {0=>[3, 6], 1=>[1, 4], 2=>[2, 5]}
+
+  def group_by(&block)
+    h = {}
+    self.each do |e|
+      key = block.call(e)
+      h.key?(key) ? (h[key] << e) : (h[key] = [e])
+    end
+    h
+  end
+
 end
diff --git a/mrbgems/mruby-enum-ext/test/enum.rb b/mrbgems/mruby-enum-ext/test/enum.rb
index 42c9eff8c..aa56cdf84 100644
--- a/mrbgems/mruby-enum-ext/test/enum.rb
+++ b/mrbgems/mruby-enum-ext/test/enum.rb
@@ -35,3 +35,10 @@ assert("Enumerable#each_slice") do
   assert_equal  a, [[1, 2, 3], [4, 5, 6], [7, 8, 9], [10]]
 end
 
+assert("Enumerable#group_by") do
+  r = (1..6).group_by {|i| i % 3 }
+  assert_equal r[0], [3, 6]
+  assert_equal r[1], [1, 4]
+  assert_equal r[2], [2, 5]
+end
+