module GenericBinarySearch
  use bool.Bool
  use int.Int
  use int.ComputerDivision
  use ref.Ref

  let binary_search (pred: int -> bool) (low: int) (high: int): int
    requires { pred low = False }
    requires { pred high = True }
    requires { forall i j. i <= j -> (pred i -> pred j) }
    ensures  { pred result = False /\ pred (result + 1) = True }
  =
    let cur_low = ref low in
    let cur_high = ref high in

    while !cur_high - !cur_low > 1 do
      variant { !cur_high - !cur_low }
      invariant { pred !cur_low = False /\ pred !cur_high = True }

      let mid = div (!cur_high + !cur_low) 2 in
      if not (pred mid) then
          cur_low := mid
      else
          cur_high := mid
    done;
    !cur_low
end

module ArrayBinarySearch
  use bool.Bool
  use array.Array
  use int.Int
  use option.Option

  use GenericBinarySearch

  exception Not_found (* raised to signal a search failure *)

  (* type t is the type of sorted arrays of integers *)
  type t = {
    arr: array int
  } invariant { forall i1 i2. 0 <= i1 <= i2 < length arr -> arr[i1] <= arr[i2] }

  let function pred (a: { t }) (x: int) (i: int): bool
  = if i <= -1 then False else
    if i >= length a.arr then True else
    x <= a.arr[i]

  (* As specified above, the type `t` is the type of sorted arrays *)

  let arr_binary_search (a: t) (x: int): int
    ensures  { 0 <= result < length a.arr /\ a.arr[result] = x }
    raises   { Not_found -> forall i. 0 <= i < length a.arr -> a.arr[i] <> x }
  =
    let idx = binary_search (pred a x) (-1) (length a.arr) in
    if idx + 1 < length a.arr then if a.arr[idx + 1] = x then return idx + 1;

    raise Not_found
end

(*
 *  To prove this with why3, first save the code above to a file (say 'bin_search.mlw').
 *  You will probably need to have the z3 prover installed. Then run the following:
 *
 *  $ why3 config --detect   (if you haven't done so already after installing the provers)
 *
 *  $ why3 prove -P z3 bin_search.mlw
 *  bin_search.mlw GenericBinarySearch VC binary_search: Valid (0.02s)
 *  bin_search.mlw ArrayBinarySearch VC t: Valid (0.07s)
 *  bin_search.mlw ArrayBinarySearch VC pred: Valid (0.02s)
 *  bin_search.mlw ArrayBinarySearch VC arr_binary_search: Valid (0.03s)
 *)