How to use lambda expression or other ways to batch extract specific member of a new datatype in Dafny?

Question

I'm using dafny to verify some properties with respect to trees. Hence I defined a datatype called "Tree"

datatype Tree = Empty | Node(key: int, left: Tree, right: Tree, max_size: real, exist: bool)
// other codes

You can see here all the trees are binary tree.(In fact, all are binary search tree, promised by other lemmas and predicates omitted here)

Now I have a function to collect all nodes in a tree t

function tree_nodes(t: Tree): (s: set<Tree>)
{
  match t
  case Empty => {}
  case Node(_, l, r, _, _) => tree_nodes(l) + {t} + tree_nodes(r)
}

Also, I need collect all keys stored in nodes

function tree_elements(t: Tree): (s: set<int>)
{
  match t
  case Empty => {}
  case Node(k, l, r, _, _) => tree_elements(l) + {k} + tree_elements(r)
}

Obviously, it is easy to find |tree_elements(t)| == |tree_nodes(t)|, but dafny cannot prove that. So I changed this function like

function tree_elements(t: Tree): (s: set<int>)
{
  var node_set := tree_nodes(t);
  while true
  {
    // for each node `n` in node_set, extract the key `n.key` and add it into s
  }
}

However, the compiler said while is not allowed in function body. One easy solution is change the type from function to method. But I hope this statement to keep function type, since lots of specifications have used this function.

lemma equal_properties(p: Tree, q: Tree)
  requires BST(p) && BST(q)
  requires p == q
  requires p != Empty
  ensures tree_elements(p) == tree_elements(q)
{
}
// and many lemmas, predicates and functions followed

The goal is to pass the verifications like

function tree_nodes(t: Tree): (s: set<Tree>)
{
  match t
  case Empty => {}
  case Node(_, l, r, _, _) => tree_nodes(l) + {t} + tree_nodes(r)
}

function tree_elements(t: Tree): (s: set<int>)
{
  match t
  case Empty => {}
  case Node(k, l, r, _, _) => tree_elements(l) + {k} + tree_elements(r)
  // You can modify the implemention of this function, just promise it is a `function` rather than `method`
}

lemma tree_properties(t: Tree)
ensures |tree_nodes(t)| == |tree_elements(t)|
// this is to say, for all the nodes in tree t, their keys are all in tree_elements(t)
ensures t != Empty ==> (forall n :: n in tree_nodes(t) ==> n.key in tree_elements(t))
// this is to say, for each key in tree t, there must exist some node n in tree_nodes(t), the key is stored in n
ensures t != Empty ==> (forall k :: k in tree_elements(t) ==> (exists n :: n in tree_nodes(t) ==> n.key == k))
{}

Answer 1

Your statement is not immediately obvious because you are assuming that the tree cannot contain duplicate keys. If you make that a requirement of the lemma then dafny can verify it.

predicate validTree(t: Tree) {
    if t.Empty? then true else t.key !in tree_elements(t.left)
        && t.key !in tree_elements(t.right)
        && tree_elements(t.left) !! tree_elements(t.right) 
        && t !in tree_nodes(t.left)
        && t !in tree_nodes(t.right)
        && tree_nodes(t.left) !! tree_nodes(t.right)
        && validTree(t.left)
        && validTree(t.right)
}

lemma {:induction false} tree_properties(t: Tree)
    requires validTree(t)
    ensures |tree_nodes(t)| == |tree_elements(t)|
{
    if t.Empty? {
        assert |tree_nodes(t)| == |tree_elements(t)|;
    }else{
        tree_properties(t.left);
        tree_properties(t.right);
        assert |tree_nodes(t)| == |tree_elements(t)|;
    }
}