Mutable references in a container

In this chapter, we’ll look at a client application using a TreeMap<u64, &mut u64>. Putting mutable references in a container like this is kind of uncommon, but not unheard of. In Verus, it takes a little additional know-how.

A detour through Vec

Before looking at TreeMap, though, let’s do an example with Vec to review some of the theory. Here, we use a vector of mutable references as a way of “indexing” into a triple of local variables.

fn vec_with_mut_refs(i: usize)
    requires 0 <= i < 3
{
    let mut a = 0;
    let mut b = 0;
    let mut c = 0;

    let mut v = vec![&mut a, &mut b, &mut c];

    *v[i] = 1;

    assert(i == 0 ==> (a, b, c) == (1, 0, 0));
    assert(i == 1 ==> (a, b, c) == (0, 1, 0));
    assert(i == 2 ==> (a, b, c) == (0, 0, 1));
}

Unfortunately, all 3 assertions will fail to verify, although they are obviously true. Why is this?

Typically, when we have a mutable reference stored in a local variable (say, a: &mut u64), Verus will identify the last mutation to a and insert an assumption of the predicate, has_resolved(a). This is called the resolution assumption and this assumption is necessary for the updated value to “flow back” to the borrowed-from location.

Okay, so what happens we have a vector? In this case, the local variable is v: Vec<&mut u64>, so Verus will determine the resolution point for v. Immediately after *v[i] = 1 we have the resolution point for v and learn has_resolved(v). When applied to a container like this, has_resolved is recursive: it means that all mutable references contained within v are resolved.

So we have that v is resolved; how do we then get that v[0], v[1], and v[2] are resolved?

In this case, vstd provides the axiom axiom_vec_has_resolved:

pub broadcast axiom fn axiom_vec_has_resolved<T>(vec: Vec<T>, i: int)
    ensures
        0 <= i < vec.len() ==> has_resolved::<Vec<T>>(vec) ==> has_resolved(vec@[i]);

Thus, this law implies that has_resolved(v[0]), which should give us the updated value of a; that has_resolved(v[1]), which should give us the updated value of b; and that has_resolved(v[2]), which should give us the updated value of c.

To get the above program to verify, we just need to invoke the resolution axiom, which can be done either by calling it directly or simply triggering the quantifier:

fn vec_with_mut_refs(i: usize)
    requires 0 <= i < 3
{
    let mut a = 0;
    let mut b = 0;
    let mut c = 0;

    let mut v = vec![&mut a, &mut b, &mut c];

    *v[i] = 1;

    assert(has_resolved(v[0]));
    assert(has_resolved(v[1]));
    assert(has_resolved(v[2]));

    assert(i == 0 ==> (a, b, c) == (1, 0, 0));
    assert(i == 1 ==> (a, b, c) == (0, 1, 0));
    assert(i == 2 ==> (a, b, c) == (0, 0, 1));
}

A resolution lemma for TreeMap

Now, let’s talk about TreeMap<u64, &mut u64>.

Using what we learned in the above section, we know that making use of a container of mutable references requires some kind of has_resolved law. Thus, our TreeMap library should export a lemma like:

has_resolved(map) ==> map@.dom().contains(k) ==> has_resolved(map@[k])

How do we prove this?

Verus automatically emits axioms about has_resolved for individual structs and enums (including user-defined ones). For example, if we look at our TreeMap struct:

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

Verus will automatically emit an axiom that:

has_resolved::<TreeMap<K, V>>(map) ==>
    has_resolved::<Option<Box<Node<K, V>>>>(map.root)

(at least, when the field is visible). Likewise, there’s an axiom for Option that will give us has_resolved of the contained value for the Some case:

has_resolved::<Option<T>>(option) ==>
    (option matches Some(x) ==> has_resolved::<T>(x))

and finally for Box:

has_resolved::<Box<U>>(box) ==>
    has_resolved::<U>(*box)

So chaining these all together, we can get has_resolved::<TreeMap<K, V>>(map) ==> has_resolved::<Node<K, V>>(*map.root.unwrap()). Then we can similarly recurse within the Node struct on the left and right fields until we reach a value V. We just need to put this all together into an inductive proof:

proof fn lemma_node_has_resolved<K: TotalOrdered, V>(node: Node<K, V>, k: K)
    requires
        has_resolved(node),
        node.as_map().dom().contains(k),
    ensures
        has_resolved(node.as_map()[k])
    decreases node,
{
    if node.left.is_some() && node.left.unwrap().as_map().dom().contains(k) {
        lemma_node_has_resolved(*node.left.unwrap(), k);
    }
    if node.right.is_some() && node.right.unwrap().as_map().dom().contains(k) {
        lemma_node_has_resolved(*node.right.unwrap(), k);
    }
}

proof fn lemma_tree_map_has_resolved<K: TotalOrdered, V>(map: TreeMap<K, V>, k: K)
    requires
        has_resolved(map),
        map@.dom().contains(k),
    ensures
        has_resolved(map@[k])
{
    match map.root {
        Some(node) => {
            lemma_node_has_resolved(*node, k);
        }
        None => { }
    }
}

Now let’s try it out with a client:

fn tree_map_with_mut_refs(i: u64)
    requires 0 <= i < 3
{
    let mut a = 0;
    let mut b = 0;
    let mut c = 0;

    let mut tree_map = TreeMap::<u64, &mut u64>::new();
    tree_map.insert(0, &mut a);
    tree_map.insert(1, &mut b);
    tree_map.insert(2, &mut c);

    **tree_map.get_mut(i).unwrap() = 1;

    proof {
        lemma_tree_map_has_resolved(tree_map, 0);
        lemma_tree_map_has_resolved(tree_map, 1);
        lemma_tree_map_has_resolved(tree_map, 2);
    }

    assert(i == 0 ==> (a, b, c) == (1, 0, 0));
    assert(i == 1 ==> (a, b, c) == (0, 1, 0));
    assert(i == 2 ==> (a, b, c) == (0, 0, 1));
}