Toccata

(*
   Generate all binary trees of size n.

   Given n, the program return an array a of size n+1 such that
   a[i] contains the list of all binary trees of size i.

   TODO: tail-recursive version of combine
*)

module GenerateAllTrees

  use int.Int
  use list.List
  use list.Mem
  use list.Append
  use list.Distinct
  use array.Array
  use list.Length

  type tree = Empty | Node tree tree

  function size (t: tree) : int = match t with
    | Empty    -> 0
    | Node l r -> 1 + size l + size r
  end

  lemma size_nonneg: forall t: tree. size t >= 0

  lemma size_left:
    forall t: tree. size t > 0 ->
    exists l r: tree. t = Node l r /\ size l < size t

  predicate all_trees (n: int) (l: list tree) =
    distinct l /\
    forall t: tree. size t = n <-> mem t l

  lemma all_trees_0: all_trees 0 (Cons Empty Nil)

  lemma tree_diff:
    forall l1 l2: tree. size l1 <> size l2 ->
    forall r1 r2: tree. Node l1 r1 <> Node l2 r2

  (* combines two lists of trees l1 and l2 into the list of trees
     with a left sub-tree from l1 and a right sub-tree from l2 *)
  let combine (i1: int) (l1: list tree) (i2: int) (l2: list tree)
    requires { 0 <= i1 /\ all_trees i1 l1 /\ 0 <= i2 /\ all_trees i2 l2 }
    ensures  { distinct result }
    ensures  { forall t:tree. mem t result <->
      (exists l r: tree. t = Node l r /\ size l = i1 /\ size r = i2) }
  = let rec loop1 (l1: list tree) : list tree variant { l1 }
      requires { distinct l1 }
      ensures  { distinct result }
      ensures  { forall t:tree. mem t result <->
        (exists l r: tree. t = Node l r /\ mem l l1 /\ mem r l2) }
    = match l1 with
      | Nil -> Nil
      | Cons t1 l1 ->
          let rec loop2 (l2: list tree) : list tree variant { l2 }
            requires { distinct l2 }
            ensures  { distinct result }
            ensures  { forall t:tree. mem t result <->
              (exists r: tree. t = Node t1 r /\ mem r l2) }
          = match l2 with
            | Nil -> Nil
            | Cons t2 l2 -> Cons (Node t1 t2) (loop2 l2)
            end
         in
         loop2 l2 ++ loop1 l1
      end
    in
    loop1 l1

  let all_trees (n: int)
    requires { n >= 0 }
    ensures { forall i: int. 0 <= i <= n -> all_trees i result[i] }
  = let a = make (n+1) Nil in
    a[0] <- Cons Empty Nil;
    for i = 1 to n do
      invariant { forall k: int. 0 <= k < i -> all_trees k a[k] }
      a[i] <- Nil;
      for j = 0 to i-1 do
        invariant { forall k: int. 0 <= k < i -> all_trees k a[k] }
        invariant { distinct a[i] }
        invariant { forall t: tree. mem t a[i] <->
          (exists l r: tree. t = Node l r /\ size t = i /\ size l < j) }
        a[i] <- (combine j a[j] (i-1-j) a[i-1-j]) ++ a[i]
      done
    done;
    a

end