Merge pull request #3755 from christopheraue/array_bsearch_index

Added Array#bsearch_index

2 files changed, 81 insertions, 25 deletions

diff --git a/mrbgems/mruby-array-ext/mrblib/array.rb b/mrbgems/mruby-array-ext/mrblib/array.rb
index 1e6d4d581..716eabe06 100644
--- a/mrbgems/mruby-array-ext/mrblib/array.rb
+++ b/mrbgems/mruby-array-ext/mrblib/array.rb
@@ -595,31 +595,58 @@ class Array
   def bsearch(&block)
     return to_enum :bsearch unless block_given?
 
+    if idx = bsearch_index(&block)
+      self[idx]
+    else
+      nil
+    end
+  end
+
+  ##
+  #  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.
+
+  def bsearch_index(&block)
+    return to_enum :bsearch_index unless block_given?
+
     low = 0
-    high = self.size
+    high = size
     satisfied = false
+
     while low < high
-      mid = low + ((high - low) / 2).truncate
-      val = self[mid]
-      v = block.call(val)
-      if v.is_a?(Integer)
-        return val if v == 0
-        smaller = v < 0
-      elsif v == true
+      mid = ((low+high)/2).truncate
+      res = block.call self[mid]
+
+      case res
+      when 0 # find-any mode: Found!
+        return mid
+      when Numeric # find-any mode: Continue...
+        in_lower_half = res < 0
+      when true # find-min mode
+        in_lower_half = true
         satisfied = true
-        smaller = true
-      elsif v == false || v.nil?
-        smaller = false
+      when false, nil # find-min mode
+        in_lower_half = false
+      else
+        raise TypeError, 'invalid block result (must be numeric, true, false or nil)'
       end
-      if smaller
+
+      if in_lower_half
         high = mid
       else
         low = mid + 1
       end
     end
-    return nil if low == self.size
-    return nil unless satisfied
-    self[low]
+
+    satisfied ? low : nil
   end
 
   ##
diff --git a/mrbgems/mruby-array-ext/test/array.rb b/mrbgems/mruby-array-ext/test/array.rb
index 95a796cf9..c0db1b1cc 100644
--- a/mrbgems/mruby-array-ext/test/array.rb
+++ b/mrbgems/mruby-array-ext/test/array.rb
@@ -228,18 +228,47 @@ end
 
 assert("Array#bsearch") do
   # Find minimum mode
-  a = [0, 4, 7, 10, 12]
-  assert_include [4, 7], a.bsearch {|x| x >= 4 }
-  assert_equal 7, a.bsearch {|x| x >= 6 }
-  assert_equal 0, a.bsearch {|x| x >= -1 }
-  assert_nil a.bsearch {|x| x >= 100 }
+  a = [0, 2, 4]
+  assert_equal 0, a.bsearch{ |x| x >= -1 }
+  assert_equal 0, a.bsearch{ |x| x >= 0 }
+  assert_equal 2, a.bsearch{ |x| x >= 1 }
+  assert_equal 2, a.bsearch{ |x| x >= 2 }
+  assert_equal 4, a.bsearch{ |x| x >= 3 }
+  assert_equal 4, a.bsearch{ |x| x >= 4 }
+  assert_nil      a.bsearch{ |x| x >= 5 }
 
   # Find any mode
-  a = [0, 4, 7, 10, 12]
-  assert_include [4, 7], a.bsearch {|x| 1 - (x / 4).truncate }
-  assert_nil a.bsearch {|x| 4 - (x / 2).truncate }
-  assert_equal(nil, a.bsearch {|x| 1 })
-  assert_equal(nil, a.bsearch {|x| -1 })
+  a = [0, 4, 8]
+  def between(lo, x, hi)
+    if x < lo
+      1
+    elsif x > hi
+      -1
+    else
+      0
+    end
+  end
+  assert_nil      a.bsearch{ |x| between(-3, x, -1) }
+  assert_equal 0, a.bsearch{ |x| between(-1, x,  1) }
+  assert_nil      a.bsearch{ |x| between( 1, x,  3) }
+  assert_equal 4, a.bsearch{ |x| between( 3, x,  5) }
+  assert_nil      a.bsearch{ |x| between( 5, x,  7) }
+  assert_equal 8, a.bsearch{ |x| between( 7, x,  9) }
+  assert_nil      a.bsearch{ |x| between( 9, x, 11) }
+
+  assert_equal 0, a.bsearch{ |x| between( 0, x,  3) }
+  assert_equal 4, a.bsearch{ |x| between( 0, x,  4) }
+  assert_equal 4, a.bsearch{ |x| between( 4, x,  8) }
+  assert_equal 8, a.bsearch{ |x| between( 5, x,  8) }
+
+  # Invalid block result
+  assert_raise TypeError, 'invalid block result (must be numeric, true, false or nil)' do
+    a.bsearch{ 'I like to watch the world burn' }
+  end
+end
+
+assert("Array#bsearch_index") do
+  # tested through Array#bsearch
 end
 
 assert("Array#delete_if") do