
module Bag

  use int.Int

  type bag 'a = 'a -> int

  predicate (==) (b1 b2: bag 'a) =
    forall x: 'a. b1 x = b2 x

  let constant empty : bag 'a =
    fun _ -> 0

  let ghost function add (e: 'a) (b: bag 'a): bag 'a =
    fun x -> if pure {x = e} then b x + 1 else b x

  let ghost function remove (e: 'a) (b: bag 'a): bag 'a =
    fun x -> if pure {x = e} then b x - 1 else b x

end

module BagSpec

  use int.Int
  use Bag

  type t 'a = private {
          mutable size: int;
    ghost mutable contents: bag 'a;
  } invariant { 0 <= size }

  val create (_dummy: 'a) : t 'a
    ensures { result.size = 0 }
    ensures { result.contents == Bag.empty }

  val length (t: t 'a) : int
    ensures { result = t.size }

  val clear (t: t 'a) : unit
    writes  { t.size, t.contents }
    ensures { t.size = 0 }
    ensures { t.contents == Bag.empty }

  val add (t: t 'a) (x: 'a) : unit
    writes  { t.size, t.contents }
    ensures { t.size = old t.size + 1 }
    ensures { t.contents == Bag.add x (old t.contents) }

  val get (t: t 'a) (i: int) : 'a
    requires { 0 <= i < t.size }
    ensures  { t.contents result > 0 }

  val remove (t: t 'a) (i: int) : unit
    requires { 0 <= i < t.size }
    writes   { t.size, t.contents }
    ensures  { t.size = old t.size - 1 }
    ensures  { exists x: 'a. t.contents == Bag.remove x (old t.contents) }

end

module ResizableArraySpec

  use int.Int
  use map.Map
  use map.Const

  type rarray 'a = private {
    ghost mutable len: int;
    ghost mutable data: map int 'a;
  } invariant { 0 <= len }

  function ([]) (r: rarray 'a) (i: int) : 'a = Map.get r.data i

  val function ([<-]) (r: rarray 'a) (i: int) (v: 'a) : rarray 'a
    ensures { result.len = r.len }
    ensures { result.data = Map.set r.data i v }

  val length (r: rarray 'a) : int
    ensures { result = r.len }

  val make (len: int) (dummy: 'a) : rarray 'a
    requires { 0 <= len }
    ensures  { result.len = len }
    ensures  { result.data = Const.const dummy }

  val ([]) (r: rarray 'a) (i: int) : 'a
    requires { 0 <= i < r.len }
    ensures  { result = r[i] }

  val ([]<-) (r: rarray 'a) (i: int) (v: 'a) : unit
    requires { 0 <= i < r.len }
    writes   { r.data }
    ensures  { r = (old r)[i <- v] }

  val resize (r: rarray 'a) (len: int) : unit
    requires { 0 <= len }
    writes   { r.len, r.data }
    ensures  { r.len = len }
    ensures  { forall i: int.
               0 <= i < old r.len -> i < len -> r[i] = (old r)[i] }

end

module BagImpl

  use int.Int
  use Bag
  use import ResizableArraySpec as R
  use map.Map as Map
  use int.NumOf as NumOf

  function numof (r: rarray 'a) (x: 'a) (l u: int) : int =
    NumOf.numof (fun i -> Map.get r.R.data i = x) l u

  type t 'a = {
    mutable       size:     int;
                  data:     rarray 'a;
    mutable ghost contents: bag 'a;
  }
    invariant { 0 <= size = data.len }
    invariant { forall x: 'a. contents x = numof data x 0 size }
    by { size = 0; data = R.make 0 (any 'a); contents = empty }

  let create (dummy: 'a) : t 'a
    ensures { result.size = 0 }
    ensures { result.contents == Bag.empty }
  =
    { size = 0; data = make 0 dummy; contents = Bag.empty }

  let length_ (t: t 'a) : int
    ensures { result = t.size }
  = t.size

  let clear (t: t 'a) : unit
    ensures { t.size = 0 }
    ensures { t.contents == Bag.empty }
  =
    resize t.data 0;
    t.size <- 0;
    t.contents <- Bag.empty

  let add (t: t 'a) (x: 'a) : unit
    ensures { t.size = old t.size + 1 }
    ensures { t.contents == Bag.add x (old t.contents) }
  =
    let n = t.size in
    t.size <- n + 1;
    resize t.data (n + 1);
    assert { forall v: 'a. t.contents v = numof t.data v 0 n by
               forall i. 0 <= i < n -> t.data[i] = (old t.data)[i] };
    t.data[n] <- x;
    assert { forall v: 'a. t.contents v = numof t.data v 0 n };
    t.contents <- Bag.add x t.contents

  let get (t: t 'a) (i: int) : 'a
    requires { 0 <= i < t.size }
    ensures  { result = t.data[i] }
  = t.data[i]

  let remove (t: t 'a) (i: int) : unit
    requires { 0 <= i < t.size }
    ensures  { t.size = old t.size - 1 }
    ensures  { forall x: 'a. x = old t.data[i] ->
               t.contents == Bag.remove x (old t.contents) }
  =
    let n = t.size - 1 in
    let ghost x = t.data[i] in
    t.contents <- Bag.remove x t.contents;
    if i < n then t.data[i] <- t.data[n];
    assert { forall v: 'a. t.contents v = numof t.data v 0 n by
      numof t.data v 0 n = numof (old t.data) v 0 i +
                           numof (old t.data) v (i+1) (n+1) by
        numof t.data v 0 i = numof (old t.data) v 0 i /\
        numof t.data v i n = numof (old t.data) v (i+1) (n+1) by
          numof t.data v (i+1) n = numof (old t.data) v (i+1) n };
    resize t.data n;
    t.size <- n

  clone BagSpec with
    type t = t,
    val create = create,
    val length = length_,
    val clear = clear,
    val add = add,
    val get = get,
    val remove = remove

end

module Harness

  use int.Int
  use Bag
  use BagImpl

  let test1 () =
    let b = create 0 in
    add b 17;
    add b 42;
    assert { b.contents 42 = 1 };
    add b 55;
    add b 89;
    add b 42;
    assert { b.contents 42 = 2 };
    ()

end
