datatype Tree<T> = Leaf | Node(Tree<T>, Tree<T>, T)
datatype List<T> = Nil | Cons(T, List<T>)

function flatten<T>(tree: Tree<T>): List<T>
{
	match(tree)
    case Leaf() =>  Nil
    case Node(tree1, tree2, x) => Cons(x, append(flatten(tree1), flatten(tree2)))
}

function append<T>(xs: List<T>, ys: List<T>): List<T>
  ensures xs == Nil ==> append(xs, ys) == ys
  ensures (ys == Nil) == (append(xs, ys) == xs)
{
	match(xs)
		case Nil => ys
		case Cons(x, xs') =>  Cons(x, append(xs', ys))
}

function listContains<T>(xs: List<T>, element: T): bool
{
	match(xs)
    case Nil => false
		case Cons(x, xs') => (x == element) || listContains(xs', element)
}

function treeContains<T>(tree: Tree<T>, element: T): bool
{
	match(tree)
    case Leaf() => false
    case Node(tree1, tree2, x) => (x == element) || treeContains(tree1, element) || treeContains(tree2, element)
}

lemma littleAppendsAreBigAppend<T>(xs: List<T>, ys: List<T>, element: T)
  ensures listContains(xs, element) || listContains(ys, element) <==> listContains(append(xs, ys) ,element)
{

}

lemma sameElements<T>(tree: Tree<T>, element: T)
  ensures treeContains(tree, element) <==> listContains(flatten(tree), element)
{
  match(tree) {
     case Leaf() => {
       assert treeContains(tree, element)
           == treeContains(Leaf(), element)
           == false
           == listContains<T>(Nil, element)
           == listContains(flatten(tree), element);
      }
      case Node(tree1, tree2, x) => {
        assert treeContains(tree, element)
          == treeContains(Node(tree1, tree2, x), element)
          == (x == element) || treeContains(tree1, element) || treeContains(tree2, element);

        sameElements(tree1, element);
        sameElements(tree2, element);

        assert (x == element) || treeContains(tree1, element) || treeContains(tree2, element)
          == (x == element) || listContains(flatten(tree1), element) || listContains(flatten(tree2), element);

        littleAppendsAreBigAppend(flatten(tree1), flatten(tree2), element);

        assert listContains(flatten(tree), element)
          == listContains(flatten(Node(tree1, tree2, x)), element)
          == listContains(append(Cons(x, flatten(tree1)), flatten(tree2)), element)
          == listContains(Cons(x, append(flatten(tree1), flatten(tree2))), element)
          == (x == element) || listContains(append(flatten(tree1), flatten(tree2)), element)
          == (x == element) || listContains(flatten(tree1), element) || listContains(flatten(tree2), element);
    }
  }
}