Making it generic

In the previous sections, we devised a TreeMap<V> which a used fixed key type (u64). In this section, we’ll show to make a TreeMap<K, V> which is generic over the key type K.

Defining a “total order”

The main reason this is challenging is that the BST requires a way of comparing values of K, both for equality, and to obtain an ordering. This comparison is used both in the implementation (to find the node for a given key, or to figure out where such a node should be inserted) and in the well-formedness invariants that enforce the BST ordering property.

We can define the concept of “total order” generically by creating a trait.

pub enum Cmp { Less, Equal, Greater }

pub trait TotalOrdered : Sized {
    spec fn le(self, other: Self) -> bool;

    proof fn reflexive(x: Self)
        ensures Self::le(x, x);

    proof fn transitive(x: Self, y: Self, z: Self)
        requires Self::le(x, y), Self::le(y, z),
        ensures Self::le(x, z);

    proof fn antisymmetric(x: Self, y: Self)
        requires Self::le(x, y), Self::le(y, x),
        ensures x == y;

    proof fn total(x: Self, y: Self)
        ensures Self::le(x, y) || Self::le(y, x);

    fn compare(&self, other: &Self) -> (c: Cmp)
        ensures (match c {
            Cmp::Less => self.le(*other) && self != other,
            Cmp::Equal => self == other,
            Cmp::Greater => other.le(*self) && self != other,
        });
}

This trait simultaneously:

• Requires a binary relation le to exist
• Requires it to satisfy the properties of a total order
• Requires an executable three-way comparison function to exist

There’s one simplification we’ve made here: we’re assuming that “equality” in the comparison function is the same as spec equality. This isn’t always suitable; some datatypes may have more than one way to represent the same logical value. A more general specification would allow an ordering that respects some arbitrary equivalence relation. This is how vstd::hash_map::HashMapWithView works, for example. To keep things simple for this demonstration though, we’ll use a total ordering that respects spec equality.

Updating the struct and definitions

We’ll start by updating the structs to take a generic parameter K: TotalOrdered.

struct Node<K: TotalOrdered, V> {
    key: K,
    value: V,
    left: Option<Box<Node<K, V>>>,
    right: Option<Box<Node<K, V>>>,
}

pub struct TreeMap<K: TotalOrdered, V> {
    root: Option<Box<Node<K, V>>>,
}

We’ll also update the well-formedness condition to use the generic K::le instead of integer <=. Where the original definition used a < b, we now use a.le(b) && a != b.

impl<K: TotalOrdered, V> Node<K, V> {
    pub closed spec fn well_formed(self) -> bool
        decreases self
    {
        &&& (forall |elem| #[trigger] Node::<K, V>::optional_as_map(self.left).dom().contains(elem) ==> elem.le(self.key) && elem != self.key)
        &&& (forall |elem| #[trigger] Node::<K, V>::optional_as_map(self.right).dom().contains(elem) ==> self.key.le(elem) && elem != self.key)
        &&& (match self.left {
            Some(left_node) => left_node.well_formed(),
            None => true,
        })
        &&& (match self.right {
            Some(right_node) => right_node.well_formed(),
            None => true,
        })
    }
}

impl<K: TotalOrdered, V> TreeMap<K, V> {
    #[verifier::type_invariant]
    spec fn well_formed(self) -> bool {
        match self.root {
            Some(node) => node.well_formed(),
            None => true, // empty tree always well-formed
        }
    }
}

Meanwhile, the definition of as_map doesn’t rely on the ordering function, so it can be left alone, the same as before.

Updating the implementations and proofs

Updating the implementations take a bit more work, since we need more substantial proof code. Whereas Verus has good automation for integer inequalities (<), it has no such automation for our new, hand-made TotalOrdered trait. Thus, we need to add proof code to invoke its properties manually.

Let’s take a look at Node::get.

The meat of the proof roughly goes as follows:

Supoose we’re looking for the key key which compares less than self.key. Then we need to show that recursing into the left subtree gives the correct answer; for this, it suffices to show that key is not in the right subtree.

Suppose (for contradiction) that key were in the right subtree. Then (by the well-formedness invariant), we must have key > self.key. But we already established that key < self.key. Contradiction. (Formally, this contradiction can be obtained by invoking antisymmetry.)

impl<K: TotalOrdered, V> Node<K, V> {
    fn get_from_optional(node: &Option<Box<Node<K, V>>>, key: K) -> Option<&V>
        requires node.is_some() ==> node.unwrap().well_formed(),
        returns (match node {
            Some(node) => (if node.as_map().dom().contains(key) { Some(&node.as_map()[key]) } else { None }),
            None => None,
        }),
    {
        match node {
            None => None,
            Some(node) => {
                node.get(key)
            }
        }
    }

    fn get(&self, key: K) -> Option<&V>
        requires self.well_formed(),
        returns (if self.as_map().dom().contains(key) { Some(&self.as_map()[key]) } else { None })
    {
        match key.compare(&self.key) {
            Cmp::Equal => {
                Some(&self.value)
            }
            Cmp::Less => {
                proof {
                    if Node::<K, V>::optional_as_map(self.right).dom().contains(key) {
                        TotalOrdered::antisymmetric(self.key, key);
                        assert(false);
                    }
                    assert(key != self.key);
                    assert((match self.left {
                            Some(node) => (if node.as_map().dom().contains(key) { Some(&node.as_map()[key]) } else { None }),
                            None => None,
                        }) == (if self.as_map().dom().contains(key) { Some(&self.as_map()[key]) } else { None }));
                }
                Self::get_from_optional(&self.left, key)
            }
            Cmp::Greater => {
                proof {
                    if Node::<K, V>::optional_as_map(self.left).dom().contains(key) {
                        TotalOrdered::antisymmetric(self.key, key);
                        assert(false);
                    }
                    assert(key != self.key);
                    assert((match self.right {
                            Some(node) => (if node.as_map().dom().contains(key) { Some(&node.as_map()[key]) } else { None }),
                            None => None,
                        }) == (if self.as_map().dom().contains(key) { Some(&self.as_map()[key]) } else { None }));
                }
                Self::get_from_optional(&self.right, key)
            }
        }
    }
}

We can update insert and delete similarly, manually inserting lemma calls to invoke the total-ordering properties where necessary.

Full source

The full source for this example can be found here.