
(** {1 Implementation of logarithmic-time ordered associative tables on top
       of AVL trees}

    Author: Martin Clochard *)


(** {2 Shared base implementation between set and map structures} *)
module MapBase

  use int.Int
  use avl.SelectionTypes
  use option.Option
  use ref.Ref
  use seq.Seq
  use mach.peano.Peano as I

  (** {3 Implementation parameters *)

  (** The level of balancing is left abstract. *)
  val constant balancing : I.t
    ensures { result > 0 }

  (** Stored elements are identified by totally ordered keys *)
  scope D type t 'a end scope K type t end
  clone export key_type.KeyType with type t = D.t, type key = K.t
  clone preorder.Computable as CO with type t = K.t, axiom .
  scope D
    let function measure 'a : unit = ()
  end

  (** As we do not need any extra information in order to perform search
      for this particular instance, we instantiate the monoid by unit. *)
  scope M
    type t = unit
    constant zero : unit = ()
    function op (_x _y:unit) : unit = ()
    let lemma neutral_ (x:unit) : unit
      ensures { op zero x = x = op x zero }
    = match x with _ -> () end
    clone export monoid.Monoid with type t = t,
      constant zero = zero,function op = op,lemma assoc,lemma neutral
    clone export monoid.MonoidSumDef with type M.t = t,
      constant M.zero = zero,function M.op = op,goal M.assoc,goal M.neutral
    let zero () : unit ensures { result = () } = ()
    let op (_x _y:unit) : unit ensures { result = () } = ()
  end

  (** In associative tables, elements are selected
      with respect to their keys. *)
  type selector = K.t

  (** Efficient search for ordered keys can be carried out efficiently
      only in strictly increasing sequences. *)
  predicate selection_possible (_:'b) (s:seq (D.t 'a)) =
    forall i j. 0 <= i < j < length s -> CO.lt s[i].key s[j].key

  predicate upper_bound_s (k:K.t) (s:seq (D.t 'a)) =
    forall i. 0 <= i < length s -> CO.lt s[i].key k
  predicate lower_bound_s (k:K.t) (s:seq (D.t 'a)) =
    forall i. 0 <= i < length s -> CO.lt k s[i].key

  (** Selected splits correspond to sorted sequence, in which:

      1) the left part precede the selector key

      2) the selector key precede the right part

      3) If it exists, the middle element has key equivalent to the selector *)
  predicate selected (k:K.t) (e:split (D.t 'a)) =
    let _s = rebuild e in
    upper_bound_s k e.left /\ lower_bound_s k e.right /\
    match e.middle with None -> true | Some d -> CO.eq k d.key end

  (** Comparison-based binary search *)
  let selected_part (ghost lseq rseq:seq (D.t 'a))
    (k:K.t) (_l:'e) (d:D.t 'a) (_r:'f) : part_base K.t
    requires { selection_possible k (lseq ++ cons d rseq) }
    returns { Here -> let e2 = { left = lseq;
        middle = Some d;
        right = rseq } in selected k e2
      | Left sl -> selection_possible sl lseq /\
        forall e. selected sl e /\ rebuild e = lseq ->
          selected k (right_extend e d rseq)
      | Right sr -> selection_possible sr rseq /\
        forall e. selected sr e /\ rebuild e = rseq ->
          selected k (left_extend lseq d e) }
  = let kd = d.key in
    let ghost s = (lseq ++ cons d rseq) in
    assert { forall i. 0 <= i < length rseq -> rseq[i] = s[i+length lseq+1] };
    assert { d = s[length lseq] };
    let cmp = CO.compare k kd in
    if cmp < 0
    then Left k
    else if cmp > 0
    then Right k
    else Here

  (** Full instantiation of the avl module. *)
  clone avl.AVL as Sel with type selector = selector,
    predicate selection_possible = selection_possible,
    predicate selected = selected,
    val selected_part = selected_part,
    goal selection_empty,
    val balancing = balancing,
    type D.t = D.t,
    val D.measure = D.measure,
    type M.t = M.t,
    constant M.zero = M.zero,
    function M.op = M.op,
    goal M.assoc,
    goal M.neutral,
    function M.agg = M.agg,
    goal M.agg_empty,
    goal M.agg_sing,
    goal M.agg_cat,
    val M.zero = M.zero,
    val M.op = M.op

  (** {3 Adaptation of the specification to associative tables} *)

  (** Extra invariant: sequence sorted *)
  type t 'a = { field : Sel.t 'a } invariant { selection_possible () field }
  by { field = Sel.empty () }

  type m 'a = {
    domn : K.t -> bool;
    func : K.t -> D.t 'a;
    card : int;
  }

  predicate domain (s:seq (D.t 'a)) (k:K.t) =
    exists i. 0 <= i < length s /\ CO.eq s[i].key k

  let ghost function make_func (s:seq (D.t 'a)) (k:K.t) : D.t 'a
    requires { selection_possible () s /\ domain s k }
    ensures { forall i. 0 <= i < length s /\ CO.eq s[i].key k -> result = s[i] }
  = let j = ref 0 in
    while not CO.eq s[!j].key k do
      invariant { 0 <= !j < length s }
      invariant { CO.le s[!j].key k }
      variant { length s - !j }
      j := !j + 1
    done;
    s[!j]

  function m (t:t 'a) : m 'a =
    { domn = fun y -> domain t.field y;
      func = fun y -> make_func t.field y;
      card = length t.field }
  meta coercion function m

  let lemma correction (t:t 'a)
    ensures { forall k1 k2. CO.eq k1 k2 -> t.domn k1 <-> t.domn k2 }
    ensures { forall k1 k2. CO.eq k1 k2 /\ t.domn k1 -> t.func k1 = t.func k2 }
    ensures { forall k. t.domn k -> CO.eq k (t.func k).key }
    ensures { t.card >= 0 }
  = ()

  (* Link model and selection. *)
  let lemma selected_model (k0:K.t) (e:split (D.t 'a))
    requires { selection_possible k0 (rebuild e) }
    requires { selected k0 e }
    ensures {
      let s = rebuild e in
      match e.middle with
      | None -> forall k. CO.eq k k0 -> not domain s k
      | Some d -> forall k. CO.eq k k0 -> domain s k /\ make_func s k = d
        by s[length e.left] = d so CO.eq d.key k
      end
      /\ (forall k. CO.le k k0 -> not domain e.right k)
      /\ (forall k. CO.le k0 k -> not domain e.left k)
      /\ (forall k. CO.lt k0 k -> domain e.right k <-> domain s k)
      /\ (forall k. CO.lt k k0 -> domain e.left k <-> domain s k)
      /\ (forall k. CO.lt k0 k /\ domain e.right k ->
        make_func e.right k = make_func s k)
      /\ (forall k. CO.lt k k0 /\ domain e.left k ->
        make_func e.left k = make_func s k)
    }
  = ()

  (* Link selection and sortedness. *)
  let lemma selected_sorted (k0:K.t) (e:split (D.t 'a))
    requires { selected k0 e }
    ensures { selection_possible k0 (rebuild e) <->
      selection_possible k0 e.left /\ selection_possible k0 e.right }
  = ()

  (** Create an empty associative table. *)
  let empty () : t 'a
    ensures { result.card = 0 /\ forall k. not result.domn k }
  = { field = Sel.empty () }

  (** Create an associative table with a single element. *)
  let singleton (d:D.t 'a) : t 'a
    ensures { forall k. result.domn k <-> CO.eq k d.key }
    ensures { forall k. CO.eq k d.key -> result.func k = d }
    ensures { result.card = 1 }
  = { field = Sel.singleton d }

  (** Check emptyness of an associative table. *)
  let is_empty (ghost rk:ref K.t) (t:t 'a) : bool
    ensures { result -> forall k. not t.domn k }
    ensures { not result -> t.domn !rk }
    ensures { result <-> t.card = 0 }
  = let res = Sel.is_empty t.field in
    ghost if not res then rk := t.field.Sel.m.Sel.seq[0].key;
    res

  let min (t:t 'a) : D.t 'a
    requires { t.card > 0 }
    ensures  { t.domn result.key /\ t.func result.key = result }
    ensures  { forall k. t.domn k -> CO.le result.key k }
  = Sel.front t.field

  let max (t:t 'a) : D.t 'a
    requires { t.card > 0 }
    ensures  { t.domn result.key /\ t.func result.key = result }
    ensures  { forall k. t.domn k -> CO.le k result.key }
  = Sel.back t.field

  (** Get and extract the element with minimum key from the table. *)
  let decompose_min (t:t 'a) : option (D.t 'a,t 'a)
    returns { None -> t.card = 0 /\ forall k. not t.domn k
      | Some (d,r:t 'a) -> t.card = r.card + 1 /\
        t.domn d.key /\ t.func d.key = d /\ not r.domn d.key /\
        (forall k. CO.lt k d.key -> not t.domn k) /\
        (forall k. CO.eq k d.key -> t.func k = d) /\
        (forall k. CO.le k d.key -> not r.domn k) /\
        (forall k. CO.lt d.key k -> r.domn k <-> t.domn k) /\
        (forall k. CO.lt d.key k /\ r.domn k -> r.func k = t.func k) }
  = match Sel.decompose_front t.field with
    | None -> None
    | Some (d,r) -> Some (d,{ field = r })
    end

  (** Get and extract the element with maximum key from the table. *)
  let decompose_max (t:t 'a) : option (t 'a,D.t 'a)
    returns { None -> t.card = 0 /\ forall k. not t.domn k
      | Some (l:t 'a,d) -> t.card = l.card + 1 /\
        t.domn d.key /\ t.func d.key = d /\ not l.domn d.key /\
        (forall k. CO.lt d.key k -> not t.domn k) /\
        (forall k. CO.eq k d.key -> t.func k = d) /\
        (forall k. CO.le d.key k -> not l.domn k) /\
        (forall k. CO.lt k d.key -> l.domn k <-> t.domn k) /\
        (forall k. CO.lt k d.key /\ l.domn k -> l.func k = t.func k) }
  = match Sel.decompose_back t.field with
    | None -> None
    | Some (l,d) -> Some ({ field = l },d)
    end

  (** Optimized insertion of an element with minimum key. *)
  let add_min (d:D.t 'a) (t:t 'a) : t 'a
    requires { forall k. t.domn k -> CO.lt d.key k }
    ensures { result.card = t.card + 1 }
    ensures { forall k. CO.lt k d.key -> not result.domn k }
    ensures { forall k. CO.eq k d.key -> result.domn k /\ result.func k = d }
    ensures { forall k. CO.lt d.key k -> result.domn k <-> t.domn k }
    ensures { forall k. CO.lt d.key k /\ result.domn k ->
      result.func k = t.func k }
  = { field = Sel.cons d t.field }

  (** Optimized insertion of an element with maximal key. *)
  let add_max (t:t 'a) (d:D.t 'a) : t 'a
    requires { forall k. t.domn k -> CO.lt k d.key }
    ensures { result.card = t.card + 1 }
    ensures { forall k. CO.lt d.key k -> not result.domn k }
    ensures { forall k. CO.eq k d.key -> result.domn k /\ result.func k = d }
    ensures { forall k. CO.lt k d.key -> result.domn k <-> t.domn k }
    ensures { forall k. CO.lt k d.key /\ result.domn k ->
      result.func k = t.func k }
  = { field = Sel.snoc t.field d }

  (** Ordered union of two associative tables. *)
  let ordered_union (l r:t 'a) : t 'a
    requires { forall k1 k2. l.domn k1 /\ r.domn k2 -> CO.lt k1 k2 }
    ensures  { forall k. result.domn k <-> l.domn k \/ r.domn k }
    ensures  { forall k. result.domn k /\ l.domn k -> result.func k = l.func k }
    ensures  { forall k. result.domn k /\ r.domn k -> result.func k = r.func k }
    ensures  { result.card = l.card + r.card }
  = { field = Sel.concat l.field r.field }

  (** Get the value associated to a key in the table, if it exists. *)
  let get (k:K.t) (t:t 'a) : option (D.t 'a)
    returns { None -> not t.domn k
      | Some d -> t.card > 0 /\ CO.eq d.key k /\ t.domn k /\ t.func k = d }
  = Sel.get (Sel.default_split ()) k t.field

  (** Insert a value in the table. Erase binding if the key was
      already present. *)
  let insert (d:D.t 'a) (t:t 'a) : t 'a
    ensures { result.card = if t.domn d.key then t.card else t.card + 1 }
    ensures { forall k. CO.eq k d.key -> result.domn k /\ result.func k = d }
    ensures { result.domn d.key /\ result.func d.key = d }
    ensures { forall k. not CO.eq k d.key -> result.domn k <-> t.domn k }
    ensures { forall k. not CO.eq k d.key /\ result.domn k ->
      result.func k = t.func k }
  = let r0 = Sel.default_split () in
    let t2 = Sel.insert r0 d.key d t.field in
    selected_model d.key !r0;
    let r1 = { !r0 with middle = Some d } in
    assert { rebuild r1 = t2 };
    selected_model d.key r1;
    { field = t2 }

  (** Erase any binding associated to the given key. *)
  let remove (k0:K.t) (t:t 'a) : t 'a
    ensures { result.card = if t.domn k0 then t.card - 1 else t.card }
    ensures { forall k. CO.eq k k0 -> not result.domn k }
    ensures { not result.domn k0 }
    ensures { forall k. not CO.eq k k0 -> result.domn k <-> t.domn k }
    ensures { forall k. not CO.eq k k0 /\ result.domn k ->
      result.func k = t.func k }
  = { field = Sel.remove (Sel.default_split ()) k0 t.field }

  (** Split the table in three: elements before the key,
      element associated to the key (if exists),
      elements after the key. *)
  let split (k0:K.t) (t:t 'a) : (t 'a,option (D.t 'a),t 'a)
    returns { (lf:t 'a,o,rg:t 'a) ->
      (forall k. CO.lt k k0 -> lf.domn k <-> t.domn k) /\
      (forall k. CO.lt k k0 /\ lf.domn k -> lf.func k = t.func k) /\
      (forall k. CO.le k0 k -> not lf.domn k) /\
      (forall k. CO.lt k0 k -> rg.domn k <-> t.domn k) /\
      (forall k. CO.lt k0 k /\ rg.domn k -> rg.func k = t.func k) /\
      (forall k. CO.le k k0 -> not rg.domn k) /\
      match o with
      | None -> t.card = lf.card + rg.card /\
        (forall k. CO.eq k k0 -> not t.domn k) /\
        not t.domn k0
      | Some d -> t.card = lf.card + rg.card + 1 /\
        CO.eq d.key k0 /\
        (forall k. CO.eq k k0 -> t.domn k /\ t.func k = d)
        /\ t.domn k0 /\ t.func k0 = d
      end }
  = let lf,o,rg = Sel.split (Sel.default_split ()) k0 t.field in
    ({ field = lf }, o, { field = rg })

  (** {3 Extension with set-theoretic routines }

      Those routines go beyond single-call to the AVL trees function.
      Also, unlike previous routines they are not logarithmic-time but
      linear-time instead. *)

  (** Internal specification wrappers over the AVL view and join routines. *)
  type view 'a =
    | AEmpty
    | ANode (t 'a) (D.t 'a) (t 'a) I.t

  let view (t:t 'a) : view 'a
    returns { AEmpty -> t.card = 0 /\ forall k:K.t. not t.domn k
      | ANode l d r h ->
        t.card = 1 + l.card + r.card
        /\ h = t.field.Sel.hgt
        /\ let k0 = d.key in
           (forall k. CO.lt k k0 -> l.domn k <-> t.domn k)
        /\ (forall k. CO.lt k k0 /\ l.domn k -> l.func k = t.func k)
        /\ (forall k. CO.le k0 k -> not l.domn k)
        /\ (forall k. CO.lt k0 k -> r.domn k <-> t.domn k)
        /\ (forall k. CO.lt k0 k /\ r.domn k -> r.func k = t.func k)
        /\ (forall k. CO.le k k0 -> not r.domn k)
        /\ (forall k. CO.eq k k0 -> t.domn k /\ t.func k = d)
        /\ t.domn k0 /\ t.func k0 = d }
  = match Sel.view t.field with
    | Sel.AEmpty -> AEmpty
    | Sel.ANode l d r h _ -> ANode { field = l } d { field = r } h
    end

  let join (l:t 'a) (d:D.t 'a) (r:t 'a) : t 'a
    requires { forall k. l.domn k -> CO.lt k d.key }
    requires { forall k. r.domn k -> CO.lt d.key k }
    ensures { forall k. CO.lt k d.key -> result.domn k <-> l.domn k }
    ensures { forall k. CO.lt k d.key /\ result.domn k ->
      result.func k = l.func k }
    ensures { forall k. CO.lt d.key k -> result.domn k <-> r.domn k }
    ensures { forall k. CO.lt d.key k /\ result.domn k ->
      result.func k = r.func k }
    ensures { forall k. CO.eq k d.key -> result.domn k /\ result.func k = d }
    ensures { result.card = 1 + l.card + r.card }
  = let ghost s = { left = l.field.Sel.m.Sel.seq;
                    middle = Some d;
                    right = r.field.Sel.m.Sel.seq }
    in
    assert { selected d.key s };
    { field = Sel.join l.field d r.field }

  (** Add every element from `a` into `t`. *)
  let rec add_all (a:t 'a) (t:t 'a) : t 'a
    ensures { forall k. result.domn k <-> a.domn k \/ t.domn k }
    ensures { forall k. a.domn k -> result.func k = a.func k }
    ensures { forall k. t.domn k /\ not a.domn k -> result.func k = t.func k }
    ensures { result.card >= a.card /\ result.card >= t.card }
    ensures { result.card <= a.card + t.card }
    variant { a.m.card + t.m.card }
  = match view a with
    | AEmpty -> t
    | ANode al ad ar ah -> match view t with
      | AEmpty -> a
      | ANode tl td tr th -> if I.le ah th
        then let (al,ad,ar) = split td.key a in
          let ul = add_all al tl in
          let ur = add_all ar tr in
          let ud = match ad with
            | None -> td
            | Some ad -> ad
            end in
          join ul ud ur
        else let (tl,_,tr) = split ad.key t in
          let ul = add_all al tl in
          let ur = add_all ar tr in
          join ul ad ur
      end
    end

  (** Create the table with the elements of `a`
      whose key appear in `p`. *)
  let rec filter (p:t 'b) (a:t 'a) : t 'a
    ensures { forall k. result.domn k <-> a.domn k /\ p.domn k }
    ensures { forall k. result.domn k -> result.func k = a.func k }
    ensures { result.m.card <= a.m.card /\ result.m.card <= p.m.card }
    variant { a.m.card + p.m.card }
  = match view a with
    | AEmpty -> a
    | ANode al ad ar ah -> match view p with
      | AEmpty -> empty ()
      | ANode pl pd pr ph -> if I.le ah ph
        then let (al,ad,ar) = split pd.key a in
          let fl = filter pl al in
          let fr = filter pr ar in
          match ad with
          | None -> ordered_union fl fr
          | Some ad -> join fl ad fr
          end
        else let (pl,pd,pr) = split ad.key p in
          let fl = filter pl al in
          let fr = filter pr ar in
          match pd with
          | None -> ordered_union fl fr
          | _ -> join fl ad fr
          end
      end
    end

  (** Complement of `filter`: remove from `a` every element whose
      key appear in `p`. *)
  let rec remove_all (p:t 'b) (a:t 'a) : t 'a
    ensures { forall k. result.domn k <-> a.domn k /\ not p.domn k }
    ensures { forall k. result.domn k -> result.func k = a.func k }
    ensures { result.m.card <= a.m.card }
    variant { a.m.card + p.m.card }
  = match view a with
    | AEmpty -> a
    | ANode al ad ar ah -> match view p with
      | AEmpty -> a
      | ANode pl pd pr ph -> if I.le ah ph
        then let (al,_,ar) = split pd.key a in
          let fl = remove_all pl al in
          let fr = remove_all pr ar in
          ordered_union fl fr
        else let (pl,pd,pr) = split ad.key p in
          let fl = remove_all pl al in
          let fr = remove_all pr ar in
          match pd with
          | None -> join fl ad fr
          | _ -> ordered_union fl fr
          end
      end
    end

  (** Create a table with the elements that appear
      exactly in one of `a` and `b`, but not both. *)
  let rec symdiff (a b:t 'a) : t 'a
    ensures { forall k. result.domn k <-> not (a.domn k <-> b.domn k) }
    ensures { forall k. result.domn k /\ a.domn k -> result.func k = a.func k }
    ensures { forall k. result.domn k /\ b.domn k -> result.func k = b.func k }
    ensures { result.m.card <= a.m.card + b.m.card }
    variant { a.m.card + b.m.card }
  = match view a with
    | AEmpty -> b
    | ANode al ad ar ah -> match view b with
      | AEmpty -> a
      | ANode bl bd br bh -> if I.le ah bh
        then let (al,ad,ar) = split bd.key a in
          let sl = symdiff al bl in
          let sr = symdiff ar br in
          match ad with
          | None -> join sl bd sr
          | _ -> ordered_union sl sr
          end
        else let (bl,bd,br) = split ad.key b in
          let sl = symdiff al bl in
          let sr = symdiff ar br in
          match bd with
          | None -> join sl ad sr
          | _ -> ordered_union sl sr
          end
      end
    end

end

(** {2 Instantiation of the base to key-value ordered associative tables} *)
module Map

  use int.Int
  use option.Option
  use ref.Ref
  use mach.peano.Peano as I

  val constant balancing : I.t
    ensures { result > 0 }

  (** Parameter: key type with computable total preorder. *)
  scope K type t end
  clone preorder.Computable as CO with type t = K.t, axiom .

  (** Elements are key-value pairs *)
  scope D
    type t 'a = (K.t,'a)
    let function key (t:t 'a) : K.t = let (a,_) = t in a
  end

(** Direct instantiation. *)
  clone MapBase as MB with val balancing = balancing,
    type K.t = K.t,
    type D.t = D.t,
    val key = D.key,
    predicate CO.le = CO.le,
    goal CO.Refl,
    goal CO.Trans,
    goal CO.Total,
    predicate CO.lt = CO.lt,
    goal CO.lt_def,
    predicate CO.eq = CO.eq,
    goal CO.eq_def,
    val CO.compare = CO.compare

  (** Slight adaptation of the logical model
      (get rid of keys in the output) *)
  type t 'a = MB.t 'a
  type m 'a = {
    domn : K.t -> bool;
    func : K.t -> 'a;
    card : int;
  }
  function m (t:t 'a) : m 'a =
    { domn = t.MB.domn;
      func = fun k -> let (_,v) = t.MB.func k in v;
      card = t.MB.m.MB.card }
  meta coercion function m

  let lemma correction (t:t 'a)
    ensures { forall k1 k2. CO.eq k1 k2 -> t.domn k1 <-> t.domn k2 }
    ensures { forall k1 k2. CO.eq k1 k2 /\ t.domn k1 -> t.func k1 = t.func k2 }
  = ()

  let empty () : t 'a
    ensures { forall k. not result.domn k }
    ensures { result.card = 0 }
  = MB.empty ()

  (** Create a table with a single key->value binding. *)
  let singleton (k0:K.t) (v:'a) : t 'a
    ensures { forall k. result.domn k <-> CO.eq k k0 }
    ensures { forall k. CO.eq k k0 -> result.func k = v }
    ensures { result.domn k0 /\ result.func k0 = v }
    ensures { result.card = 1 }
  = MB.singleton (k0,v)

  (** Check emptyness of a table. *)
  let is_empty (ghost rk:ref K.t) (t:t 'a) : bool
    ensures { result -> forall k. not t.domn k }
    ensures { not result -> t.domn !rk }
    ensures { result <-> t.m.card = 0 }
  = MB.is_empty rk t

  let min_binding (t:t 'a) : (K.t,'a)
    requires { t.card > 0 }
    returns  { (k,v) -> t.domn k /\ t.func k = v
      /\ forall k2. t.domn k2 -> CO.le k k2 }
  = MB.min t

  let max_binding (t:t 'a) : (K.t,'a)
    requires { t.card > 0 }
    returns  { (k,v) -> t.domn k /\ t.func k = v
      /\ forall k2. t.domn k2 -> CO.le k2 k }
  = MB.max t

  (** Get and extract the (key->value) binding with minimum key. *)
  let decompose_min (t:t 'a) : option ((K.t,'a),t 'a)
    returns { None -> t.card = 0 /\ forall k. not t.domn k
      | Some ((k0,v),r:t 'a) -> t.card = r.card + 1 /\
        t.domn k0 /\ t.func k0 = v /\ not r.domn k0 /\
        (forall k. CO.lt k k0 -> not t.domn k) /\
        (forall k. CO.eq k k0 -> t.func k = v) /\
        (forall k. CO.le k k0 -> not r.domn k) /\
        (forall k. CO.lt k0 k -> r.domn k <-> t.domn k) /\
        (forall k. CO.lt k0 k /\ r.domn k -> r.func k = t.func k) }
  = MB.decompose_min t

  (** Get and extract the (key->value) binding with maximum key. *)
  let decompose_max (t:t 'a) : option (t 'a,(K.t,'a))
    returns { None -> t.card = 0 /\ forall k. not t.domn k
      | Some (l:t 'a,(k0,v)) -> t.card = l.card + 1 /\
        t.domn k0 /\ t.func k0 = v /\ not l.domn k0 /\
        (forall k. CO.lt k0 k -> not t.domn k) /\
        (forall k. CO.eq k k0 -> t.func k = v) /\
        (forall k. CO.le k0 k -> not l.domn k) /\
        (forall k. CO.lt k k0 -> l.domn k <-> t.domn k) /\
        (forall k. CO.lt k k0 /\ l.domn k -> l.func k = t.func k) }
  = MB.decompose_max t

  (** Add a key->value binding with minimal key. *)
  let add_min_binding (k0:K.t) (v:'a) (t:t 'a) : t 'a
    requires { forall k. t.domn k -> CO.lt k0 k }
    ensures { result.card = t.card + 1 }
    ensures { forall k. CO.lt k k0 -> not result.domn k }
    ensures { forall k. CO.eq k k0 -> result.domn k /\ result.func k = v }
    ensures { forall k. CO.lt k0 k -> result.domn k <-> t.domn k }
    ensures { forall k. CO.lt k0 k /\ result.domn k ->
      result.func k = t.func k }
    ensures { result.domn k0 /\ result.func k0 = v }
  = MB.add_min (k0,v) t

  (** Add a key->value binding with maximal key. *)
  let add_max_binding (k0:K.t) (v:'a) (t:t 'a) : t 'a
    requires { forall k. t.domn k -> CO.lt k k0 }
    ensures { result.card = t.card + 1 }
    ensures { forall k. CO.lt k0 k -> not result.domn k }
    ensures { forall k. CO.eq k k0 -> result.domn k /\ result.func k = v }
    ensures { forall k. CO.lt k k0 -> result.domn k <-> t.domn k }
    ensures { forall k. CO.lt k k0 /\ result.domn k ->
      result.func k = t.func k }
    ensures { result.domn k0 /\ result.func k0 = v }
  = MB.add_max t (k0,v)

  (** Ordered fusion of two associative tables. *)
  let ordered_join (l r:t 'a) : t 'a
    requires { forall k1 k2. l.domn k1 /\ r.domn k2 -> CO.lt k1 k2 }
    ensures  { forall k. result.domn k <-> l.domn k \/ r.domn k }
    ensures  { forall k. result.domn k /\ l.domn k -> result.func k = l.func k }
    ensures  { forall k. result.domn k /\ r.domn k -> result.func k = r.func k }
    ensures  { result.card = l.card + r.card }
  = MB.ordered_union l r

  (** Extract the value associated to some key. *)
  let get (k:K.t) (t:t 'a) : option 'a
    returns { None -> not t.domn k
      | Some v -> t.card > 0 /\ t.domn k /\ t.func k = v }
  = match MB.get k t with
    | None -> None
    | Some (_,v) -> Some v
    end

  (** Set the binding for key `k`, erasing any such previous binding. *)
  let insert (k0:K.t) (v:'a) (t:t 'a) : t 'a
    ensures { result.card = if t.domn k0 then t.card else t.card + 1 }
    ensures { forall k. CO.eq k k0 -> result.domn k /\ result.func k = v }
    ensures { result.domn k0 /\ result.func k0 = v }
    ensures { forall k. not CO.eq k k0 -> result.domn k <-> t.domn k }
    ensures { forall k. not CO.eq k k0 /\ result.domn k ->
      result.func k = t.func k }
  = MB.insert (k0,v) t

  (** Erase any potential binding for key `k`. *)
  let remove (k0:K.t) (t:t 'a) : t 'a
    ensures { result.card = if t.domn k0 then t.card - 1 else t.card }
    ensures { forall k. CO.eq k k0 -> not result.domn k }
    ensures { not result.domn k0 }
    ensures { forall k. not CO.eq k k0 -> result.domn k <-> t.domn k }
    ensures { forall k. not CO.eq k k0 /\ result.domn k ->
      result.func k = t.func k }
  = MB.remove k0 t

  (** Split the table in three parts:
      Bindings with key lower than `k`, value associated to `k`,
      and bindings with key greater than `k`. *)
  let split (k0:K.t) (t:t 'a) : (t 'a,option 'a,t 'a)
    returns { (lf:t 'a,o,rg:t 'a) ->
      (forall k. CO.lt k k0 -> lf.domn k <-> t.domn k) /\
      (forall k. CO.lt k k0 /\ lf.domn k -> lf.func k = t.func k) /\
      (forall k. CO.le k0 k -> not lf.domn k) /\
      (forall k. CO.lt k0 k -> rg.domn k <-> t.domn k) /\
      (forall k. CO.lt k0 k /\ rg.domn k -> rg.func k = t.func k) /\
      (forall k. CO.le k k0 -> not rg.domn k) /\
      match o with
      | None -> t.card = lf.card + rg.card /\
        (forall k. CO.eq k k0 -> not t.domn k) /\
        not t.domn k0
      | Some v -> t.card = lf.card + rg.card + 1 /\
        (forall k. CO.eq k k0 -> t.domn k /\ t.func k = v)
        /\ t.domn k0 /\ t.func k0 = v
      end }
  = let (lf,o,rg) = MB.split k0 t in
    let o = match o with None -> None | Some (_,v) -> Some v end in
    (lf,o,rg)

end

(** {2 Instantiation of the base to ordered sets} *)
module Set

  use int.Int
  use option.Option
  use ref.Ref
  use import mach.peano.Peano as I

  (** The balancing level is left abstract. *)
  val constant balancing : I.t
    ensures { result > 0 }

  (** Parameter: comparable elements. *)
  scope K type t end
  clone preorder.Computable as CO with type t = K.t, axiom .

  (** Elements are themselves the keys. *)
  scope D
    type t 'a = K.t
    let function key (x:'a) : 'a = x
  end

  (** Actual instantiation. *)
  clone MapBase as MB with val balancing = balancing,
    type K.t = K.t,
    type D.t = D.t,
    val key = D.key,
    predicate CO.le = CO.le,
    goal CO.Refl,
    goal CO.Trans,
    goal CO.Total,
    predicate CO.lt = CO.lt,
    goal CO.lt_def,
    predicate CO.eq = CO.eq,
    goal CO.eq_def,
    val CO.compare = CO.compare

  (** Minor adaptation of the logical model: only keep the set of elements.
      In fact, the functions from MapBase can been used as-is
      (but may be over-specifications)
      The following specifications wrappers are only given for demonstration. *)
  type t = MB.t unit
  function set (m:MB.m unit) : K.t -> bool = m.MB.domn
  function card (m:MB.m unit) : int = m.MB.card

  (** Invariant on the logical model. *)
  let lemma correction (t:t) : unit
    ensures { forall k1 k2:K.t. CO.eq k1 k2 -> t.set k1 <-> t.set k2 }
    ensures { t.card >= 0 }
  = ()

  (** Create an empty set. *)
  let empty () : t
    ensures { forall k. not result.set k }
    ensures { result.card = 0 }
  = MB.empty ()

  (** Create a singleton set. *)
  let singleton (k0:K.t) : t
    ensures { forall k. result.set k <-> CO.eq k k0 }
    ensures { result.set k0 }
    ensures { result.card = 1 }
  = MB.singleton k0

  (** Test emptyness of a set. *)
  let is_empty (ghost rk:ref K.t) (t:t) : bool
    ensures { result -> forall k. not t.set k }
    ensures { not result -> t.set !rk }
    ensures { result <-> t.card = 0 }
  = MB.is_empty rk t

  (** Minimum element of a set. *)
  let min_elt (t:t) : K.t
    requires { t.card > 0 }
    ensures  { t.set result }
    ensures  { forall k. t.set k -> CO.le result k }
  = MB.min t

  (** Maximum element of a set. *)
  let max_elt (t:t) : K.t
    requires { t.card > 0 }
    ensures  { t.set result }
    ensures  { forall k. t.set k -> CO.le k result }
  = MB.max t

  (** Get and remove minimum element from a set. *)
  let decompose_min (t:t) : option (K.t,t)
    returns { None -> t.card = 0 /\ forall k. not t.set k
      | Some (k0,r:t) -> t.card = r.card + 1 /\
        t.set k0 /\ not r.set k0 /\
        (forall k. CO.lt k k0 -> not t.set k) /\
        (forall k. CO.le k k0 -> not r.set k) /\
        (forall k. CO.lt k0 k -> r.set k <-> t.set k) }
  = MB.decompose_min t

  (** Get and remove maximum element from a set. *)
  let decompose_max (t:t) : option (t,K.t)
    returns { None -> t.card = 0 /\ forall k. not t.set k
      | Some (l:t,k0) -> t.card = l.card + 1 /\
        t.set k0 /\ not l.set k0 /\
        (forall k. CO.lt k0 k -> not t.set k) /\
        (forall k. CO.le k0 k -> not l.set k) /\
        (forall k. CO.lt k0 k -> l.set k <-> t.set k) }
  = MB.decompose_max t

  (** Add minimal element to a set. *)
  let add_min_elt (k0:K.t) (t:t) : t
    requires { forall k. t.set k -> CO.lt k0 k }
    ensures { result.card = t.card + 1 }
    ensures { forall k. CO.lt k k0 -> not result.set k }
    ensures { forall k. CO.lt k0 k -> result.set k <-> t.set k }
  = MB.add_min k0 t

  (** Add maximal element to a set. *)
  let add_max_elt (t:t) (k0:K.t) : t
    requires { forall k. t.set k -> CO.lt k k0 }
    ensures { result.card = t.card + 1 }
    ensures { forall k. CO.lt k0 k -> not result.set k }
    ensures { forall k. CO.lt k k0 -> result.set k <-> t.set k }
  = MB.add_max t k0

  (** Ordered union of two sets. *)
  let ordered_union (l r:t) : t
    requires { forall k1 k2. l.set k1 /\ r.set k2 -> CO.lt k1 k2 }
    ensures  { forall k. result.set k <-> l.set k \/ r.set k }
    ensures  { result.card = l.card + r.card }
  = MB.ordered_union l r

  (** Test membership of an element. *)
  let mem (k0:K.t) (t:t) : bool
    ensures { result <-> t.set k0 }
    ensures { result -> forall k. CO.eq k k0 -> t.set k }
    ensures { forall k. CO.eq k k0 /\ t.set k -> result }
    ensures { result -> t.card > 0 }
  = match MB.get k0 t with None -> false | _ -> true end

  (** Add an element to a set. *)
  let add (k0:K.t) (t:t) : t
    ensures { result.card = if t.set k0 then t.card else t.card + 1 }
    ensures { forall k. CO.eq k k0 -> result.set k }
    ensures { result.set k0 }
    ensures { forall k. not CO.eq k k0 -> result.set k <-> t.set k }
  = MB.insert k0 t

  (** Remove an element from a set. *)
  let remove (k0:K.t) (t:t) : t
    ensures { result.card = if t.set k0 then t.card - 1 else t.card }
    ensures { forall k. CO.eq k k0 -> not result.set k }
    ensures { not result.set k0 }
    ensures { forall k. not CO.eq k k0 -> result.set k <-> t.set k }
  = MB.remove k0 t

  (** Split the set into three parts: elements lower than `k`,
      elements equal to `k`,
      and elements bigger than `k` *)
  let split (k0:K.t) (t:t) : (t,bool,t)
    returns { (lf:t,b,rg:t) ->
      (forall k. CO.lt k k0 -> lf.set k <-> t.set k) /\
      (forall k. CO.le k0 k -> not lf.set k) /\
      (forall k. CO.lt k0 k -> rg.set k <-> t.set k) /\
      (forall k. CO.le k k0 -> not rg.set k) /\
      t.card = lf.card + rg.card + (if b then 1 else 0) /\
      (b <-> t.set k0) /\
      (b -> forall k. CO.eq k k0 -> t.set k) /\
      (forall k. CO.eq k k0 /\ t.set k -> b)
    }
  = let (lf,o,rg) = MB.split k0 t in
    let b = match o with None -> false | _ -> true end in
    (lf,b,rg)

  (** Extension: set-theoretic routines. *)

  (** Compute the union of two sets. *)
  let union (a b:t) : t
    ensures { forall k. result.set k <-> a.set k \/ b.set k }
    ensures { result.card >= a.card /\ result.card >= b.card }
    ensures { result.card <= a.card + b.card }
  = MB.add_all a b

  (** Compute the intersection of two sets. *)
  let inter (a b:t) : t
    ensures { forall k. result.set k <-> a.set k /\ b.set k }
    ensures { result.card <= a.card /\ result.card <= b.card }
  = MB.filter a b

  (** Compute the difference of two sets. *)
  let diff (a b:t) : t
    ensures { forall k. result.set k <-> a.set k /\ not b.set k }
    ensures { result.card <= a.card }
  = MB.remove_all b a

  (** Compute the symmetrical difference of two sets. *)
  let symdiff (a b:t) : t
    ensures { forall k. result.set k <-> not (a.set k <-> b.set k) }
    ensures { result.card <= a.card + b.card }
  = MB.symdiff a b

end

(** Example instances: integer keys/elements *)
module IMapAndSet

  use int.Int
  use mach.peano.Peano as I

  scope K type t = int end

  let constant balancing : I.t = I.succ I.zero

  predicate le (x y:int) = x <= y
  predicate eq (x y:int) = x = y
  predicate lt (x y:int) = x < y

  let compare (x y:int) : int ensures { result = x - y } = x - y

  clone Map as M with val balancing = balancing,
    type K.t = K.t,
    predicate CO.le = le,
    predicate CO.lt = lt,
    predicate CO.eq = eq,
    goal CO.lt_def,
    goal CO.eq_def,
    goal CO.Refl,
    goal CO.Trans,
    goal CO.Total,
    val CO.compare = compare

  clone Set as S with val balancing = balancing,
    type K.t = K.t,
    predicate CO.le = le,
    predicate CO.lt = lt,
    predicate CO.eq = eq,
    goal CO.lt_def,
    goal CO.eq_def,
    goal CO.Refl,
    goal CO.Trans,
    goal CO.Total,
    val CO.compare = compare

end
