array.c: bsearch_index

* array.c (rb_ary_bsearch_index): Implement Array#bsearch_index
  method, which is similar to bsearch and returns the index or
  nil.  [Feature #10730]

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

ChangeLog

@@ -1,3 +1,9 @@
+Fri Jun 12 16:28:17 2015  Radan Skoric  <radan.skoric@gmail.com>
+
+	* array.c (rb_ary_bsearch_index): Implement Array#bsearch_index
+	  method, which is similar to bsearch and returns the index or
+	  nil.  [Feature #10730]
+
 Thu Jun 11 19:11:22 2015  SHIBATA Hiroshi  <hsbt@ruby-lang.org>
 
 	* ext/zlib/zlib.c: Fix indentation for rdoc.

array.c

@@ -2547,6 +2547,8 @@ rb_ary_sort(VALUE ary)
     return ary;
 }
 
+static long ary_bsearch_index(VALUE ary);
+
 /*
  *  call-seq:
  *     ary.bsearch {|x| block }  -> elem
@@ -2602,10 +2604,40 @@ rb_ary_sort(VALUE ary)
 
 static VALUE
 rb_ary_bsearch(VALUE ary)
+{
+    long index_result = ary_bsearch_index(ary);
+
+    if (index_result < 0) return rb_ary_entry(ary, index_result);
+    return INT2FIX(index_result);
+}
+
+/*
+ *  call-seq:
+ *     ary.bsearch_index {|x| block }  -> int or nil
+ *
+ *  By using binary search, finds an index of a value from this array which
+ *  meets the given condition in O(log n) where n is the size of the array.
+ *
+ *  It supports two modes, depending on the nature of the block and they are
+ *  exactly the same as in the case of #bsearch method with the only difference
+ *  being that this method returns the index of the element instead of the
+ *  element itself. For more details consult the documentation for #bsearch.
+ */
+
+static VALUE
+rb_ary_bsearch_index(VALUE ary)
+{
+    long index_result = ary_bsearch_index(ary);
+
+    returns INT2FIX(index_result);
+}
+
+static long
+ary_bsearch_index(VALUE ary)
 {
     long low = 0, high = RARRAY_LEN(ary), mid;
-    int smaller = 0;
+    int smaller = 0, satisfied = 0;
-    VALUE v, val, satisfied = Qnil;
+    VALUE v, val;
 
     RETURN_ENUMERATOR(ary, 0, 0);
     while (low < high) {
@@ -2613,11 +2645,11 @@ rb_ary_bsearch(VALUE ary)
 	val = rb_ary_entry(ary, mid);
 	v = rb_yield(val);
 	if (FIXNUM_P(v)) {
-	    if (v == INT2FIX(0)) return val;
+	    if (v == INT2FIX(0)) return mid;
 	    smaller = (SIGNED_VALUE)v < 0; /* Fixnum preserves its sign-bit */
 	}
 	else if (v == Qtrue) {
-	    satisfied = val;
+	    satisfied = 1;
 	    smaller = 1;
 	}
 	else if (v == Qfalse || v == Qnil) {
@@ -2626,7 +2658,7 @@ rb_ary_bsearch(VALUE ary)
 	else if (rb_obj_is_kind_of(v, rb_cNumeric)) {
 	    const VALUE zero = INT2FIX(0);
 	    switch (rb_cmpint(rb_funcallv(v, id_cmp, 1, &zero), v, zero)) {
-	      case 0: return val;
+	      case 0: return mid;
 	      case 1: smaller = 1; break;
 	      case -1: smaller = 0;
 	    }
@@ -2643,7 +2675,8 @@ rb_ary_bsearch(VALUE ary)
 	    low = mid + 1;
 	}
     }
-    return satisfied;
+    if (!satisfied) return -1;
+    return low;
 }
 
 
@@ -5858,6 +5891,7 @@ Init_Array(void)
     rb_define_method(rb_cArray, "drop", rb_ary_drop, 1);
     rb_define_method(rb_cArray, "drop_while", rb_ary_drop_while, 0);
     rb_define_method(rb_cArray, "bsearch", rb_ary_bsearch, 0);
+    rb_define_method(rb_cArray, "bsearch_index", rb_ary_bsearch_index, 0);
     rb_define_method(rb_cArray, "any?", rb_ary_any_p, 0);
 
     id_cmp = rb_intern("<=>");

test/ruby/test_array.rb

@@ -2529,6 +2529,41 @@ class TestArray < Test::Unit::TestCase
     assert_include([4, 7], a.bsearch {|x| (2**100).coerce((1 - x / 4) * (2**100)).first })
   end
 
+  def test_bsearch_index_typechecks_return_values
+    assert_raise(TypeError) do
+      [1, 2, 42, 100, 666].bsearch_index {"not ok"}
+    end
+    assert_equal [1, 2, 42, 100, 666].bsearch_index {}, [1, 2, 42, 100, 666].bsearch_index {false}
+  end
+
+  def test_bsearch_index_with_no_block
+    enum = [1, 2, 42, 100, 666].bsearch_index
+    assert_nil enum.size
+    assert_equal 2, enum.each{|x| x >= 33 }
+  end
+
+  def test_bsearch_index_in_find_minimum_mode
+    a = [0, 4, 7, 10, 12]
+    assert_equal(1, a.bsearch_index {|x| x >=   4 })
+    assert_equal(2, a.bsearch_index {|x| x >=   6 })
+    assert_equal(0, a.bsearch_index {|x| x >=  -1 })
+    assert_equal(nil, a.bsearch_index {|x| x >= 100 })
+  end
+
+  def test_bsearch_index_in_find_any_mode
+    a = [0, 4, 7, 10, 12]
+    assert_include([1, 2], a.bsearch_index {|x| 1 - x / 4 })
+    assert_equal(nil, a.bsearch_index {|x| 4 - x / 2 })
+    assert_equal(nil, a.bsearch_index {|x| 1 })
+    assert_equal(nil, a.bsearch_index {|x| -1 })
+
+    assert_include([1, 2], a.bsearch_index {|x| (1 - x / 4) * (2**100) })
+    assert_equal(nil, a.bsearch_index {|x|   1  * (2**100) })
+    assert_equal(nil, a.bsearch_index {|x| (-1) * (2**100) })
+
+    assert_include([1, 2], a.bsearch_index {|x| (2**100).coerce((1 - x / 4) * (2**100)).first })
+  end
+
   def test_shared_marking
     reduce = proc do |s|
       s.gsub(/(verify_internal_consistency_reachable_i:\sWB\smiss\s\S+\s\(T_ARRAY\)\s->\s)\S+\s\((proc|T_NONE)\)\n