From 2836acf045735771c12cb2395c8b8aea80571dbf Mon Sep 17 00:00:00 2001
From: yui-knk <spiketeika@gmail.com>
Date: Sun, 3 Aug 2014 12:45:54 +0900
Subject: [PATCH] implement flatten_merge, flatten, flatten!

---
 README.md         |  3 --
 mrblib/mrb_set.rb | 37 ++++++++++++++----
 test/mrb_set.rb   | 98 ++++++++++++++++++++++++++++++++++++++++++++++-
 3 files changed, 125 insertions(+), 13 deletions(-)

diff --git a/README.md b/README.md
index e88817fc9..fba04d7fb 100644
--- a/README.md
+++ b/README.md
@@ -40,9 +40,6 @@ These methods are not implemented yet.
 + taint
 + untaint
 + to_set
-+ flatten_merge
-+ flatten
-+ flatten!
 + classify
 + divide
 
diff --git a/mrblib/mrb_set.rb b/mrblib/mrb_set.rb
index e35db1b14..b1e5ca7f2 100644
--- a/mrblib/mrb_set.rb
+++ b/mrblib/mrb_set.rb
@@ -82,14 +82,35 @@ class Set
 #  def to_set
 #  end
 #
-#  def flatten_merge
-#  end
-#
-#  def flatten
-#  end
-#
-#  def flatten!
-#  end
+  def flatten_merge(set, seen = Set.new)
+    set.each { |e|
+      if e.is_a?(Set)
+        if seen.include?(e_id = e.object_id)
+          raise ArgumentError, "tried to flatten recursive Set"
+        end
+
+        seen.add(e_id)
+        flatten_merge(e, seen)
+        seen.delete(e_id)
+      else
+        add(e)
+      end
+    }
+
+    self
+  end
+
+  def flatten
+    self.class.new.flatten_merge(self)
+  end
+
+  def flatten!
+    if detect { |e| e.is_a?(Set) }
+      replace(flatten())
+    else
+      nil
+    end
+  end
 
   def include?(o)
     @hash.include?(o)
diff --git a/test/mrb_set.rb b/test/mrb_set.rb
index 216eb70ab..559c2c749 100644
--- a/test/mrb_set.rb
+++ b/test/mrb_set.rb
@@ -113,8 +113,102 @@ assert("Set#to_a") do
   assert_equal([1,2,3], ary.sort)
 end
 
-#assert("Set#flatten") do
-#end
+assert("Set#flatten") do
+  # test1
+  set1 = Set[
+    1,
+    Set[
+      5,
+      Set[7,
+        Set[0]
+      ],
+      Set[6,2],
+      1
+    ],
+    3,
+    Set[3,4]
+  ]
+
+  set2 = set1.flatten
+  set3 = Set.new(0..7)
+
+  assert_false set1.equal?(set2) # assert_not_same
+  assert_equal(set3, set2)
+
+
+  # test2; multiple occurrences of a set in an set
+  set1 = Set[1, 2]
+  set2 = Set[set1, Set[set1, 4], 3]
+
+  assert_nothing_raised {
+    set3 = set2.flatten
+  }
+
+  assert_equal(Set.new(1..4), set3)
+
+
+  # test3; recursion
+  set2 = Set[]
+  set1 = Set[1, set2]
+  set2.add(set1)
+
+  assert_raise(ArgumentError) {
+    set1.flatten
+  }
+end
+
+assert("Set#flatten!") do
+  # test1
+  set1 = Set[
+    1,
+    Set[
+      5,
+      Set[7,
+        Set[0]
+      ],
+      Set[6,2],
+      1
+    ],
+    3,
+    Set[3,4]
+  ]
+
+  set3 = Set.new(0..7)
+  orig_set1 = set1
+  set1.flatten!
+
+  assert_true orig_set1.equal?(set1) # assert_same
+  assert_equal(set3, set1)
+
+
+  # test2; multiple occurrences of a set in an set
+  set1 = Set[1, 2]
+  set2 = Set[set1, Set[set1, 4], 3]
+
+  assert_nothing_raised {
+    set2.flatten!
+  }
+
+  assert_equal(Set.new(1..4), set2)
+
+
+  # test3; recursion
+  set2 = Set[]
+  set1 = Set[1, set2]
+  set2.add(set1)
+
+  assert_raise(ArgumentError) {
+    set1.flatten!
+  }
+
+  # test5; miscellaneous
+  assert_nil(Set.new(0..31).flatten!)
+
+  x = Set[Set[],Set[1,2]].flatten!
+  y = Set[1,2]
+
+  assert_equal(x, y)
+end
 
 assert("Set#include?") do
   set = Set[1,2,3]