
(* Tree relabelling.

   Given a tree t, with values at the leaves, build a new tree with the
   same structure and leaves labelled with distinct integers.
 *)

module Relabel

  use list.List
  use list.Mem
  use list.Append
  use list.Distinct

  type tree 'a =
    | Leaf 'a
    | Node (tree 'a) (tree 'a)

  function labels (t : tree 'a) : list 'a = match t with
    | Leaf x -> Cons x Nil
    | Node l r -> labels l ++ labels r
  end

  lemma labels_Leaf :
    forall x y : 'a. mem x (labels (Leaf y)) <-> x=y

  lemma labels_Node :
    forall x : 'a, l r : tree 'a.
      mem x (labels (Node l r)) <-> (mem x (labels l) \/ mem x (labels r))

  inductive same_shape (t1 : tree 'a) (t2 : tree 'b) =
    | same_shape_Leaf :
        forall x1 : 'a, x2 : 'b.
        same_shape (Leaf x1) (Leaf x2)
    | same_shape_Node :
        forall l1 r1 : tree 'a, l2 r2 : tree 'b.
        same_shape l1 l2 -> same_shape r1 r2 ->
        same_shape (Node l1 r1) (Node l2 r2)

  use int.Int
  use ref.Ref

  val r : ref int

  let fresh () ensures { !r = old !r + 1 /\ result = !r } =
    r := !r + 1; !r

  let rec relabel (t : tree 'a) : tree int
    variant { t }
    ensures { same_shape t result /\ distinct (labels result) /\
      old !r <= !r /\
      (forall x:int. mem x (labels result) -> old !r < x <= !r) }
  = match t with
    | Leaf _   -> Leaf (fresh ())
    | Node l r -> Node (relabel l) (relabel r)
    end

end
