From 93a8b08f59c60de60be7eb4aa4f72b7da3a9b3aa Mon Sep 17 00:00:00 2001
From: yui-knk <spiketeika@gmail.com>
Date: Sun, 3 Aug 2014 13:06:30 +0900
Subject: [PATCH] implement #classify

---
 README.md         |  1 -
 mrblib/mrb_set.rb | 13 +++++++++++--
 test/mrb_set.rb   | 15 ++++++++++++---
 3 files changed, 23 insertions(+), 6 deletions(-)

diff --git a/README.md b/README.md
index fba04d7fb..8d9818e51 100644
--- a/README.md
+++ b/README.md
@@ -40,7 +40,6 @@ These methods are not implemented yet.
 + taint
 + untaint
 + to_set
-+ classify
 + divide
 
 ## License
diff --git a/mrblib/mrb_set.rb b/mrblib/mrb_set.rb
index 2b0d751c6..0b45e9a2f 100644
--- a/mrblib/mrb_set.rb
+++ b/mrblib/mrb_set.rb
@@ -284,8 +284,17 @@ class Set
     @hash.eql?(o.instance_variable_get(:@hash))
   end
 
-  # def classify
-  # end
+  def classify
+    return to_enum :classify unless block_given?
+    h = {}
+
+    each { |i|
+      x = yield(i)
+      (h[x] ||= self.class.new).add(i)
+    }
+
+    h
+  end
 
   # def divide
   # end
diff --git a/test/mrb_set.rb b/test/mrb_set.rb
index 0751d08d7..0234b7692 100644
--- a/test/mrb_set.rb
+++ b/test/mrb_set.rb
@@ -549,9 +549,18 @@ assert("Set#==") do
   assert_equal(set1.clone, set1)
 end
 
-#assert("Set#classify") do
-#end
-#
+assert("Set#classify") do
+  set = Set.new(1..10)
+  ret = set.classify { |i| i % 3 }
+
+  assert_equal(3, ret.size)
+  assert_equal(Hash, ret.class)
+  ret.each_value { |v| assert_equal(Set, v.class) }
+  assert_equal(Set[3,6,9], ret[0])
+  assert_equal(Set[1,4,7,10], ret[1])
+  assert_equal(Set[2,5,8], ret[2])
+end
+
 #assert("Set#divide") do
 #end