(** {1 Range Sum Queries}

We are interested in specifying and proving correct
data structures that support efficient computation of the sum of the
values over an arbitrary range of an array.
Concretely, given an array of integers `a`, and given a range
delimited by indices `i` (inclusive) and `j` (exclusive), we wish
to compute the value: `\sum_{k=i}^{j-1} a[k]`.

In the first part, we consider a simple loop
for computing the sum in linear time.

In the second part, we introduce a cumulative sum array
that allows answering arbitrary range queries in constant time.

In the third part, we explore a tree data structure that
supports modification of values from the underlying array `a`,
with logarithmic time operations.

*)


(** {2 Specification of Range Sum} *)

module ArraySum

  use int.Int
  use array.Array

  (** `sum a i j` denotes the sum `\sum_{i <= k < j} a[k]`.
      It is axiomatizated by the two following axioms expressing
      the recursive definition

      if `i <= j` then `sum a i j = 0`

      if `i < j` then `sum a i j = a[i] + sum a (i+1) j`

  *)
  let rec function sum (a:array int) (i j:int) : int
   requires { 0 <= i <= j <= a.length }
   variant { j - i }
   = if j <= i then 0 else a[i] + sum a (i+1) j

  (** lemma for summation from the right:

      if `i < j` then `sum a i j = sum a i (j-1) + a[j-1]`

 *)
  lemma sum_right : forall a : array int, i j : int.
    0 <= i < j <= a.length  ->
    sum a i j = sum a i (j-1) + a[j-1]

end




(** {2 First algorithm, a linear one} *)

module Simple

  use int.Int
  use array.Array
  use ArraySum
  use ref.Ref

  (** `query a i j` returns the sum of elements in `a` between
      index `i` inclusive and index `j` exclusive *)
  let query (a:array int) (i j:int) : int
    requires { 0 <= i <= j <= a.length }
    ensures { result = sum a i j }
  = let s = ref 0 in
    for k=i to j-1 do
      invariant { !s = sum a i k }
      s := !s + a[k]
    done;
    !s

end




(** {2 Additional lemmas on `sum`}
  needed in the remaining code *)

module ExtraLemmas

  use int.Int
  use array.Array
  use ArraySum

  (** summation in adjacent intervals *)
  lemma sum_concat : forall a:array int, i j k:int.
    0 <= i <= j <= k <= a.length ->
    sum a i k = sum a i j + sum a j k

  (** Frame lemma for `sum`, that is `sum a i j` depends only
      of values of `a[i..j-1]` *)
  lemma sum_frame : forall a1 a2 : array int, i j : int.
    0 <= i <= j ->
    j <= a1.length ->
    j <= a2.length ->
    (forall k : int. i <= k < j -> a1[k] = a2[k]) ->
    sum a1 i j = sum a2 i j

  (** Updated lemma for `sum`: how does `sum a i j` changes when
      `a[k]` is changed for some `k` in `[i..j-1]` *)
  lemma sum_update : forall a:array int, i v l h:int.
    0 <= l <= i < h <= a.length ->
    sum (a[i<-v]) l h = sum a l h + v - a[i]


end




(** {2 Algorithm 2: using a cumulative array}

   creation of cumulative array is linear

   query is in constant time

   array update is linear

*)


module CumulativeArray

  use int.Int
  use array.Array
  use ArraySum
  use ExtraLemmas

  predicate is_cumulative_array_for (c:array int) (a:array int) =
    c.length = a.length + 1 /\
    forall i. 0 <= i < c.length -> c[i] = sum a 0 i

  (** `create a` builds the cumulative array associated with `a`. *)
  let create (a:array int) : array int
    ensures { is_cumulative_array_for result a }
  = let l = a.length in
    let s = Array.make (l+1) 0 in
    for i=1 to l do
      invariant { forall k. 0 <= k < i -> s[k] = sum a 0 k }
      s[i] <- s[i-1] + a[i-1]
    done;
    s

  (** `query c i j a` returns the sum of elements in `a` between
      index `i` inclusive and index `j` exclusive, in constant time *)
  let query (c:array int) (i j:int) (ghost a:array int): int
    requires { is_cumulative_array_for c a }
    requires { 0 <= i <= j < c.length }
    ensures { result = sum a i j }
  = c[j] - c[i]


  (** `update c i v a` updates cell `a[i]` to value `v` and updates
      the cumulative array `c` accordingly *)
  let update (c:array int) (i:int) (v:int) (ghost a:array int) : unit
    requires { is_cumulative_array_for c a }
    requires { 0 <= i < a.length }
    writes  { c, a }
    ensures { is_cumulative_array_for c a }
    ensures { a[i] = v }
    ensures { forall k. 0 <= k < a.length /\ k <> i ->
              a[k] = (old a)[k] }
  = let incr = v - c[i+1] + c[i] in
    a[i] <- v;
    for j=i+1 to c.length-1 do
      invariant { forall k. j <= k < c.length -> c[k] = sum a 0 k - incr }
      invariant { forall k. 0 <= k < j -> c[k] = sum a 0 k }
      c[j] <- c[j] + incr
    done

end






(** {2 Algorithm 3: using a cumulative tree}

  creation is linear

  query is logarithmic

  update is logarithmic

*)




module CumulativeTree

  use int.Int
  use array.Array
  use ArraySum
  use ExtraLemmas
  use int.ComputerDivision

  type indexes =
    { low : int;
      high : int;
      isum : int;
    }

  type tree = Leaf indexes | Node indexes tree tree

  let function indexes (t:tree) : indexes =
    match t with
    | Leaf ind -> ind
    | Node ind _ _ -> ind
    end

  predicate is_indexes_for (ind:indexes) (a:array int) (i j:int) =
    ind.low = i /\ ind.high = j /\
    0 <= i < j <= a.length /\
    ind.isum = sum a i j

  predicate is_tree_for (t:tree) (a:array int) (i j:int) =
    match t with
    | Leaf ind ->
        is_indexes_for ind a i j /\ j = i+1
    | Node ind l r ->
        is_indexes_for ind a i j /\
        i = l.indexes.low /\ j = r.indexes.high /\
        let m = l.indexes.high in
        m = r.indexes.low /\
        i < m < j /\ m = div (i+j) 2 /\
        is_tree_for l a i m /\
        is_tree_for r a m j
    end

  (** {3 creation of cumulative tree} *)

  let rec tree_of_array (a:array int) (i j:int) : tree
    requires { 0 <= i < j <= a.length }
    variant { j - i }
    ensures { is_tree_for result a i j }
    = if i+1=j then begin
       Leaf { low = i; high = j; isum = a[i] }
       end
      else
        begin
        let m = div (i+j) 2 in
        assert { i < m < j };
        let l = tree_of_array a i m in
        let r = tree_of_array a m j in
        let s = l.indexes.isum + r.indexes.isum in
        assert { s = sum a i j };
        Node { low = i; high = j; isum = s} l r
        end


  let create (a:array int) : tree
    requires { a.length >= 1 }
    ensures { is_tree_for result a 0 a.length }
  = tree_of_array a 0 a.length


(** {3 query using cumulative tree} *)


  let rec query_aux (t:tree) (ghost a: array int)
      (i j:int) : int
    requires { is_tree_for t a t.indexes.low t.indexes.high }
    requires { 0 <= t.indexes.low <= i < j <= t.indexes.high <= a.length }
    variant { t }
    ensures { result = sum a i j }
  = match t with
    | Leaf ind ->
      ind.isum
    | Node ind l r ->
      let k1 = ind.low in
      let k3 = ind.high in
      if i=k1 && j=k3 then ind.isum else
      let m = l.indexes.high in
      if j <= m then query_aux l a i j else
      if i >= m then query_aux r a i j else
      query_aux l a i m + query_aux r a m j
    end


  let query (t:tree) (ghost a: array int) (i j:int) : int
    requires { 0 <= i <= j <= a.length }
    requires { is_tree_for t a 0 a.length }
    ensures { result = sum a i j }
  = if i=j then 0 else query_aux t a i j


  (** frame lemma for predicate `is_tree_for` *)
  lemma is_tree_for_frame : forall t:tree, a:array int, k v i j:int.
    0 <= k < a.length ->
    k < i \/ k >= j ->
    is_tree_for t a i j ->
    is_tree_for t a[k<-v] i j

(** {3 update cumulative tree} *)


  let rec update_aux
      (t:tree) (i:int) (ghost a :array int) (v:int) : (t': tree, delta: int)
    requires { is_tree_for t a t.indexes.low t.indexes.high }
    requires { t.indexes.low <= i < t.indexes.high }
    variant { t }
    ensures {
        delta = v - a[i] /\
        t'.indexes.low = t.indexes.low /\
        t'.indexes.high = t.indexes.high /\
        is_tree_for t' a[i<-v] t'.indexes.low t'.indexes.high }
  = match t with
    | Leaf ind ->
        assert { i = ind.low };
        (Leaf { ind with isum = v }, v - ind.isum)
    | Node ind l r ->
        let m = l.indexes.high in
      if i < m then
        let l',delta = update_aux l i a v in
        assert { is_tree_for l' a[i<-v] t.indexes.low m };
        assert { is_tree_for r a[i<-v] m t.indexes.high };
        (Node {ind with isum = ind.isum + delta } l' r, delta)
      else
        let r',delta = update_aux r i a v in
        assert { is_tree_for l a[i<-v] t.indexes.low m };
        assert { is_tree_for r' a[i<-v] m t.indexes.high };
        (Node {ind with isum = ind.isum + delta} l r',delta)
    end

  let update (t:tree) (ghost a:array int) (i v:int) : tree
     requires { 0 <= i < a.length }
     requires { is_tree_for t a 0 a.length }
     writes { a }
     ensures { a[i] = v }
     ensures { forall k. 0 <= k < a.length /\ k <> i -> a[k] = (old a)[k] }
     ensures { is_tree_for result a 0 a.length }
  = let t,_ = update_aux t i a v in
    assert { is_tree_for t a[i <- v] 0 a.length };
    a[i] <- v;
    t


(** {2 complexity analysis}

  We would like to prove that `query` is really logarithmic. This is
  non-trivial because there are two recursive calls in some cases.

  So far, we are only able to prove that `update` is logarithmic

  We express the complexity by passing a "credit" as a ghost
  parameter. We pose the precondition that the credit is at least
  equal to the depth of the tree.

*)

  (** preliminaries: definition of the depth of a tree, and showing
      that it is indeed logarithmic in function of the number of its
      elements *)

  use int.MinMax

  function depth (t:tree) : int =
    match t with
    | Leaf _ -> 1
    | Node _ l r -> 1 + max (depth l) (depth r)
    end

  lemma depth_min : forall t. depth t >= 1

  use bv.Pow2int

  let rec lemma depth_is_log (t:tree) (a :array int) (k:int)
     requires { k >= 0 }
     requires { is_tree_for t a t.indexes.low t.indexes.high }
     requires { t.indexes.high - t.indexes.low <= pow2 k }
     variant { t }
     ensures { depth t <= k+1 }
  = match t with
    | Leaf _ -> ()
    | Node _ l r ->
       depth_is_log l a (k-1);
       depth_is_log r a (k-1)
    end


  (** `update_aux` function instrumented with a credit *)

  use ref.Ref

  let rec update_aux_complexity
        (t:tree) (i:int) (ghost a :array int)
        (v:int) (ghost c:ref int) : (t': tree, delta: int)
     requires { is_tree_for t a t.indexes.low t.indexes.high }
     requires { t.indexes.low <= i < t.indexes.high }
     variant { t }
     ensures { !c - old !c <= depth t }
     ensures {
        delta = v - a[i] /\
        t'.indexes.low = t.indexes.low /\
        t'.indexes.high = t.indexes.high /\
        is_tree_for t' a[i<-v] t'.indexes.low t'.indexes.high }
  = c := !c + 1;
    match t with
    | Leaf ind ->
      assert { i = ind.low };
      (Leaf { ind with isum = v }, v - ind.isum)
    | Node ind l r ->
      let m = l.indexes.high in
      if i < m then
        let l',delta = update_aux_complexity l i a v c in
        assert { is_tree_for l' a[i<-v] t.indexes.low m };
        assert { is_tree_for r a[i<-v] m t.indexes.high };
        (Node {ind with isum = ind.isum + delta } l' r, delta)
      else
        let r',delta = update_aux_complexity r i a v c in
        assert { is_tree_for l a[i<-v] t.indexes.low m };
        assert { is_tree_for r' a[i<-v] m t.indexes.high };
        (Node {ind with isum = ind.isum + delta} l r',delta) (*>*)
    end

  (** `query_aux` function instrumented with a credit *)

  let rec query_aux_complexity (t:tree) (ghost a: array int)
      (i j:int) (ghost c:ref int) : int
    requires { is_tree_for t a t.indexes.low t.indexes.high }
    requires { 0 <= t.indexes.low <= i < j <= t.indexes.high <= a.length }
    variant { t }
    ensures { !c - old !c <=
         if i = t.indexes.low /\ j = t.indexes.high then 1 else
         if i = t.indexes.low \/ j = t.indexes.high then 2 * depth t else
          4 * depth t }
    ensures { result = sum a i j }
  = c := !c + 1;
    match t with
    | Leaf ind ->
      ind.isum
    | Node ind l r ->
      let k1 = ind.low in
      let k3 = ind.high in
      if i=k1 && j=k3 then ind.isum else
      let m = l.indexes.high in
      if j <= m then query_aux_complexity l a i j c else
      if i >= m then query_aux_complexity r a i j c else
      query_aux_complexity l a i m c + query_aux_complexity r a m j c
    end

end