(*
  let rec heigth (t: tree 'a) (k: int -> 'b) : 'b
    ensures { post k (heigth t) result }
  = match t with
    | Empty -> k 0
    | Node l x r ->
       heigth l (fun hl -> ensures { post k (1 + max hl (height r)) result }
       heigth r (fun hr -> ensures { post k (1 + max hl hr) result }
         k (1 + max hl hr)))
    end

  let main (t: tree 'a) : int = heigth t (fun x -> x)
*)

module TreeHeight

  use bintree.Tree, bintree.Height
  use int.MinMax, int.Int

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

  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 >= 0

  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

  predicate post (k: kont 'a) (arg res: int) = match k with
     | Kid -> arg = res
     | Kleft r k' -> post k' (1 + max arg (height r)) res
     | Kright hl k' -> post k' (1 + max hl arg) res
     end

  let rec height (t: tree 'a) (k: kont 'a) : int
    variant { var_tree t + var_kont k }
    ensures { post k (height t) result }
  = match t with
    | Empty -> apply k 0
    | Node l _ r -> height l (Kleft r k)
    end
  with apply (k: kont 'a) (arg: int) : int
    variant  { var_kont k }
    ensures  { post k arg result }
  = match k with
    | Kid -> arg
    | Kleft r k -> height r (Kright arg k)
    | Kright hl k -> apply k (1 + max hl arg)
    end

  let main (t: tree 'a) : int
    ensures { result = height t }
  = height t Kid

end