* lib/set.rb (Set#{delete_if,keep_if,collect!,reject!,select!,classify,divide},

  SortedSet#{delete_if,keep_if}): Return sized enumerators.

* test/test_set.rb: add test for above.

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

3 files changed, 79 insertions, 9 deletions

diff --git a/ChangeLog b/ChangeLog
index 35b8abf4ae..aac5b30851 100644
--- a/ChangeLog
+++ b/ChangeLog
@@ -1,3 +1,10 @@
+Wed May  4 17:53:15 2016  Kazuki Tsujimoto  <kazuki@callcc.net>
+
+	* lib/set.rb (Set#{delete_if,keep_if,collect!,reject!,select!,classify,divide},
+	  SortedSet#{delete_if,keep_if}): Return sized enumerators.
+
+	* test/test_set.rb: add test for above.
+
 Tue May  3 23:25:48 2016  Kazuhiro NISHIYAMA  <zn@mbf.nifty.com>
 
 	* numeric.c: [DOC] Update result of 123456789 ** -2.
diff --git a/lib/set.rb b/lib/set.rb
index a8f4345f35..70abeeb892 100644
--- a/lib/set.rb
+++ b/lib/set.rb
@@ -337,7 +337,7 @@ class Set
   # Deletes every element of the set for which block evaluates to
   # true, and returns self.
   def delete_if
-    block_given? or return enum_for(__method__)
+    block_given? or return enum_for(__method__) { size }
     # @hash.delete_if should be faster, but using it breaks the order
     # of enumeration in subclasses.
     select { |o| yield o }.each { |o| @hash.delete(o) }
@@ -347,7 +347,7 @@ class Set
   # Deletes every element of the set for which block evaluates to
   # false, and returns self.
   def keep_if
-    block_given? or return enum_for(__method__)
+    block_given? or return enum_for(__method__) { size }
     # @hash.keep_if should be faster, but using it breaks the order of
     # enumeration in subclasses.
     reject { |o| yield o }.each { |o| @hash.delete(o) }
@@ -356,7 +356,7 @@ class Set
 
   # Replaces the elements with ones returned by collect().
   def collect!
-    block_given? or return enum_for(__method__)
+    block_given? or return enum_for(__method__) { size }
     replace(self.class.new(self) { |o| yield(o) })
   end
   alias map! collect!
@@ -364,7 +364,7 @@ class Set
   # Equivalent to Set#delete_if, but returns nil if no changes were
   # made.
   def reject!(&block)
-    block or return enum_for(__method__)
+    block or return enum_for(__method__) { size }
     n = size
     delete_if(&block)
     self if size != n
@@ -373,7 +373,7 @@ class Set
   # Equivalent to Set#keep_if, but returns nil if no changes were
   # made.
   def select!(&block)
-    block or return enum_for(__method__)
+    block or return enum_for(__method__) { size }
     n = size
     keep_if(&block)
     self if size != n
@@ -468,7 +468,7 @@ class Set
   #             #     2001=>#<Set: {"c.rb", "d.rb", "e.rb"}>,
   #             #     2002=>#<Set: {"f.rb"}>}
   def classify # :yields: o
-    block_given? or return enum_for(__method__)
+    block_given? or return enum_for(__method__) { size }
 
     h = {}
 
@@ -496,7 +496,7 @@ class Set
   #             #            #<Set: {3, 4}>,
   #             #            #<Set: {6}>}>
   def divide(&func)
-    func or return enum_for(__method__)
+    func or return enum_for(__method__) { size }
 
     if func.arity == 2
       require 'tsort'
@@ -649,7 +649,7 @@ class SortedSet < Set
           end
 
           def delete_if
-            block_given? or return enum_for(__method__)
+            block_given? or return enum_for(__method__) { size }
             n = @hash.size
             super
             @keys = nil if @hash.size != n
@@ -657,7 +657,7 @@ class SortedSet < Set
           end
 
           def keep_if
-            block_given? or return enum_for(__method__)
+            block_given? or return enum_for(__method__) { size }
             n = @hash.size
             super
             @keys = nil if @hash.size != n
diff --git a/test/test_set.rb b/test/test_set.rb
index cb0727558e..9439974fea 100644
--- a/test/test_set.rb
+++ b/test/test_set.rb
@@ -432,6 +432,12 @@ class TC_Set < Test::Unit::TestCase
     ret = set.delete_if { |i| i % 3 == 0 }
     assert_same(set, ret)
     assert_equal(Set[1,2,4,5,7,8,10], set)
+
+    set = Set.new(1..10)
+    enum = set.delete_if
+    assert_equal(set.size, enum.size)
+    assert_same(set, enum.each { |i| i % 3 == 0 })
+    assert_equal(Set[1,2,4,5,7,8,10], set)
   end
 
   def test_keep_if
@@ -444,6 +450,12 @@ class TC_Set < Test::Unit::TestCase
     ret = set.keep_if { |i| i % 3 != 0 }
     assert_same(set, ret)
     assert_equal(Set[1,2,4,5,7,8,10], set)
+
+    set = Set.new(1..10)
+    enum = set.keep_if
+    assert_equal(set.size, enum.size)
+    assert_same(set, enum.each { |i| i % 3 != 0 })
+    assert_equal(Set[1,2,4,5,7,8,10], set)
   end
 
   def test_collect!
@@ -462,6 +474,22 @@ class TC_Set < Test::Unit::TestCase
 
     assert_same(set, ret)
     assert_equal(Set[2,4,6,'A','B','C',nil], set)
+
+    set = Set[1,2,3,'a','b','c',-1..1,2..4]
+    enum = set.collect!
+
+    assert_equal(set.size, enum.size)
+    assert_same(set, enum.each  { |i|
+      case i
+      when Numeric
+        i * 2
+      when String
+        i.upcase
+      else
+        nil
+      end
+    })
+    assert_equal(Set[2,4,6,'A','B','C',nil], set)
   end
 
   def test_reject!
@@ -474,6 +502,12 @@ class TC_Set < Test::Unit::TestCase
     ret = set.reject! { |i| i % 3 == 0 }
     assert_same(set, ret)
     assert_equal(Set[1,2,4,5,7,8,10], set)
+
+    set = Set.new(1..10)
+    enum = set.reject!
+    assert_equal(set.size, enum.size)
+    assert_same(set, enum.each { |i| i % 3 == 0 })
+    assert_equal(Set[1,2,4,5,7,8,10], set)
   end
 
   def test_select!
@@ -486,6 +520,12 @@ class TC_Set < Test::Unit::TestCase
     ret = set.select! { |i| i % 3 != 0 }
     assert_same(set, ret)
     assert_equal(Set[1,2,4,5,7,8,10], set)
+
+    set = Set.new(1..10)
+    enum = set.select!
+    assert_equal(set.size, enum.size)
+    assert_equal(nil, enum.each { |i| i <= 10 })
+    assert_equal(Set.new(1..10), set)
   end
 
   def test_merge
@@ -563,6 +603,18 @@ class TC_Set < Test::Unit::TestCase
     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])
+
+    set = Set.new(1..10)
+    enum = set.classify
+
+    assert_equal(set.size, enum.size)
+    ret = enum.each { |i| i % 3 }
+    assert_equal(3, ret.size)
+    assert_instance_of(Hash, ret)
+    ret.each_value { |value| assert_instance_of(Set, value) }
+    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
 
   def test_divide
@@ -596,6 +648,17 @@ class TC_Set < Test::Unit::TestCase
         raise "unexpected group: #{s.inspect}"
       end
     }
+
+    set = Set.new(1..10)
+    enum = set.divide
+    ret = enum.each { |i| i % 3 }
+
+    assert_equal(set.size, enum.size)
+    assert_equal(3, ret.size)
+    n = 0
+    ret.each { |s| n += s.size }
+    assert_equal(set.size, n)
+    assert_equal(set, ret.flatten)
   end
 
   def test_taintness