module SparseArray

(*
      If the sparse array contains three elements x y z, at index
      a b c respectively, then the three arrays look like this:

              b     a      c
values +-----+-+---+-+----+-+----+
       |     |y|   |x|    |z|    |
       +-----+-+---+-+----+-+----+

index  +-----+-+---+-+----+-+----+
       |     |1|   |0|    |2|    |
       +-----+-+---+-+----+-+----+

        0 1 2  n=3
back   +-+-+-+-------------------+
       |a|b|c|                   |
       +-+-+-+-------------------+

*)

  use int.Int
  use array.Array

  constant maxlen : int = 1000

  type sparse_array 'a = { values : array 'a;
                           index  : array int;
                           back   : array int;
                   mutable card   : int;
                           def    : 'a; }
  invariant {
    0 <= card <= length values <= maxlen /\
    length values = length index = length back /\
    forall i : int.
      0 <= i < card ->
      0 <= back[i] < length values /\ index[back[i]] = i
  } by {
    values = make 0 (any 'a);
    index  = make 0 0;
    back   = make 0 0;
    card   = 0;
    def    = any 'a
  }

  predicate is_elt (a: sparse_array 'a) (i: int) =
    0 <= a.index[i] < a.card /\ a.back[a.index[i]] = i

  function value (a: sparse_array 'a) (i: int) : 'a =
    if is_elt a i then
      a.values[i]
    else
      a.def

  function length (a: sparse_array 'a) : int = Array.length a.values

  (* creation *)

  val malloc (n:int) : array 'a ensures { Array.length result = n }

  let create (sz: int) (d: 'a)
    requires { 0 <= sz <= maxlen }
    ensures { result.card = 0 /\ result.def = d /\ length result = sz }
  = { values = malloc sz;
      index  = malloc sz;
      back   = malloc sz;
      card   = 0;
      def    = d }

  (* access *)

  let test (a: sparse_array 'a) i
    requires { 0 <= i < length a }
    ensures { result=True <-> is_elt a i }
  = 0 <= a.index[i] < a.card && a.back[a.index[i]] = i

  let get (a: sparse_array 'a) i
    requires { 0 <= i < length a }
    ensures { result = value a i }
  = if test a i then
      a.values[i]
    else
      a.def

  (* assignment *)

  use map.MapInjection as MI

  lemma permutation :
    forall a: sparse_array 'a.
    (* sparse_array invariant *)
      Array.(0 <= a.card <= Array.length a.values <= maxlen /\
      length a.values = length a.index = length a.back /\
      forall i : int.
        0 <= i < a.card ->
        0 <= a.back[i] < length a.values /\ a.index[a.back[i]] = i) ->
    (* sparse_array invariant *)
    a.card = a.length ->
    forall i: int. 0 <= i < a.length -> is_elt a i
      by MI.surjective a.back.elts a.card
      so exists j. 0 <= j < a.card /\ a.back[j] = i


  let set (a: sparse_array 'a) i v
    requires { 0 <= i < length a }
    ensures { value a i = v /\
      forall j:int. j <> i -> value a j = value (old a) j }
  = a.values[i] <- v;
    if not (test a i) then begin
      assert { a.card < length a };
      a.index[i] <- a.card;
      a.back[a.card] <- i;
      a.card <- a.card + 1
    end

end

module Harness

  use SparseArray

  type elt
  val constant default : elt

  val constant c1 : elt
  val constant c2 : elt

  let harness () =
    let a = create 10 default in
    let b = create 20 default in
    let get_a_5 = get a 5 in assert { get_a_5 = default };
    let get_b_7 = get b 7 in assert { get_b_7 = default };
    set a 5 c1;
    set b 7 c2;
    let get_a_5 = get a 5 in assert { get_a_5 = c1 };
    let get_b_7 = get b 7 in assert { get_b_7 = c2 };
    let get_a_7 = get a 7 in assert { get_a_7 = default };
    let get_b_5 = get b 5 in assert { get_b_5 = default };
    let get_a_0 = get a 0 in assert { get_a_0 = default };
    let get_b_0 = get b 0 in assert { get_b_0 = default };
    ()

  val predicate (!=) (x y : elt)
    ensures { result <-> x <> y }

  exception BenchFailure

  let bench () raises { BenchFailure -> true } =
    let a = create 10 default in
    let b = create 20 default in
    if get a 5 != default then raise BenchFailure;
    if get b 7 != default then raise BenchFailure;
    set a 5 c1;
    set b 7 c2;
    if get a 5 != c1 then raise BenchFailure;
    if get b 7 != c2 then raise BenchFailure;
    if get a 7 != default then raise BenchFailure;
    if get b 5 != default then raise BenchFailure;
    if get a 0 != default then raise BenchFailure;
    if get b 0 != default then raise BenchFailure

end