(** {1 Arrays} *)

(** {2 Generic Arrays}

The length is a non-mutable field, so that we get for free that
modification of an array does not modify its length.

*)

module Array

  use int.Int
  use map.Map

  type array [@extraction:array] 'a = private {
    mutable ghost elts : int -> 'a;
                length : int
  } invariant { 0 <= length }

  function ([]) (a: array 'a) (i: int) : 'a = a.elts i

  val ([]) (a: array 'a) (i: int) : 'a
    requires { [@expl:index in array bounds] 0 <= i < length a }
    ensures  { result = a[i] }

  val ghost function ([<-]) (a: array 'a) (i: int) (v: 'a): array 'a
    ensures { result.length = a.length }
    ensures { result.elts = Map.set a.elts i v }

  val ([]<-) (a: array 'a) (i: int) (v: 'a) : unit writes {a}
    requires { [@expl:index in array bounds] 0 <= i < length a }
    ensures  { a.elts = Map.set (old a).elts i v }
    ensures  { a = (old a)[i <- v] }

  (** unsafe get/set operations with no precondition *)

  exception OutOfBounds

  let defensive_get (a: array 'a) (i: int)
    ensures { 0 <= i < length a /\ result = a[i] }
    raises  { OutOfBounds -> i < 0 \/ i >= length a }
  = if i < 0 || i >= length a then raise OutOfBounds;
    a[i]

  let defensive_set (a: array 'a) (i: int) (v: 'a)
    ensures { 0 <= i < length a }
    ensures { a = (old a)[i <- v] }
    raises  { OutOfBounds -> (i < 0 \/ i >= length a) /\ a = old a }
  = if i < 0 || i >= length a then raise OutOfBounds;
    a[i] <- v

  function make (n: int) (v: 'a) : array 'a

  axiom make_spec : forall n:int, v:'a.
    n >= 0 ->
    (forall i:int. 0 <= i < n -> (make n v)[i] = v) /\
    length (make n v) = n

  val make [@extraction:array_make] (n: int) (v: 'a) : array 'a
    requires { [@expl:array creation size] n >= 0 }
    ensures { forall i:int. 0 <= i < n -> result[i] = v }
    ensures { result.length = n }

  val empty () : array 'a
    ensures { result.length = 0 }

  let copy (a: array 'a) : array 'a
    ensures  { length result = length a }
    ensures  { forall i:int. 0 <= i < length result -> result[i] = a[i] }
  =
    let len = length a in
    if len = 0 then empty ()
    else begin
      let b = make len a[0] in
      for i = 1 to len - 1 do
        invariant { forall k. 0 <= k < i -> b[k] = a[k] }
        b[i] <- a[i]
      done;
      b
    end

  let sub (a: array 'a) (ofs: int) (len: int) : array 'a
    requires { 0 <= ofs /\ 0 <= len /\ ofs + len <= length a }
    ensures  { length result = len }
    ensures  { forall i:int. 0 <= i < len -> result[i] = a[ofs + i] }
  =
    if length a = 0 then begin
     assert { len = 0 };
     empty ()
    end else begin
      let b = make len a[0] in
      for i = 0 to len-1 do
        invariant { forall k. 0 <= k < i -> b[k] = a[ofs+k] }
        b[i] <- a[ofs+i];
      done;
      b
    end

  let fill (a: array 'a) (ofs: int) (len: int) (v: 'a)
    requires { 0 <= ofs /\ 0 <= len /\ ofs + len <= length a }
    ensures  { forall i:int.
      (0 <= i < ofs \/ ofs + len <= i < length a) -> a[i] = old a[i] }
    ensures  { forall i:int. ofs <= i < ofs + len -> a[i] = v }
  =
    for k = 0 to len - 1 do
      invariant { forall i:int.
        (0 <= i < ofs \/ ofs + len <= i < length a) -> a[i] = old a[i] }
      invariant { forall i:int. ofs <= i < ofs + k -> a[i] = v }
      a[ofs + k] <- v
    done

  let blit (a1: array 'a) (ofs1: int)
                 (a2: array 'a) (ofs2: int) (len: int) : unit writes {a2}
    requires { 0 <= ofs1 /\ 0 <= len /\ ofs1 + len <= length a1 }
    requires { 0 <= ofs2 /\             ofs2 + len <= length a2 }
    ensures  { forall i:int.
      (0 <= i < ofs2 \/ ofs2 + len <= i < length a2) -> a2[i] = old a2[i] }
    ensures  { forall i:int.
      ofs2 <= i < ofs2 + len -> a2[i] = a1[ofs1 + i - ofs2] }
  =
    for i = 0 to len - 1 do
      invariant { forall k. not (0 <= k < i) -> a2[ofs2 + k] = old a2[ofs2 + k] }
      invariant { forall k. 0 <= k < i -> a2[ofs2 + k] = a1[ofs1 + k] }
      a2[ofs2 + i] <- a1[ofs1 + i];
    done

  let append (a1: array 'a) (a2: array 'a) : array 'a
    ensures { length result = length a1 + length a2 }
    ensures { forall i. 0 <= i < length a1 -> result[i] = a1[i] }
    ensures { forall i. 0 <= i < length a2 -> result[length a1 + i] = a2[i] }
  =
    if length a1 = 0 then copy a2
    else begin
      let a = make (length a1 + length a2) a1[0] in
      blit a1 0 a 0 (length a1);
      blit a2 0 a (length a1) (length a2);
      a
    end

  let self_blit (a: array 'a) (ofs1: int) (ofs2: int) (len: int) : unit
    writes {a}
    requires { 0 <= ofs1 /\ 0 <= len /\ ofs1 + len <= length a }
    requires { 0 <= ofs2 /\             ofs2 + len <= length a }
    ensures  { forall i:int.
      (0 <= i < ofs2 \/ ofs2 + len <= i < length a) -> a[i] = old a[i] }
    ensures  { forall i:int.
      ofs2 <= i < ofs2 + len -> a[i] = old a[ofs1 + i - ofs2] }
  =
    if ofs1 <= ofs2 then (* from right to left *)
      for k = len - 1 downto 0 do
        invariant  { forall i:int.
          (0 <= i <= ofs2 + k \/ ofs2 + len <= i < length a) ->
          a[i] = (old a)[i] }
        invariant  { forall i:int.
          ofs2 + k < i < ofs2 + len -> a[i] = (old a)[ofs1 + i - ofs2] }
        a[ofs2 + k] <- a[ofs1 + k]
      done
    else (* from left to right *)
      for k = 0 to len - 1 do
        invariant  { forall i:int.
          (0 <= i < ofs2 \/ ofs2 + k <= i < length a) ->
          a[i] = (old a)[i] }
        invariant  { forall i:int.
          ofs2 <= i < ofs2 + k -> a[i] = (old a)[ofs1 + i - ofs2] }
        a[ofs2 + k] <- a[ofs1 + k]
      done

  (*** TODO?
     - concat : 'a array list -> 'a array
     - to_list
     - of_list
  *)

end

module Init

  use int.Int
  use export Array

  let init (n: int) (f: int -> 'a) : array 'a
    requires { [@expl:array creation size] n >= 0 }
    ensures { forall i:int. 0 <= i < n -> result[i] = f i }
    ensures { result.length = n }
  =
    if n = 0 then empty ()
    else begin
      let a = make n (f 0) in
      for i = 1 to n - 1 do
        invariant { forall k. 0 <= k < i -> a[k] = f k }
        a[i] <- f i
      done;
      a
    end


end

(** {2 Sorted Arrays} *)

module IntArraySorted

  use int.Int
  use Array
  clone map.MapSorted as M with type elt = int, predicate le = (<=)

  predicate sorted_sub (a : array int) (l u : int) =
    M.sorted_sub a.elts l u
  (** `sorted_sub a l u` is true whenever the array segment `a(l..u-1)`
      is sorted w.r.t order relation `le` *)

  predicate sorted (a : array int) =
    M.sorted_sub a.elts 0 a.length
  (** `sorted a` is true whenever the array `a` is sorted w.r.t `le` *)

end

module Sorted

  use int.Int
  use Array

  type elt

  predicate le elt elt

  predicate sorted_sub (a: array elt) (l u: int) =
    forall i1 i2 : int. l <= i1 < i2 < u -> le a[i1] a[i2]
  (** `sorted_sub a l u` is true whenever the array segment `a(l..u-1)`
      is sorted w.r.t order relation `le` *)

  predicate sorted (a: array elt) =
    forall i1 i2 : int. 0 <= i1 < i2 < length a -> le a[i1] a[i2]
  (** `sorted a` is true whenever the array `a` is sorted w.r.t `le` *)

end

(** {2 Arrays Equality} *)

module ArrayEq

  use int.Int
  use Array
  use map.MapEq

  predicate array_eq_sub (a1 a2: array 'a) (l u: int) =
    a1.length = a2.length /\ 0 <= l <= a1.length /\ 0 <= u <= a1.length /\
    map_eq_sub a1.elts a2.elts l u

  predicate array_eq (a1 a2: array 'a) =
    a1.length = a2.length /\ map_eq_sub a1.elts a2.elts 0 (length a1)

end

module ArrayExchange

  use int.Int
  use Array
  use map.MapExchange as M

  predicate exchange (a1 a2: array 'a) (i j: int) =
    a1.length = a2.length /\
    M.exchange a1.elts a2.elts 0 a1.length i j
  (** `exchange a1 a2 i j` means that arrays `a1` and `a2` only differ
      by the swapping of elements at indices `i` and `j` *)

end

(** {2 Permutation} *)

module ArrayPermut

  use int.Int
  use Array
  use map.MapPermut as M
  use map.MapEq
  use ArrayEq
  use export ArrayExchange

  predicate permut (a1 a2: array 'a) (l u: int) =
    a1.length = a2.length /\ 0 <= l <= a1.length /\ 0 <= u <= a1.length /\
    M.permut a1.elts a2.elts l u
  (** `permut a1 a2 l u` is true when the segment
      `a1(l..u-1)` is a permutation of the segment `a2(l..u-1)`.
      Values outside of the interval `(l..u-1)` are ignored. *)

  predicate permut_sub (a1 a2: array 'a) (l u: int) =
    map_eq_sub a1.elts a2.elts 0 l /\
    permut a1 a2 l u /\
    map_eq_sub a1.elts a2.elts u (length a1)
  (** `permut_sub a1 a2 l u` is true when the segment
      `a1(l..u-1)` is a permutation of the segment `a2(l..u-1)`
      and values outside of the interval `(l..u-1)` are equal. *)

  predicate permut_all (a1 a2: array 'a) =
    a1.length = a2.length /\ M.permut a1.elts a2.elts 0 a1.length
  (** `permut_all a1 a2 l u` is true when array `a1` is a permutation
      of array `a2`. *)

  lemma exchange_permut_sub:
    forall a1 a2: array 'a, i j l u: int.
    exchange a1 a2 i j -> l <= i < u -> l <= j < u ->
    0 <= l -> u <= length a1 -> permut_sub a1 a2 l u

  lemma permut_sub_trans:
    forall a1 a2 a3: array 'a, l u: int.
    0 <= l -> u <= length a1 -> permut_sub a1 a2 l u ->
    permut_sub a2 a3 l u -> permut_sub a1 a3 l u

  (** we can always enlarge the interval *)
  lemma permut_sub_weakening:
    forall a1 a2: array 'a, l1 u1 l2 u2: int.
    permut_sub a1 a2 l1 u1 -> 0 <= l2 <= l1 -> u1 <= u2 <= length a1 ->
    permut_sub a1 a2 l2 u2

  lemma exchange_permut_all:
    forall a1 a2: array 'a, i j: int.
    exchange a1 a2 i j -> permut_all a1 a2

end

module ArraySwap

  use int.Int
  use Array
  use export ArrayExchange

  let swap (a:array 'a) (i:int) (j:int) : unit
    requires { 0 <= i < length a /\ 0 <= j < length a }
    writes   { a }
    ensures  { exchange (old a) a i j }
  = let v = a[i] in
    a[i] <- a[j];
    a[j] <- v

end

(** {2 Sum of elements} *)

module ArraySum

  use Array
  use int.Sum as S

  (** `sum a l h` is the sum of `a[i]` for `l <= i < h` *)
  function sum (a: array int) (l h: int) : int = S.sum a.elts l h

end

(** {2 Number of array elements satisfying a given predicate} *)

module NumOf
  use Array
  use int.NumOf as N

  (** the number of `a[i]` such that `l <= i < u` and `pr i a[i]` *)
  function numof (pr: int -> 'a -> bool) (a: array 'a) (l u: int) : int =
    N.numof (fun i -> pr i a[i]) l u

end

module NumOfEq
  use Array
  use int.NumOf as N

  (** the number of `a[i]` such that `l <= i < u` and `a[i] = v` *)
  function numof (a: array 'a) (v: 'a) (l u: int) : int =
    N.numof (fun i -> a[i] = v) l u

end

module ToList
  use int.Int
  use Array
  use list.List

  let rec function to_list (a: array 'a) (l u: int) : list 'a
    requires { l >= 0 /\ u <= a.length }
    variant  { u - l }
  = if u <= l then Nil else Cons a[l] (to_list a (l+1) u)

  use list.Append

  let rec lemma to_list_append (a: array 'a) (l m u: int)
    requires { 0 <= l <= m <= u <= a.length }
    variant  { m - l }
    ensures  { to_list a l m ++ to_list a m u = to_list a l u }
  = if l < m then to_list_append a (l+1) m u

end

module ToSeq
  use int.Int
  use Array
  use seq.Seq as S

  let rec function to_seq_sub (a: array 'a) (l u: int) : S.seq 'a
    requires { l >= 0 /\ u <= a.length }
    variant { u - l }
  = if u <= l then S.empty else S.cons a[l] (to_seq_sub a (l+1) u)

  let rec lemma to_seq_length (a: array 'a) (l u: int)
    requires { 0 <= l <= u <= length a }
    variant  { u - l }
    ensures  { S.length (to_seq_sub a l u) = u - l }
  = if l < u then to_seq_length a (l+1) u

  let rec lemma to_seq_nth (a: array 'a) (l i u: int)
    requires { 0 <= l <= i < u <= length a }
    variant  { i - l }
    ensures  { S.get (to_seq_sub a l u) (i - l) = a[i] }
  = if l < i then to_seq_nth a (l+1) i u

  let function to_seq (a: array 'a) : S.seq 'a = to_seq_sub a 0 (length a)
  meta coercion function to_seq

end

(** {2 Number of inversions in an array of integers}

    We show that swapping two elements that are ill-sorted decreases
    the number of inversions. Useful to prove the termination of
    sorting algorithms that use swaps. *)

module Inversions

  use Array
  use ArrayExchange
  use int.Int
  use int.Sum
  use int.NumOf

  (* to prove termination, we count the total number of inversions *)
  predicate inversion (a: array int) (i j: int) =
    a[i] > a[j]

  function inversions_for (a: array int) (i: int) : int =
    numof (inversion a i) i (length a)

  function inversions (a: array int) : int =
    sum (inversions_for a) 0 (length a)

  (* the key lemma to prove termination: whenever we swap two consecutive
     values that are ill-sorted, the total number of inversions decreases *)
  let lemma exchange_inversion (a1 a2: array int) (i0: int)
    requires { 0 <= i0 < length a1 - 1 }
    requires { a1[i0] > a1[i0 + 1] }
    requires { exchange a1 a2 i0 (i0 + 1) }
    ensures  { inversions a2 < inversions a1 }
  = assert { inversion a1 i0 (i0+1) };
    assert { not (inversion a2 i0 (i0+1)) };
    assert { forall i. 0 <= i < i0 ->
             inversions_for a2 i = inversions_for a1 i
             by numof (inversion a2 i) i (length a2)
              = numof (inversion a2 i) i i0
              + numof (inversion a2 i) i0 (i0+1)
              + numof (inversion a2 i) (i0+1) (i0+2)
              + numof (inversion a2 i) (i0+2) (length a2)
             /\ numof (inversion a1 i) i (length a1)
              = numof (inversion a1 i) i i0
              + numof (inversion a1 i) i0 (i0+1)
              + numof (inversion a1 i) (i0+1) (i0+2)
              + numof (inversion a1 i) (i0+2) (length a1)
             /\ numof (inversion a2 i) i0 (i0+1)
                = numof (inversion a1 i) (i0+1) (i0+2)
             /\ numof (inversion a2 i) (i0+1) (i0+2)
                = numof (inversion a1 i) i0 (i0+1)
             /\ numof (inversion a2 i) i i0 = numof (inversion a1 i) i i0
             /\ numof (inversion a2 i) (i0+2) (length a2)
                = numof (inversion a1 i) (i0+2) (length a1)
              };
    assert { forall i. i0 + 1 < i < length a1 ->
             inversions_for a2 i = inversions_for a1 i };
    assert { inversions_for a2 i0 = inversions_for a1 (i0+1)
             by numof (inversion a1 (i0+1)) (i0+2) (length a1)
              = numof (inversion a2 i0    ) (i0+2) (length a1) };
    assert { 1 + inversions_for a2 (i0+1) = inversions_for a1 i0
             by numof (inversion a1 i0) i0 (length a1)
              = numof (inversion a1 i0) (i0+1) (length a1)
              = 1 + numof (inversion a1 i0) (i0+2) (length a1)
              = 1 + numof (inversion a2 (i0+1)) (i0+2) (length a2) };
    let sum_decomp (a: array int) (i j k: int)
      requires { 0 <= i <= j <= k <= length a = length a1 }
      ensures  { sum (inversions_for a) i k =
                 sum (inversions_for a) i j + sum (inversions_for a) j k }
    = () in
    let decomp (a: array int)
      requires { length a = length a1 }
      ensures  { inversions a = sum (inversions_for a) 0 i0
                              + inversions_for a i0
                              + inversions_for a (i0+1)
                              + sum (inversions_for a) (i0+2) (length a) }
    = sum_decomp a 0 i0 (length a);
      sum_decomp a i0 (i0+1) (length a);
      sum_decomp a (i0+1) (i0+2) (length a);
    in
    decomp a1; decomp a2;
    ()

end