A simple binary search tree

In this section, we’re going to be implementing and verifying a Binary Search Tree (BST).

In the study of data structures, there are many known ways to balance a binary search tree. To keep things simple, we won’t be implementing any of them—instead, we’ll be implementing a straightforward, unbalanced binary search tree. Improving the design to be more efficient will be left as an exercise.

Furthermore, our first draft of an implementation is going to map keys of the fixed orderable type, u64, to values of type V. In a later chapter, we’ll change the keys to also be generic, thus mapping K to V for arbitrary types K and V.

The implementation

The structs

We’ll start by defining the tree shape itself, which contains one (key, value) pair at every node. We make no distinction between “leaf nodes” and “interior nodes”. Rather, every node has an optional left child and an optional right child. Furthermore, the tree might be entirely empty, in which case there is no root.

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

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

Note that only TreeMap is marked pub. Its field, root, as well as the Node type as a whole, are implementation details, and thus are private to the module.

The abstract view

When creating a new data structure, there are usually two important first steps:

Establish an interpretation of the data structure as some abstract datatype that will be used to write specifications.
Establish the well-formedness invariants of the data structure.

We’ll do the first one first (in part because it will actually help with the second one). In this case, we want to interpret the data structure as a Map<u64, V>. We can define such a function recursively.

impl<V> Node<V> {
    spec fn optional_as_map(node_opt: Option<Box<Node<V>>>) -> Map<u64, V>
        decreases node_opt,
    {
        match node_opt {
            None => Map::empty(),
            Some(node) => node.as_map(),
        }
    }

    spec fn as_map(self) -> Map<u64, V>
        decreases self,
    {
        Node::<V>::optional_as_map(self.left)
          .union_prefer_right(Node::<V>::optional_as_map(self.right))
          .insert(self.key, self.value)
    }
}

impl<V> TreeMap<V> {
    pub closed spec fn as_map(self) -> Map<u64, V> {
        Node::<V>::optional_as_map(self.root)
    }
}

Again note that only TreeMap::as_map is marked pub, and furthermore, that it’s marked closed. The definition of as_map is, again, an implementation detail.

It is customary to also implement the View trait as a convenience. This lets clients refer to the map implementation using the @ notation, e.g., tree_map@ as a shorthand for tree_map.view(). We’ll be writing our specifications in terms of tree_map.view().

impl<V> View for TreeMap<V> {
    type V = Map<u64, V>;

    open spec fn view(&self) -> Map<u64, V> {
        self.as_map()
    }
}

Establishing well-formedness

Next we establish well-formedness. This amounts to upholding the BST ordering property, namely, that for every node N, the nodes in N’s left subtree have keys less than N, while the nodes in N’s right subtree have keys greater than N. Again, this can be defined by a recursive spec function.

impl<V> Node<V> {
    spec fn well_formed(self) -> bool
        decreases self
    {
        &&& (forall |elem| Node::<V>::optional_as_map(self.left).dom().contains(elem) ==> elem < self.key)
        &&& (forall |elem| Node::<V>::optional_as_map(self.right).dom().contains(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<V> TreeMap<V> {
    pub closed spec fn well_formed(self) -> bool {
        match self.root {
            Some(node) => node.well_formed(),
            None => true, // empty tree always well-formed
        }
    }
}

Implementing a constructor: `TreeMap::new()`

Defining a constructor is simple; we create an empty tree with no root. The specification indicates that the returned object must represent the empty map.

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

Recall that tree_map@ is equivalent to tree_map.as_map(). An inspection of the definition of tree_map.as_map() and Node::optional_as_map() should make it apparent this will be the empty map when root is None.

Observe again that this specification does not refer to the tree internals at all, only that it is well-formed and that its abstract view is the empty map.

Implementing the `insert` operation

We can also implement insert using a recursive traversal. We search for the given node, using the well-formedness conditions to prove that we’re doing the right thing. During this traversal, we’ll either find a node with the right key, in which case we update the value, or we’ll reach a leaf without ever finding the desired node, in which case we create a new node.

(Aside: One slight snag has to do with a limitation of Verus’s handing of mutable references. Specifically, Verus doesn’t yet support an easy way to get a &mut T out of a &mut Option<T>. To get around this, we use std::mem::swap to get ownership of the node.)

impl<V> Node<V> {
    fn insert_into_optional(node: &mut Option<Box<Node<V>>>, key: u64, value: V)
        requires
            old(node).is_some() ==> old(node).unwrap().well_formed(),
        ensures
            node.is_some() ==> node.unwrap().well_formed(),
            Node::<V>::optional_as_map(*node) =~= Node::<V>::optional_as_map(*old(node)).insert(key, value)
    {
        if node.is_none() {
            *node = Some(Box::new(Node::<V> {
                key: key,
                value: value,
                left: None,
                right: None,
            }));
        } else {
            let mut tmp = None;
            std::mem::swap(&mut tmp, node);
            let mut boxed_node = tmp.unwrap();

            (&mut *boxed_node).insert(key, value);

            *node = Some(boxed_node);
        }
    }

    fn insert(&mut self, key: u64, value: V)
        requires
            old(self).well_formed(),
        ensures
            self.well_formed(),
            self.as_map() =~= old(self).as_map().insert(key, value),
    {
        if key == self.key {
            self.value = value;

            assert(!Node::<V>::optional_as_map(self.left).dom().contains(key));
            assert(!Node::<V>::optional_as_map(self.right).dom().contains(key));
        } else if key < self.key {
            Self::insert_into_optional(&mut self.left, key, value);

            assert(!Node::<V>::optional_as_map(self.right).dom().contains(key));
        } else {
            Self::insert_into_optional(&mut self.right, key, value);

            assert(!Node::<V>::optional_as_map(self.left).dom().contains(key));
        }
    }
}

impl<V> TreeMap<V> {
    pub fn insert(&mut self, key: u64, value: V)
        requires
            old(self).well_formed()
        ensures
            self.well_formed(),
            self@ == old(self)@.insert(key, value)
    {
        Node::<V>::insert_into_optional(&mut self.root, key, value);
    }
}

Observe that the specification of TreeMap::insert is given in terms of Map::insert.

Implementing the `delete` operation

Implementing delete is a little harder, because if we need to remove an interior node, we might have to reshape the tree a bit. However, since we aren’t trying to follow any particular balancing strategy, it’s still not that bad:

impl<V> Node<V> {
    fn delete_from_optional(node: &mut Option<Box<Node<V>>>, key: u64)
        requires
            old(node).is_some() ==> old(node).unwrap().well_formed(),
        ensures
            node.is_some() ==> node.unwrap().well_formed(),
            Node::<V>::optional_as_map(*node) =~= Node::<V>::optional_as_map(*old(node)).remove(key)
    {
        if node.is_some() {
            let mut tmp = None;
            std::mem::swap(&mut tmp, node);
            let mut boxed_node = tmp.unwrap();

            if key == boxed_node.key {
                assert(!Node::<V>::optional_as_map(boxed_node.left).dom().contains(key));
                assert(!Node::<V>::optional_as_map(boxed_node.right).dom().contains(key));

                if boxed_node.left.is_none() {
                    *node = boxed_node.right;
                } else {
                    if boxed_node.right.is_none() {
                        *node = boxed_node.left;
                    } else {
                        let (popped_key, popped_value) = Node::<V>::delete_rightmost(&mut boxed_node.left);
                        boxed_node.key = popped_key;
                        boxed_node.value = popped_value;
                        *node = Some(boxed_node);
                    }
                }
            } else if key < boxed_node.key {
                assert(!Node::<V>::optional_as_map(boxed_node.right).dom().contains(key));
                Node::<V>::delete_from_optional(&mut boxed_node.left, key);
                *node = Some(boxed_node);
            } else {
                assert(!Node::<V>::optional_as_map(boxed_node.left).dom().contains(key));
                Node::<V>::delete_from_optional(&mut boxed_node.right, key);
                *node = Some(boxed_node);
            }
        }
    }

    fn delete_rightmost(node: &mut Option<Box<Node<V>>>) -> (popped: (u64, V))
        requires
            old(node).is_some(),
            old(node).unwrap().well_formed(),
        ensures
            node.is_some() ==> node.unwrap().well_formed(),
            Node::<V>::optional_as_map(*node) =~= Node::<V>::optional_as_map(*old(node)).remove(popped.0),
            Node::<V>::optional_as_map(*old(node)).dom().contains(popped.0),
            Node::<V>::optional_as_map(*old(node))[popped.0] == popped.1,
            forall |elem| Node::<V>::optional_as_map(*old(node)).dom().contains(elem) ==> popped.0 >= elem,
    {
        let mut tmp = None;
        std::mem::swap(&mut tmp, node);
        let mut boxed_node = tmp.unwrap();

        if boxed_node.right.is_none() {
            *node = boxed_node.left;
            assert(Node::<V>::optional_as_map(boxed_node.right) =~= Map::empty());
            assert(!Node::<V>::optional_as_map(boxed_node.left).dom().contains(boxed_node.key));
            return (boxed_node.key, boxed_node.value);
        } else {
            let (popped_key, popped_value) = Node::<V>::delete_rightmost(&mut boxed_node.right);
            assert(!Node::<V>::optional_as_map(boxed_node.left).dom().contains(popped_key));
            *node = Some(boxed_node);
            return (popped_key, popped_value);
        }
    }
}

impl<V> TreeMap<V> {
    pub fn delete(&mut self, key: u64)
        requires old(self).well_formed()
        ensures
            self.well_formed(),
            self@ == old(self)@.remove(key)
    {
        Node::<V>::delete_from_optional(&mut self.root, key);
    }
}

Observe that the specification of TreeMap::delete is given in terms of Map::remove.

Implementing the `get` operation

Finally, we implement and verify TreeMap::get. This function looks up a key and returns an Option<&V> (None if the key isn’t in the TreeMap).

impl<V> Node<V> {
    fn get_from_optional(node: &Option<Box<Node<V>>>, key: u64) -> 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: u64) -> Option<&V>
        requires self.well_formed(),
        returns (if self.as_map().dom().contains(key) { Some(&self.as_map()[key]) } else { None })
    {
        if key == self.key {
            Some(&self.value)
        } else if key < self.key {
            proof { assert(!Node::<V>::optional_as_map(self.right).dom().contains(key)); }
            Self::get_from_optional(&self.left, key)
        } else {
            proof { assert(!Node::<V>::optional_as_map(self.left).dom().contains(key)); }
            Self::get_from_optional(&self.right, key)
        }
    }
}

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

Using the `TreeMap` as a client

A short client program illustrates how we can reason about the TreeMap as if it were a Map.

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));
}

Full source

The full source for this example can be found here.

Verus Tutorial and Reference