1 files changed, 91 insertions, 24 deletions
diff --git a/lib/pure/algorithm.nim b/lib/pure/algorithm.nim
index 22793ab1c..81badfae6 100644
--- a/lib/pure/algorithm.nim
+++ b/lib/pure/algorithm.nim
@@ -64,40 +64,72 @@ proc reversed*[T](a: openArray[T]): seq[T] =
   ## returns the reverse of the array `a`.
   reversed(a, 0, a.high)
 
-proc binarySearch*[T](a: openArray[T], key: T): int =
+proc binarySearch*[T, K](a: openArray[T], key: K,
+              cmp: proc (x: T, y: K): int {.closure.}): int =
   ## binary search for `key` in `a`. Returns -1 if not found.
-  if ((a.len - 1) and a.len) == 0 and a.len > 0:
-    # when `a.len` is a power of 2, a faster div can be used.
-    var step = a.len div 2
+  ##
+  ## `cmp` is the comparator function to use, the expected return values are
+  ## the same as that of system.cmp.
+  if a.len == 0:
+    return -1
+
+  let len = a.len
+
+  if len == 1:
+    if cmp(a[0], key) == 0:
+      return 0
+    else:
+      return -1
+
+  if (len and (len - 1)) == 0:
+    # when `len` is a power of 2, a faster shr can be used.
+    var step = len shr 1
+    var cmpRes: int
     while step > 0:
-      if a[result or step] <= key:
-        result = result or step
+      let i = result or step
+      cmpRes = cmp(a[i], key)
+      if cmpRes == 0:
+        return i
+
+      if cmpRes < 1:
+        result = i
       step = step shr 1
-    if a[result] != key: result = -1
+    if cmp(a[result], key) != 0: result = -1
   else:
-    var b = len(a)
+    var b = len
+    var cmpRes: int
     while result < b:
-      var mid = (result + b) div 2
-      if a[mid] < key: result = mid + 1
-      else: b = mid
-    if result >= len(a) or a[result] != key: result = -1
+      var mid = (result + b) shr 1
+      cmpRes = cmp(a[mid], key)
+      if cmpRes == 0:
+        return mid
+
+      if cmpRes < 0:
+        result = mid + 1
+      else:
+        b = mid
+    if result >= len or cmp(a[result], key) != 0: result = -1
+
+proc binarySearch*[T](a: openArray[T], key: T): int =
+  ## binary search for `key` in `a`. Returns -1 if not found.
+  binarySearch(a, key, cmp[T])
 
 proc smartBinarySearch*[T](a: openArray[T], key: T): int {.deprecated.} =
   ## **Deprecated since version 0.18.1**; Use ``binarySearch`` instead.
-  binarySearch(a,key)
+  binarySearch(a, key, cmp[T])
 
 const
   onlySafeCode = true
 
 proc lowerBound*[T, K](a: openArray[T], key: K, cmp: proc(x: T, k: K): int {.closure.}): int =
-  ## same as binarySearch except that if key is not in `a` then this
-  ## returns the location where `key` would be if it were. In other
-  ## words if you have a sorted sequence and you call
+  ## Returns a position to the first element in the `a` that is greater than `key`, or last
+  ## if no such element is found. In other words if you have a sorted sequence and you call
   ## insert(thing, elm, lowerBound(thing, elm))
   ## the sequence will still be sorted.
   ##
-  ## `cmp` is the comparator function to use, the expected return values are
+  ## The first version uses `cmp` to compare the elements. The expected return values are
   ## the same as that of system.cmp.
+  ## The second version uses the default comparison function `cmp`.
   ##
   ## example::
   ##
@@ -108,7 +140,7 @@ proc lowerBound*[T, K](a: openArray[T], key: K, cmp: proc(x: T, k: K): int {.clo
   var count = a.high - a.low + 1
   var step, pos: int
   while count != 0:
-    step = count div 2
+    step = count shr 1
     pos = result + step
     if cmp(a[pos], key) < 0:
       result = pos + 1
@@ -118,6 +150,36 @@ proc lowerBound*[T, K](a: openArray[T], key: K, cmp: proc(x: T, k: K): int {.clo
 
 proc lowerBound*[T](a: openArray[T], key: T): int = lowerBound(a, key, cmp[T])
 
+proc upperBound*[T, K](a: openArray[T], key: K, cmp: proc(x: T, k: K): int {.closure.}): int =
+  ## Returns a position to the first element in the `a` that is not less
+  ## (i.e. greater or equal to) than `key`, or last if no such element is found.
+  ## In other words if you have a sorted sequence and you call
+  ## insert(thing, elm, upperBound(thing, elm))
+  ## the sequence will still be sorted.
+  ##
+  ## The first version uses `cmp` to compare the elements. The expected return values are
+  ## the same as that of system.cmp.
+  ## The second version uses the default comparison function `cmp`.
+  ##
+  ## example::
+  ##
+  ##   var arr = @[1,2,3,4,6,7,8,9]
+  ##   arr.insert(5, arr.upperBound(4))
+  ##   # after running the above arr is `[1,2,3,4,5,6,7,8,9]`
+  result = a.low
+  var count = a.high - a.low + 1
+  var step, pos: int
+  while count != 0:
+    step = count shr 1
+    pos = result + step
+    if cmp(a[pos], key) <= 0:
+      result = pos + 1
+      count -= step + 1
+    else:
+      count = step
+
+proc upperBound*[T](a: openArray[T], key: T): int = upperBound(a, key, cmp[T])
+
 template `<-` (a, b) =
   when false:
     a = b
@@ -514,21 +576,26 @@ when isMainModule:
     doAssert lowerBound([1,2,2,3], 4, system.cmp[int]) == 4
     doAssert lowerBound([1,2,3,10], 11) == 4
 
+  block upperBound:
+    doAssert upperBound([1,2,4], 3, system.cmp[int]) == 2
+    doAssert upperBound([1,2,2,3], 3, system.cmp[int]) == 4
+    doAssert upperBound([1,2,3,5], 3) == 3
+
   block fillEmptySeq:
     var s = newSeq[int]()
     s.fill(0)
 
   block testBinarySearch:
     var noData: seq[int]
-    doAssert binarySearch(noData, 7)    == -1
+    doAssert binarySearch(noData, 7) == -1
     let oneData = @[1]
-    doAssert binarySearch(oneData, 1)   == 0
-    doAssert binarySearch(onedata, 7)   == -1
+    doAssert binarySearch(oneData, 1) == 0
+    doAssert binarySearch(onedata, 7) == -1
     let someData = @[1,3,4,7]
-    doAssert binarySearch(someData, 1)  == 0
-    doAssert binarySearch(somedata, 7)  == 3
+    doAssert binarySearch(someData, 1) == 0
+    doAssert binarySearch(somedata, 7) == 3
     doAssert binarySearch(someData, -1) == -1
-    doAssert binarySearch(someData, 5)  == -1
+    doAssert binarySearch(someData, 5) == -1
     doAssert binarySearch(someData, 13) == -1
     let moreData = @[1,3,5,7,4711]
     doAssert binarySearch(moreData, -1) == -1