Merge pull request #8956 from headius/range_count_block

Handle argument or block passed to Range#count

Author: headius

core/src/main/java/org/jruby/RubyRange.java

@@ -1214,16 +1214,25 @@ public IRubyObject call(ThreadContext ctx, IRubyObject larg, Block blk) {
 
     @JRubyMethod
     public IRubyObject count(ThreadContext context, Block block) {
-        if (isBeginless || isEndless) return asFloat(context, RubyFloat.INFINITY);
+        if (!block.isGiven()) {
+            if (isBeginless || isEndless) return asFloat(context, RubyFloat.INFINITY);
 
-        if (begin instanceof RubyInteger) {
-            IRubyObject size = size(context);
-            if (!size.isNil()) return size;
+            if (begin instanceof RubyInteger) {
+                IRubyObject size = size(context);
+                if (!size.isNil()) return size;
+            }
+
+            // fall back on Enumerable logic for other cases
         }
 
         return RubyEnumerable.count(context, this, block);
     }
 
+    @JRubyMethod
+    public IRubyObject count(ThreadContext context, IRubyObject arg, Block block) {
+        return RubyEnumerable.count(context, this, arg, block);
+    }
+
     @JRubyMethod
     public IRubyObject minmax(ThreadContext context, Block block) {
         if (block.isGiven()) return Helpers.invokeSuper(context, this, context.runtime.getRange(), "minmax", NULL_ARRAY, block);

spec/ruby/core/range/count_spec.rb

@@ -9,4 +9,18 @@
     (...nil).count.should == inf
     (..10.0).count.should == inf
   end
+
+  it "accepts an argument for comparison using ==" do
+    (1..10).count(2).should == 1
+  end
+
+  it "uses a block for comparison" do
+    (1..10).count{|x| x%2==0 }.should == 5
+  end
+
+  it "ignores the block when given an argument" do
+    -> {
+      (1..10).count(4){|x| x%2==0 }.should == 1
+    }.should complain(/given block not used/)
+  end
 end