(**

 {1 The VerifyThis competition at FM2012: Problem 2}

  see {h this web page (http://fm2012.verifythis.org/challenges)}

*)

module PrefixSumRec

   use int.Int
   use int.ComputerDivision
   use int.Power

   use map.Map
   use array.Array
   use array.ArraySum

(** {2 Preliminary lemmas on division by 2 and power of 2} *)

  (** Needed for the proof of phase1_frame and phase1_frame2 *)
  lemma Div_mod_2:
    forall x:int. x >= 0 -> x >= 2 * div x 2 >= x-1

  (** The step must be a power of 2 *)
  predicate is_power_of_2 (x:int) = exists k:int. (k >= 0 /\ x = power 2 k)

  (* needed *)
  lemma is_power_of_2_1:
    forall x:int. is_power_of_2 x -> x > 1 -> 2 * div x 2 = x

(**

{2 Modeling the "upsweep" phase }

*)

  (** Shortcuts *)
  function go_left (left right:int) : int =
    let space = right - left in left - div space 2

  function go_right (left right:int) : int =
    let space = right - left in right - div space 2

(**
   Description in a purely logic way the effect of the first phase
     "upsweep" of the algorithm.  The second phase "downsweep" then
     traverses the array in the same way than the first phase. Hence,
     the inductive nature of this definition is not an issue. *)

  inductive phase1 (left:int) (right:int) (a0: array int) (a: array int) =
    | Leaf : forall left right:int, a0 a : array int.
       right = left + 1 ->
       a[left] = a0[left] ->
       phase1 left right a0 a
    | Node : forall left right:int, a0 a : array int.
       right > left + 1 ->
       phase1 (go_left left right) left a0 a ->
       phase1 (go_right left right) right a0 a ->
       a[left] = sum a0 (left-(right-left)+1) (left+1) ->
       phase1 left right a0 a

  (** Frame properties of the "phase1" inductive *)

  (** frame lemma for "phase1" on fourth argument.
     needed to prove both upsweep, downsweep and compute_sums
   *)
  let rec lemma phase1_frame (left right:int) (a0 a a' : array int) : unit
    requires { forall i:int. left-(right-left) < i < right ->
               a[i] = a'[i]}
    requires { phase1 left right a0 a }
    variant { right-left }
    ensures { phase1 left right a0 a' } =
    if right > left + 1 then begin
       phase1_frame (go_left left right) left a0 a a';
       phase1_frame (go_right left right) right a0 a a'
    end

  (** frame lemma for "phase1" on third argument.
     needed to prove upsweep and compute_sums
   *)
  let rec lemma phase1_frame2 (left right:int) (a0 a0' a : array int) : unit
    requires { forall i:int. left-(right-left) < i < right ->
               a0[i] = a0'[i]}
    requires { phase1 left right a0 a }
    variant { right-left }
    ensures { phase1 left right a0' a } =
    if right > left + 1 then begin
       phase1_frame2 (go_left left right) left a0 a0' a;
       phase1_frame2 (go_right left right) right a0 a0' a
    end

(** {2 The upsweep phase}

  First function: modify partially the table and compute some
      intermediate partial sums

*)

  let rec upsweep (left right: int) (a: array int)
    requires { 0 <= left < right < a.length }
    requires { -1 <= left - (right - left) }
    requires { is_power_of_2 (right - left) }
    variant { right - left }
    ensures { phase1 left right (old a) a }
    ensures { let space = right - left in
      a[right] = sum (old a) (left-space+1) (right+1) /\
      (forall i: int. i <= left-space -> a[i] = (old a)[i]) /\
      (forall i: int. i > right -> a[i] = (old a)[i]) }
  = let space = right - left in
    if right > left+1 then begin
      upsweep (left - div space 2) left a;
      upsweep (right - div space 2) right a;
      assert { phase1 (left - div space 2) left (old a) a };
      assert { phase1 (right - div space 2) right (old a) a };
      assert { a[left] = sum (old a) (left-(right-left)+1) (left+1) };
      assert { a[right] = sum (old a) (left+1) (right+1) }
    end;
    a[right] <- a[left] + a[right];
    assert {
      right > left+1 -> phase1 (left - div space 2) left (old a) a };
    assert {
      right > left+1 -> phase1 (right - div space 2) right (old a) a }

(** {2 The downsweep phase} *)

  (** The property we ultimately want to prove *)
  predicate partial_sum (left:int) (right:int) (a0 a : array int) =
    forall i : int. (left-(right-left)) < i <= right -> a[i] = sum a0 0 i

  (** Second function: complete the partial using the remaining intial
      value and the partial sum already computed *)
  let rec downsweep (left right: int) (ghost a0 : array int) (a : array int)
    requires { 0 <= left < right < a.length }
    requires { -1 <= left - (right - left) }
    requires { is_power_of_2 (right - left) }
    requires { a[right] = sum a0 0 (left-(right-left) + 1) }
    requires { phase1 left right a0 a }
    variant  { right - left }
    ensures { partial_sum left right a0 a }
    ensures { forall i: int. i <= left-(right-left) -> a[i] = (old a)[i] }
    ensures { forall i: int. i > right -> a[i] = (old a)[i] }
  = let tmp = a[right] in
    assert { a[right] = sum a0 0 (left-(right-left) + 1) };
    assert { a[left] = sum a0 (left-(right-left)+1) (left+1) };
    a[right] <- a[right] + a[left];
    a[left] <- tmp;
    assert { a[right] = sum a0 0 (left + 1) };
    if right > left+1 then
    let space = right - left in
    assert { phase1 (go_left left right) left a0 (old a) };
    assert { phase1 (go_right left right) right a0 (old a) };
    assert { phase1 (go_right left right) right a0 a };
    downsweep (left - div space 2) left a0 a;
    assert { phase1 (go_right left right) right a0 a };
    downsweep (right - div space 2) right a0 a;
    assert { partial_sum (left - div space 2) left a0 a };
    assert { partial_sum (right - div space 2) right a0 a }

(** {2 The main procedure} *)

  let compute_sums a
    requires { a.length >= 2 }
    requires { is_power_of_2 a.length }
    ensures { forall i : int. 0 <= i < a.length -> a[i] = sum (old a) 0 i }
  = let a0 = ghost (copy a) in
    let l = a.length in
    let left = div l 2 - 1 in
    let right = l - 1 in
    upsweep left right a;
    (* needed for the precondition of downsweep *)
    assert { phase1 left right a0 a };
    a[right] <- 0;
    downsweep left right a0 a;
    (* needed to prove the post-condition *)
    assert { forall i : int.
      left-(right-left) < i <= right -> a[i] = sum a0 0 i }

(** {2 A simple test} *)

  let test_harness () =
    let a = make 8 0 in
    (* needed for the precondition of compute_sums *)
    assert { power 2 3 = a.length };
    a[0] <- 3; a[1] <- 1; a[2] <- 7; a[3] <- 0;
    a[4] <- 4; a[5] <- 1; a[6] <- 6; a[7] <- 3;
    compute_sums a;
    assert { a[0] = 0 };
    assert { a[1] = 3 };
    assert { a[2] = 4 };
    assert { a[3] = 11 };
    assert { a[4] = 11 };
    assert { a[5] = 15 };
    assert { a[6] = 16 };
    assert { a[7] = 22 }

  exception BenchFailure

  let bench () raises { BenchFailure -> true } =
    let a = make 8 0 in
    (* needed for the precondition of compute_sums *)
    assert { power 2 3 = a.length };
    a[0] <- 3; a[1] <- 1; a[2] <- 7; a[3] <- 0;
    a[4] <- 4; a[5] <- 1; a[6] <- 6; a[7] <- 3;
    compute_sums a;
    if a[0] <> 0 then raise BenchFailure;
    if a[1] <> 3 then raise BenchFailure;
    if a[2] <> 4 then raise BenchFailure;
    if a[3] <> 11 then raise BenchFailure;
    if a[4] <> 11 then raise BenchFailure;
    if a[5] <> 15 then raise BenchFailure;
    if a[6] <> 16 then raise BenchFailure;
    if a[7] <> 22 then raise BenchFailure;
    a

end
