(*
  let n_distinct_elements (t: tree 'a) : int
    ensures { result = cardinal (set_of_tree t empty) }
  = let h = ref empty in
    let rec distinct_elements_loop (t: tree 'a) (k: unit -> unit) : unit
      ensures { post k () (set_of_tree t (old !h)) !h () }
    = match t with
      | Empty -> k ()
      | Node l x r ->
         h := add x !h;
         distinct_elements_loop l (fun () ->
           ensures { post k () (set_of_tree r (old !h)) !h () }
         distinct_elements_loop r (fun () ->
           ensures { post k () (old !h) !h () }
           k ()))
      end
    in
    distinct_elements_loop t (fun x -> x);
    cardinal !h
*)

module DiffElementsCPS

  use set.Fset
  use ref.Ref
  use bintree.Tree
  use int.Int

  function set_of_tree (t: tree 'a) (s: fset 'a) : fset 'a =
    match t with
    | Empty -> s
    | Node l x r ->
       let s1 = set_of_tree l (add x s) in
       set_of_tree r s1
    end

  function n_set_of_tree (t: tree 'a) : int =
    cardinal (set_of_tree t empty)

  type kont 'a =
    | Kid
    | Kleft (tree 'a) (kont 'a)
    | Kright (kont 'a)

  type state 'a = fset 'a

  predicate post (k: kont 'a) (arg: unit) (old_s cur_s: state 'a) (res: unit) =
    match k with
    | Kid ->
       old_s == cur_s
    | Kleft r k ->
       post k arg (set_of_tree r old_s) cur_s res
    | Kright k ->
       post k arg old_s cur_s res
    end

  function var_tree (t: tree 'a) : int = match t with
    | Empty -> 1
    | Node l _ r -> 3 + var_tree l + var_tree r
    end

  lemma var_tree_nonneg: forall t: tree 'a.
    var_tree t >= 1

  function var_kont (k: kont 'a) : int = match k with
    | Kid -> 0
    | Kleft r k -> 2 + var_tree r + var_kont k
    | Kright k -> 1 + var_kont k
    end

  lemma var_kont_k_nonneg: forall k: kont 'a.
    var_kont k >= 0

  type elt

  clone set.SetApp with type elt = elt

  let n_distinct_elements (t: tree elt) : int
    ensures { result = n_set_of_tree t }
  = let h = ref (empty ()) in
    let rec distinct_elements_loop (t: tree elt) (k: kont elt) : unit
      ensures { post k () (set_of_tree t (old !h)) !h () }
      variant { var_tree t + var_kont k }
    = match t with
      | Empty -> apply k ()
      | Node l x r ->
         h := add x !h;
         distinct_elements_loop l (Kleft r k)
      end
    with apply (k: kont elt) (arg: unit) : unit
      ensures { post k arg (old !h) !h () }
      variant { var_kont k }
    = match k with
      | Kid ->
          arg
      | Kleft r k ->
         distinct_elements_loop r (Kright k)
      | Kright k ->
         apply k arg
      end
    in
    distinct_elements_loop t Kid;
    cardinal !h

end