Encapsulating well-formedness with type invariants

Recall our specifications from the previous chapter:

impl<V> TreeMap<V> {
    pub fn new() -> (tree_map: Self)
        ensures
            tree_map.well_formed(),
            tree_map@ == Map::<u64, V>::empty()

    pub fn insert(&mut self, key: u64, value: V)
        requires
            old(self).well_formed()
        ensures
            self.well_formed(),
            self@ == old(self)@.insert(key, value)

    pub fn delete(&mut self, key: u64)
        requires old(self).well_formed()
        ensures
            self.well_formed(),
            self@ == old(self)@.remove(key)

    pub fn get(&self, key: u64) -> Option<&V>
        requires self.well_formed(),
        returns (if self@.dom().contains(key) { Some(&self@[key]) } else { None })
}

Observe the presenence of this tree_map.well_formed() predicate, especially in the requires clauses. As a result of this, the client needs to work with this tree_map.well_formed() predicate all throughout their own code. For example:

fn test2() {
    let mut tree_map = TreeMap::<bool>::new();
    test_callee(tree_map);
}

fn test_callee(tree_map: TreeMap<bool>)
    requires tree_map.well_formed(),
{
    let mut tree_map = tree_map;
    tree_map.insert(25, true);
    tree_map.insert(100, true);
}

Without the requires clause, the above snippet would fail to verify.

Intuitively, however, one might wonder why we have to carry this predicate around at all. After all, due to encapsulation, it isn’t ever possible for the client to create a tree_map where well_formed() doesn’t hold.

In this section, we’ll show how to use Verus’s type invariants feature to remedy this burden from the client.

Applying the `type_invariant` attribute.

In order to tell Verus that we want the well_formed() predicate to be inferred automatically, we can mark it with the #[verifier::type_invariant] attribute:

impl<V> TreeMap<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
        }
    }
}

This has two effects:

It adds an extra constraint that all TreeMap objects satsify the well_formed() condition at all times. This constraint is checked by Verus whenever a TreeMap object is constructed or modified.
It allows the programmer to assume the well_formed() condition at all times, even when it isn’t present in a requires clause.

Note that in addition to adding the type_invariant attribute, we have also removed the pub specifier from well_formed. Now not only is the body invisible to the client, even the name is as well. After all, our intent is to prevent the client from needing to reason about it, at which point there is no reason to expose it through the public interface at all.

Of course, for this to be possible, we’ll need to update the specifications of TreeMap’s various pub methods.

Updating the code: `new`

Let’s start with an easy one: new.

impl<V> TreeMap<V> {
    pub fn new() -> (s: Self)
        ensures
            s@ == Map::<u64, V>::empty()
    {
        TreeMap::<V> { root: None }
    }
}

All we’ve done here is remove the s.well_formed() postcondition, which as discussed, is no longer necessary.

Crucially, Verus still requires us to prove that s.well_formed() holds. Specifically, since well_formed has been marked with #[verifier::type_invariant], Verus checks that well_formed() holds when the TreeMap constructor returns. As before, Verus can check this condition fairly trivially.

Updating the code: `get`

Now let’s take a look at get. The first thing to notice is that we remove the requires self.well_formed() clause.

impl<V> TreeMap<V> {
    pub fn get(&self, key: u64) -> Option<&V>
        returns (if self@.dom().contains(key) { Some(&self@[key]) } else { None })
    {
        proof { use_type_invariant(&*self); }
        Node::<V>::get_from_optional(&self.root, key)
    }
}

Given that we no longer have the precondition, how do we deduce self.well_formed() (which is needed to prove self.root is well-formed and call Node::get_from_optional)?

This can be done with the built-in pseudo-lemma use_type_invariant. When called on any object, this feature guarantees that the provided object satisfies its type invariants.

Updating the code: `insert`

Now let’s check TreeMap::insert, which if you recall, has to modify the tree.

impl<V> TreeMap<V> {
    pub fn insert(&mut self, key: u64, value: V)
        ensures
            self@ == old(self)@.insert(key, value)
    {
        proof { use_type_invariant(&*self); }
        let mut root = None;
        std::mem::swap(&mut root, &mut self.root);
        Node::<V>::insert_into_optional(&mut root, key, value);
        self.root = root;
    }
}

As before, we use use_type_invariant to establish that self.well_formed() holds at the beginning of the function, even without the requires clause.

One slight challenge that arises from the use of #[verifier::type_invariant] is that it enforces type invariants to hold at every program point. Sometimes, this can make intermediate computation a little tricky.

In this case, an easy way to get around this is to swap the root field with None, then swap back when we’re done. This works because the empty TreeMap trivially satisfies the well-formedness, so it’s allowed.

One might wonder why we can’t just do Node::<V>::insert_into_optional(&mut self.root, key, value) without swapping. The trouble with this is that it requires us to ensure the call to insert_into_optional is “unwind-safe”, i.e., that all type invariants would be preserved even if a panic occurs and insert_into_optional has to exit early. Right now, Verus only has one way to ensure unwind-safety, which is to bluntly ensure that no unwinding happens at all. Thus, the ideal solution would be to mark insert_into_optional as no_unwind. However, this is impossible in this case, because node insertion will call Box::new.

Between this problem, and Verus’s current limitations regarding unwind-safety, the swap approach becomes the easiest solution as a way of sidestepping it. Check the reference page for more information on the limitations of the type_invariant feature.

Updating the code: `delete`

This is pretty much the same as insert.

impl<V> TreeMap<V> {
    pub fn delete(&mut self, key: u64)
        ensures
            self@ == old(self)@.remove(key)
    {
        proof { use_type_invariant(&*self); }
        let mut root = None;
        std::mem::swap(&mut root, &mut self.root);
        Node::<V>::delete_from_optional(&mut root, key);
        self.root = root;
    }
}

The new signatures and specifications

Putting it all together, we end up with the following specifications for our public API:

impl<V> TreeMap<V> {
    pub fn new() -> (s: Self)
        ensures
            s@ == Map::<u64, V>::empty()

    pub fn insert(&mut self, key: u64, value: V)
        ensures
            self@ == old(self)@.insert(key, value)

    pub fn delete(&mut self, key: u64)
        ensures
            self@ == old(self)@.remove(key)

    pub fn get(&self, key: u64) -> Option<&V>
        returns (if self@.dom().contains(key) { Some(&self@[key]) } else { None })
}

These are almost the same as what we had before; the only difference is that all the well_formed() clauses have been removed.

Conveniently, there are no longer any requires clause at all, so it’s always possible to call any of these functions. This is also important if we want to prove the API “safe” in the Rust sense (see this page).

The new client code

As before, the client code gets to reason about the TreeMap as if it were just a Map. Now, however, it’s a bit simpler because we don’t have to reason about tree_map.well_formed().

fn test() {
    let mut tree_map = TreeMap::<bool>::new();
    tree_map.insert(17, false);
    tree_map.insert(18, false);
    tree_map.insert(17, true);

    assert(tree_map@ == map![17u64 => true, 18u64 => false]);

    tree_map.delete(17);

    assert(tree_map@ == map![18u64 => false]);

    let elem17 = tree_map.get(17);
    let elem18 = tree_map.get(18);
    assert(elem17.is_none());
    assert(elem18 == Some(&false));

    test2(tree_map);
}

fn test2(tree_map: TreeMap<bool>) {
    let mut tree_map = tree_map;
    tree_map.insert(25, true);
    tree_map.insert(100, true);
}

Full source

The full source for this example can be found here.

Verus Tutorial and Reference