From efc053896816de28e11fb07edbe619440cbe6b35 Mon Sep 17 00:00:00 2001 From: Travis Hance Date: Thu, 14 Nov 2024 14:43:16 -0500 Subject: [PATCH] add guide chapter on Binary Search Tree illustrates various aspects of modularity / building a container data structure - view - well_formed() / type_invariants - genericity using traits - implementing Clone generically --- source/docs/guide/src/SUMMARY.md | 7 + source/docs/guide/src/container_bst.md | 7 + .../guide/src/container_bst_all_source.md | 23 + source/docs/guide/src/container_bst_clone.md | 109 ++++ .../guide/src/container_bst_first_draft.md | 148 +++++ .../docs/guide/src/container_bst_generic.md | 83 +++ .../guide/src/container_bst_type_invariant.md | 168 ++++++ source/rust_verify/example/guide/bst_map.rs | 334 +++++++++++ .../example/guide/bst_map_generic.rs | 547 ++++++++++++++++++ .../example/guide/bst_map_type_invariant.rs | 324 +++++++++++ 10 files changed, 1750 insertions(+) create mode 100644 source/docs/guide/src/container_bst.md create mode 100644 source/docs/guide/src/container_bst_all_source.md create mode 100644 source/docs/guide/src/container_bst_clone.md create mode 100644 source/docs/guide/src/container_bst_first_draft.md create mode 100644 source/docs/guide/src/container_bst_generic.md create mode 100644 source/docs/guide/src/container_bst_type_invariant.md create mode 100644 source/rust_verify/example/guide/bst_map.rs create mode 100644 source/rust_verify/example/guide/bst_map_generic.rs create mode 100644 source/rust_verify/example/guide/bst_map_type_invariant.rs diff --git a/source/docs/guide/src/SUMMARY.md b/source/docs/guide/src/SUMMARY.md index c0814d1ca..ed23b4974 100644 --- a/source/docs/guide/src/SUMMARY.md +++ b/source/docs/guide/src/SUMMARY.md @@ -82,6 +82,13 @@ - [Pointers](./pointers.md) - [Concurrency](concurrency.md) +- [Verifying a container library: Binary Search Tree](./container_bst.md) + - [First draft](./container_bst_first_draft.md) + - [Encapsulating well-formedness with type invariants](./container_bst_type_invariant.md) + - [Making it generic](./container_bst_generic.md) + - [Implementing `Clone`](./container_bst_clone.md) + - [Full source for examples](./container_bst_all_source.md) + - [Understanding the guarantees of a verified program](./guarantees.md) - [Assumptions and trusted components](./tcb.md) - [Memory safety is conditional on verification](./memory-safety.md) diff --git a/source/docs/guide/src/container_bst.md b/source/docs/guide/src/container_bst.md new file mode 100644 index 000000000..401f0a127 --- /dev/null +++ b/source/docs/guide/src/container_bst.md @@ -0,0 +1,7 @@ +# Verifying a container library + +In this section, we'll learn how to verify a simple container library, specifically, +via an example of a _map_ data structure using a binary search tree. +In the case study, we'll explore various considerations for +writing a modular specification +that encapsulates verfication details as well as implementation details. diff --git a/source/docs/guide/src/container_bst_all_source.md b/source/docs/guide/src/container_bst_all_source.md new file mode 100644 index 000000000..da8878094 --- /dev/null +++ b/source/docs/guide/src/container_bst_all_source.md @@ -0,0 +1,23 @@ +# Full source for examples + + * [First draft](#first-draft) + * [Version with type invariants](#version-with-type-invariants) + * [Version with generic key type and Clone implementation](#version-with-generic-key-type-and-clone-implementation) + +## First draft + +```rust +{{#include ../../../rust_verify/example/guide/bst_map.rs:all}} +``` + +## Version with type invariants + +```rust +{{#include ../../../rust_verify/example/guide/bst_map_type_invariant.rs:all}} +``` + +## Version with generic key type and Clone implementation + +```rust +{{#include ../../../rust_verify/example/guide/bst_map_generic.rs:all}} +``` diff --git a/source/docs/guide/src/container_bst_clone.md b/source/docs/guide/src/container_bst_clone.md new file mode 100644 index 000000000..5f9358597 --- /dev/null +++ b/source/docs/guide/src/container_bst_clone.md @@ -0,0 +1,109 @@ +# Implementing Clone + +As a finishing touch, let's implement `Clone` for `TreeMap`. +The main trick here will be in figuring out the correct specification for `TreeMap::::Clone`. + +Naturally, such an implementation will require both `K: Clone` and `V: Clone`. +However, to write a sensible clone implementation for the tree, we have to consider +what the implementations of `K::clone` and `V::clone` actually do. + +Generally speaking, Verus imposes no constraints on the implementations of the `Clone`, +so it is not necessarily true that a `clone()` call will return a value that is spec-equal +to its inputs. + +With this in mind, we're going to prove the following signature for `TreeMap::clone`: + +```rust +{{#include ../../../rust_verify/example/guide/bst_map_generic.rs:clone_signature}} + { + ... + } +} +``` + +### Dealing with `K::clone` + +In order to clone all the keys, we need `K::clone` to respect the ordering of elements; otherwise, +we'd need to resort all the keys on clone in order for the resulting tree to be valid. +However, it's not likely that is desirable behavior. If `Clone` doesn't respect the +`TotalOrdered` implementation, it's likely a user bug. + +A general way to handle this would be to require that `Clone` actually be compatible +with the total-ordering in some sense. +However, you'll +recall from the previous section that we're already simplifying the "total ordered" specification +a bit. Likewise, we're going to continue to keep things simple here by also requiring +that `K: Copy`. + +As a result, we'll be able to prove that our `TreeMap` clone implementation can preserve +all keys exactly, even when compared via spec equality. That is, we'll be able to +ensure that `self@.dom() =~= res@.dom()`. + +### Dealing with `V::clone` + +So what about `V`? Again, we don't know _a priori_ what `V::clone` does. It might return +values unequal to the imput; it might even be nondeterminstic. Therefore, +a cloned `TreeMap` may have different values than the original. + +In order to specify `TreeMap::::clone` as generically as possible, we choose +to write its ensures clause _in terms of_ the ensures clause for `V::clone`. +This can be done using `[`call_ensures`](./exec_funs_as_values.md)`. +The predicate `call_ensures(V::clone, (&self@[key],), res@[key])` effectively says +"`self@[key]` and `res@[key]` are a possible input-output pair for `V::clone`". + +### Understanding the implications of the signature + +Let's do a few examples. + +First, consider cloning a `TreeMap::`. The Verus standard library provides +a specification for `u32::clone`; it's the same as a copy, i.e., a cloned `u32` always +equals the input. As a result, we can deduce that cloning a `TreeMap::` will +preserve its `view` exactly. We can prove this using [extensional equality](./extensional_equality.md). + +```rust +{{#include ../../../rust_verify/example/guide/bst_map_generic.rs:clone_u32}} +``` + +We can do the same for _any_ type where `clone` guarantees spec-equality. Here's another +example with a user-defined type. + +```rust +{{#include ../../../rust_verify/example/guide/bst_map_generic.rs:clone_int_wrapper}} +``` + +This works because of the postcondition on `IntWrapper::clone`, that is, `ensures *s == self`. +If you're new to this style, it might seem initially surprising that +`IntWrapper::clone` has any effect on the verification of `test_clone_int_wrapper`, since +it doesn't directly call `IntWrapper::clone`. In this case, the postcondition is referenced +indirectly via `TreeMap:clone`. + +Let's do one more example, this time with a _less_ precise clone function. + +```rust +{{#include ../../../rust_verify/example/guide/bst_map_generic.rs:clone_weird_int}} +``` + +This example is a bit pathological; our struct, `WeirdInt`, has an extra field that doesn't +get cloned. You could imagine real-life scenarios that have this property (for example, +if every struct needs to have a unique identifier). Anyway, the postcondition of +`WeirdInt::clone` doesn't say both objects are equal, only that the `int_value` fields are equal. +This postcondition can then be inferred for each value in the map, as shown. + +### Implementing `TreeMap::::Clone`. + +As usual, we write the implementation as a recursive function. + +It's not necessary to implement `Node::Clone`; one could instead just implement a normal +recursive function as a helper for `TreeMap::Clone`; but it's more Rust-idiomatic to do it +this way. This lets us call `Option>::Clone` +in the implementation of `TreeMap::clone` (the spec for `Option::clone` is provided by +vstd). However, you can see that there are a few 'gotchas' that need +to be worked around. + +```rust +{{#include ../../../rust_verify/example/guide/bst_map_generic.rs:clone_full_impl}} +``` + +## Full source + +The full source for this example can be found [here](./container_bst_all_source.md#version-with-generic-key-type-and-clone-implementation). diff --git a/source/docs/guide/src/container_bst_first_draft.md b/source/docs/guide/src/container_bst_first_draft.md new file mode 100644 index 000000000..260e6330e --- /dev/null +++ b/source/docs/guide/src/container_bst_first_draft.md @@ -0,0 +1,148 @@ +# 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](https://en.wikipedia.org/wiki/Red%E2%80%93black_tree) +[known](https://en.wikipedia.org/wiki/AVL_tree) +[ways](https://en.wikipedia.org/wiki/Treap) +[to](https://en.wikipedia.org/wiki/Splay_tree) +[balance](https://en.wikipedia.org/wiki/B-tree) +a binary search tree. +To keep things simple, we won't be implementing of them—instead, +we'll be implemented 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 be mapping 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. + +```rust +{{#include ../../../rust_verify/example/guide/bst_map.rs:StructsDef}} +``` + +Note that only `TreeMap` is marked `pub`. Its field, `root`, as well as the `Node` type +as a whole, are implementation details, thus private to the module. + +### The abstract view + +When creating a new data structure, there are usually two things to do first: + + * 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). +In this case, we want to interpret the data structure as a +[`Map`](https://verus-lang.github.io/verus/verusdoc/vstd/map/struct.Map.html). +We can define such a function recursively. + +```rust +{{#include ../../../rust_verify/example/guide/bst_map.rs:AsMapDef}} +``` + +Again note that only `Tree::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](https://verus-lang.github.io/verus/verusdoc/vstd/view/trait.View.html) +as a convenience. This lets client 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()`. + +```rust +{{#include ../../../rust_verify/example/guide/bst_map.rs:ViewDef}} +``` + +### 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. + +```rust +{{#include ../../../rust_verify/example/guide/bst_map.rs:WellFormedDef}} +``` + +### 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. + +```rust +{{#include ../../../rust_verify/example/guide/bst_map.rs:new}} +``` + +Recall that `tree_map@` is equivalent to `tree_map.as_map()`. +An inspection the definition of `tree_map.as_map()` and `Node::optional_as_map()` should +make it apparent this will be the empty set 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 to get a +`&mut T` out of a `&mut Option`. To get around this, we use [`std::mem::swap`](https://doc.rust-lang.org/std/mem/fn.swap.html) to get ownership of the node.) + +```rust +{{#include ../../../rust_verify/example/guide/bst_map.rs:insert}} +``` + +Observe that the specification of `TreeMap::insert` can be given in terms of +[`Map::insert`](https://verus-lang.github.io/verus/verusdoc/vstd/map/struct.Map.html#method.remove). + +### 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: + +```rust +{{#include ../../../rust_verify/example/guide/bst_map.rs:delete}} +``` + +Observe that the specification of `TreeMap::delete` can be given in terms of +[`Map::remove`](https://verus-lang.github.io/verus/verusdoc/vstd/map/struct.Map.html#method.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`). + +```rust +{{#include ../../../rust_verify/example/guide/bst_map.rs:get}} +``` + +### Using the `TreeMap` as a client + +A short client program illustrates how we can reason about the `TreeMap` as if it were +a [`Map`](https://verus-lang.github.io/verus/verusdoc/vstd/map/struct.Map.html). + +```rust +{{#include ../../../rust_verify/example/guide/bst_map.rs:test}} +``` + +## Full source + +The full source for this example can be found [here](./container_bst_all_source.md#first-draft). diff --git a/source/docs/guide/src/container_bst_generic.md b/source/docs/guide/src/container_bst_generic.md new file mode 100644 index 000000000..71c214713 --- /dev/null +++ b/source/docs/guide/src/container_bst_generic.md @@ -0,0 +1,83 @@ +# Making it generic + +In the previous sections, we devised a `TreeMap` which a fixed key type (`u64`). +In this section, we'll show to make a `TreeMap` 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, for obtaining 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"](https://en.wikipedia.org/wiki/Total_order) +generically by creating a trait. + +```rust +{{#include ../../../rust_verify/example/guide/bst_map_generic.rs:trait}} +``` + +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" is the comparison +function is the same as [spec equality](./equality.md). +This isn't always suitable; some datatypes may have more that 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`](https://verus-lang.github.io/verus/verusdoc/vstd/hash_map/struct.HashMapWithView.html) 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`. + +```rust +{{#include ../../../rust_verify/example/guide/bst_map_generic.rs:structs}} +``` + +We'll also update the well-formedness condition use the generic `K::le` instead of integer `<=`. +Where the original definition used `a < b`, we now use `a.le(b) && a != b`. + +```rust +{{#include ../../../rust_verify/example/guide/bst_map_generic.rs: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.) + +```rust +{{#include ../../../rust_verify/example/guide/bst_map_generic.rs:node_get}} +``` + +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](./container_bst_all_source.md#version-with-generic-key-type-and-clone-implementation). diff --git a/source/docs/guide/src/container_bst_type_invariant.md b/source/docs/guide/src/container_bst_type_invariant.md new file mode 100644 index 000000000..ee1fef238 --- /dev/null +++ b/source/docs/guide/src/container_bst_type_invariant.md @@ -0,0 +1,168 @@ +# Encapsulating well-formedness with type invariants + +Recall our specifications from the previous chapter: + +```rust +impl TreeMap { +{{#include ../../../rust_verify/example/guide/bst_map.rs:new_signature}} + +{{#include ../../../rust_verify/example/guide/bst_map.rs:insert_signature}} + +{{#include ../../../rust_verify/example/guide/bst_map.rs:delete_signature}} + +{{#include ../../../rust_verify/example/guide/bst_map.rs:get_signature}} +} +``` + +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: + +```rust +{{#include ../../../rust_verify/example/guide/bst_map.rs:test_callee}} +``` + +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 true. + +In this section, we'll show how to use Verus's [type invariants](./reference-type-invariants.md) +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: + +```rust +{{#include ../../../rust_verify/example/guide/bst_map_type_invariant.rs:well_formed_with_attr}} +``` + +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 point is exposing 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`. + +{{#include ../../../rust_verify/example/guide/bst_map_type_invariant.rs:get}} + +All we've done here is remove the `s.well_formed()` postcondition, which as discussed, +is no longer going to be 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 we invoke the `TreeMap` constructor. +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. + +{{#include ../../../rust_verify/example/guide/bst_map_type_invariant.rs:get}} + +Given that we no longer have the precondition, how _do_ we deduce `sell.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. + +```rust +{{#include ../../../rust_verify/example/guide/bst_map_type_invariant.rs:insert}} +``` + +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`](https://doc.rust-lang.org/std/mem/fn.swap.html) 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::::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`](./reference-unwind-sig.md). 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](./reference-type-invariants.md) for more information on +the limitations of the `type_invariant` feature. + +### Updating the code: `delete` + +This is pretty much the same as `insert`. + +```rust +{{#include ../../../rust_verify/example/guide/bst_map_type_invariant.rs:delete}} +``` + +### The new signatures and specifications + +Putting it all together, we end up with the following specifications for our public API: + +```rust +impl TreeMap { +{{#include ../../../rust_verify/example/guide/bst_map_type_invariant.rs:new_signature}} + +{{#include ../../../rust_verify/example/guide/bst_map_type_invariant.rs:insert_signature}} + +{{#include ../../../rust_verify/example/guide/bst_map_type_invariant.rs:delete_signature}} + +{{#include ../../../rust_verify/example/guide/bst_map_type_invariant.rs:get_signature}} +} +``` + +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](./memory-safety.md)). + +### The new client code + +As before, the client code gets to reason about the `TreeMap` as if it were just a +[`Map`](https://verus-lang.github.io/verus/verusdoc/vstd/map/struct.Map.html). +Now, however, it's a bit simpler because we don't have to reason about `tree_map.well_formed()`. + +```rust +{{#include ../../../rust_verify/example/guide/bst_map_type_invariant.rs:example_use}} +``` + +## Full source + +The full source for this example can be found [here](./container_bst_all_source.md#version-with-type-invariants). diff --git a/source/rust_verify/example/guide/bst_map.rs b/source/rust_verify/example/guide/bst_map.rs new file mode 100644 index 000000000..79629952c --- /dev/null +++ b/source/rust_verify/example/guide/bst_map.rs @@ -0,0 +1,334 @@ +// ANCHOR: all +use vstd::prelude::*; + +verus!{ + +// ANCHOR: StructsDef +struct Node { + key: u64, + value: V, + left: Option>>, + right: Option>>, +} + +pub struct TreeMap { + root: Option>>, +} +// ANCHOR_END: StructsDef + +// ANCHOR: AsMapDef +impl Node { + spec fn optional_as_map(node_opt: Option>>) -> Map + decreases node_opt, + { + match node_opt { + None => Map::empty(), + Some(node) => node.as_map(), + } + } + + spec fn as_map(self) -> Map + decreases self, + { + Node::::optional_as_map(self.left) + .union_prefer_right(Node::::optional_as_map(self.right)) + .insert(self.key, self.value) + } +} + +impl TreeMap { + pub closed spec fn as_map(self) -> Map { + Node::::optional_as_map(self.root) + } +} +// ANCHOR_END: AsMapDef + +// ANCHOR: ViewDef +impl View for TreeMap { + type V = Map; + + open spec fn view(&self) -> Map { + self.as_map() + } +} +// ANCHOR_END: ViewDef + +// ANCHOR: WellFormedDef +impl Node { + spec fn well_formed(self) -> bool + decreases self + { + &&& (forall |elem| Node::::optional_as_map(self.left).dom().contains(elem) ==> elem < self.key) + &&& (forall |elem| Node::::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 TreeMap { + pub closed spec fn well_formed(self) -> bool { + match self.root { + Some(node) => node.well_formed(), + None => true, // empty tree always well-formed + } + } +} +// ANCHOR_END: WellFormedDef + +// ANCHOR: new +impl TreeMap { +// ANCHOR: new_signature + pub fn new() -> (tree_map: Self) + ensures + tree_map.well_formed(), + tree_map@ == Map::::empty() +// ANCHOR_END: new_signature + { + TreeMap:: { root: None } + } +} +// ANCHOR_END: new + +// ANCHOR: insert +impl Node { + fn insert_into_optional(node: &mut Option>>, key: u64, value: V) + requires + old(node).is_some() ==> old(node).unwrap().well_formed(), + ensures + node.is_some() ==> node.unwrap().well_formed(), + Node::::optional_as_map(*node) =~= Node::::optional_as_map(*old(node)).insert(key, value) + { + if node.is_none() { + *node = Some(Box::new(Node:: { + 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::::optional_as_map(self.left).dom().contains(key)); + assert(!Node::::optional_as_map(self.right).dom().contains(key)); + } else if key < self.key { + Self::insert_into_optional(&mut self.left, key, value); + + assert(!Node::::optional_as_map(self.right).dom().contains(key)); + } else { + Self::insert_into_optional(&mut self.right, key, value); + + assert(!Node::::optional_as_map(self.left).dom().contains(key)); + } + } +} + +impl TreeMap { +// ANCHOR: insert_signature + pub fn insert(&mut self, key: u64, value: V) + requires + old(self).well_formed() + ensures + self.well_formed(), + self@ == old(self)@.insert(key, value) +// ANCHOR_END: insert_signature + { + Node::::insert_into_optional(&mut self.root, key, value); + } +} +// ANCHOR_END: insert + +// ANCHOR: delete +impl Node { + fn delete_from_optional(node: &mut Option>>, key: u64) + requires + old(node).is_some() ==> old(node).unwrap().well_formed(), + ensures + node.is_some() ==> node.unwrap().well_formed(), + Node::::optional_as_map(*node) =~= Node::::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::::optional_as_map(boxed_node.left).dom().contains(key)); + assert(!Node::::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::::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::::optional_as_map(boxed_node.right).dom().contains(key)); + Node::::delete_from_optional(&mut boxed_node.left, key); + *node = Some(boxed_node); + } else { + assert(!Node::::optional_as_map(boxed_node.left).dom().contains(key)); + Node::::delete_from_optional(&mut boxed_node.right, key); + *node = Some(boxed_node); + } + } + } + + fn delete_rightmost(node: &mut Option>>) -> (popped: (u64, V)) + requires + old(node).is_some(), + old(node).unwrap().well_formed(), + ensures + node.is_some() ==> node.unwrap().well_formed(), + Node::::optional_as_map(*node) =~= Node::::optional_as_map(*old(node)).remove(popped.0), + Node::::optional_as_map(*old(node)).dom().contains(popped.0), + Node::::optional_as_map(*old(node))[popped.0] == popped.1, + forall |elem| Node::::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::::optional_as_map(boxed_node.right) =~= Map::empty()); + assert(!Node::::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::::delete_rightmost(&mut boxed_node.right); + assert(!Node::::optional_as_map(boxed_node.left).dom().contains(popped_key)); + *node = Some(boxed_node); + return (popped_key, popped_value); + } + } +} + +impl TreeMap { +// ANCHOR: delete_signature + pub fn delete(&mut self, key: u64) + requires old(self).well_formed() + ensures + self.well_formed(), + self@ == old(self)@.remove(key) +// ANCHOR_END: delete_signature + { + Node::::delete_from_optional(&mut self.root, key); + } +} +// ANCHOR_END: delete + +// ANCHOR: get +impl Node { + fn get_from_optional(node: &Option>>, 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::::optional_as_map(self.right).dom().contains(key)); } + Self::get_from_optional(&self.left, key) + } else { + proof { assert(!Node::::optional_as_map(self.left).dom().contains(key)); } + Self::get_from_optional(&self.right, key) + } + } +} + +impl TreeMap { +// ANCHOR: get_signature + pub fn get(&self, key: u64) -> Option<&V> + requires self.well_formed(), + returns (if self@.dom().contains(key) { Some(&self@[key]) } else { None }) +// ANCHOR_END: get_signature + { + Node::::get_from_optional(&self.root, key) + } +} +// ANCHOR_END: get + +// ANCHOR: test +fn test() { + let mut tree_map = TreeMap::::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)); +} +// ANCHOR_END: test + +// ANCHOR: test_callee +fn test2() { + let mut tree_map = TreeMap::::new(); + test_callee(tree_map); +} + +fn test_callee(tree_map: TreeMap) + requires tree_map.well_formed(), +{ + let mut tree_map = tree_map; + tree_map.insert(25, true); + tree_map.insert(100, true); +} +// ANCHOR_END: test_callee + + +} +// ANCHOR_END: all + +fn main() { } + diff --git a/source/rust_verify/example/guide/bst_map_generic.rs b/source/rust_verify/example/guide/bst_map_generic.rs new file mode 100644 index 000000000..4f215d935 --- /dev/null +++ b/source/rust_verify/example/guide/bst_map_generic.rs @@ -0,0 +1,547 @@ +// ANCHOR: all +use vstd::prelude::*; + +verus!{ + +// ANCHOR: 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, + }); +} +// ANCHOR_END: trait + +// ANCHOR: structs +struct Node { + key: K, + value: V, + left: Option>>, + right: Option>>, +} + +pub struct TreeMap { + root: Option>>, +} +// ANCHOR_END: structs + +impl Node { + spec fn optional_as_map(node_opt: Option>>) -> Map + decreases node_opt, + { + match node_opt { + None => Map::empty(), + Some(node) => node.as_map(), + } + } + + pub closed spec fn as_map(self) -> Map + decreases self, + { + Node::::optional_as_map(self.left) + .union_prefer_right(Node::::optional_as_map(self.right)) + .insert(self.key, self.value) + } +} + +impl TreeMap { + pub closed spec fn as_map(self) -> Map { + Node::::optional_as_map(self.root) + } +} + +impl View for TreeMap { + type V = Map; + + open spec fn view(&self) -> Map { + self.as_map() + } +} + +// ANCHOR: well_formed +impl Node { + pub closed spec fn well_formed(self) -> bool + decreases self + { + &&& (forall |elem| #[trigger] Node::::optional_as_map(self.left).dom().contains(elem) ==> elem.le(self.key) && elem != self.key) + &&& (forall |elem| #[trigger] Node::::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 TreeMap { + #[verifier::type_invariant] + spec fn well_formed(self) -> bool { + match self.root { + Some(node) => node.well_formed(), + None => true, // empty tree always well-formed + } + } +} +// ANCHOR_END: well_formed + +impl TreeMap { + pub fn new() -> (s: Self) + ensures + s@ == Map::::empty(), + { + TreeMap:: { root: None } + } +} + +impl Node { + fn insert_into_optional(node: &mut Option>>, key: K, value: V) + requires + old(node).is_some() ==> old(node).unwrap().well_formed(), + ensures + node.is_some() ==> node.unwrap().well_formed(), + Node::::optional_as_map(*node) =~= Node::::optional_as_map(*old(node)).insert(key, value) + { + if node.is_none() { + *node = Some(Box::new(Node:: { + 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: K, value: V) + requires + old(self).well_formed(), + ensures + self.well_formed(), + self.as_map() =~= old(self).as_map().insert(key, value), + { + match key.compare(&self.key) { + Cmp::Equal => { + self.value = value; + + assert(!Node::::optional_as_map(self.left).dom().contains(key)); + assert(!Node::::optional_as_map(self.right).dom().contains(key)); + } + Cmp::Less => { + Self::insert_into_optional(&mut self.left, key, value); + + proof { + if self.key.le(key) { + TotalOrdered::antisymmetric(self.key, key); + } + assert(!Node::::optional_as_map(self.right).dom().contains(key)); + } + } + Cmp::Greater => { + Self::insert_into_optional(&mut self.right, key, value); + + proof { + if key.le(self.key) { + TotalOrdered::antisymmetric(self.key, key); + } + assert(!Node::::optional_as_map(self.left).dom().contains(key)); + } + } + } + } +} + +impl TreeMap { + pub fn insert(&mut self, key: K, 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::::insert_into_optional(&mut root, key, value); + self.root = root; + } +} + +impl Node { + fn delete_from_optional(node: &mut Option>>, key: K) + requires + old(node).is_some() ==> old(node).unwrap().well_formed(), + ensures + node.is_some() ==> node.unwrap().well_formed(), + Node::::optional_as_map(*node) =~= Node::::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(); + + match key.compare(&boxed_node.key) { + Cmp::Equal => { + assert(!Node::::optional_as_map(boxed_node.left).dom().contains(key)); + assert(!Node::::optional_as_map(boxed_node.right).dom().contains(key)); + assert(boxed_node.right.is_some() ==> boxed_node.right.unwrap().well_formed()); + assert(boxed_node.left.is_some() ==> boxed_node.left.unwrap().well_formed()); + + 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::::delete_rightmost(&mut boxed_node.left); + boxed_node.key = popped_key; + boxed_node.value = popped_value; + *node = Some(boxed_node); + + proof { + assert forall |elem| #[trigger] Node::::optional_as_map(node.unwrap().right).dom().contains(elem) implies node.unwrap().key.le(elem) && elem != node.unwrap().key + by { + TotalOrdered::transitive(node.unwrap().key, key, elem); + if elem == node.unwrap().key { + TotalOrdered::antisymmetric(elem, key); + } + } + } + } + } + } + Cmp::Less => { + proof { + if Node::::optional_as_map(boxed_node.right).dom().contains(key) { + TotalOrdered::antisymmetric(boxed_node.key, key); + assert(false); + } + } + Node::::delete_from_optional(&mut boxed_node.left, key); + *node = Some(boxed_node); + } + Cmp::Greater => { + proof { + if Node::::optional_as_map(boxed_node.left).dom().contains(key) { + TotalOrdered::antisymmetric(boxed_node.key, key); + assert(false); + } + } + Node::::delete_from_optional(&mut boxed_node.right, key); + *node = Some(boxed_node); + } + } + } + } + + fn delete_rightmost(node: &mut Option>>) -> (popped: (K, V)) + requires + old(node).is_some(), + old(node).unwrap().well_formed(), + ensures + node.is_some() ==> node.unwrap().well_formed(), + Node::::optional_as_map(*node) =~= Node::::optional_as_map(*old(node)).remove(popped.0), + Node::::optional_as_map(*old(node)).dom().contains(popped.0), + Node::::optional_as_map(*old(node))[popped.0] == popped.1, + forall |elem| #[trigger] Node::::optional_as_map(*old(node)).dom().contains(elem) ==> elem.le(popped.0), + { + 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; + proof { + assert(Node::::optional_as_map(boxed_node.right) =~= Map::empty()); + assert(!Node::::optional_as_map(boxed_node.left).dom().contains(boxed_node.key)); + TotalOrdered::reflexive(boxed_node.key); + } + return (boxed_node.key, boxed_node.value); + } else { + let (popped_key, popped_value) = Node::::delete_rightmost(&mut boxed_node.right); + proof { + if Node::::optional_as_map(boxed_node.left).dom().contains(popped_key) { + TotalOrdered::antisymmetric(boxed_node.key, popped_key); + assert(false); + } + assert forall |elem| #[trigger] Node::::optional_as_map(*old(node)).dom().contains(elem) implies elem.le(popped_key) + by { + if elem.le(boxed_node.key) { + TotalOrdered::transitive(elem, boxed_node.key, popped_key); + } + } + } + *node = Some(boxed_node); + return (popped_key, popped_value); + } + } +} + +impl TreeMap { + pub fn delete(&mut self, key: K) + ensures + self@ == old(self)@.remove(key), + { + proof { use_type_invariant(&*self); } + let mut root = None; + std::mem::swap(&mut root, &mut self.root); + Node::::delete_from_optional(&mut root, key); + self.root = root; + } +} + +// ANCHOR: node_get +impl Node { + fn get_from_optional(node: &Option>>, 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::::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::::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) + } + } + } +} +// ANCHOR_END: node_get + +impl TreeMap { + pub fn get(&self, key: K) -> Option<&V> + returns (if self@.dom().contains(key) { Some(&self@[key]) } else { None }) + { + proof { use_type_invariant(&*self); } + Node::::get_from_optional(&self.root, key) + } +} + +// ANCHOR: clone_full_impl +impl Clone for Node { + fn clone(&self) -> (res: Self) + ensures + self.well_formed() ==> res.well_formed(), + self.as_map().dom() =~= res.as_map().dom(), + forall |key| #[trigger] res.as_map().dom().contains(key) ==> + call_ensures(V::clone, (&self.as_map()[key],), res.as_map()[key]) + { + // TODO(fixme): Assigning V::clone to a variable is a hack needed to work around + // this issue: https://github.com/verus-lang/verus/issues/1348 + let v_clone = V::clone; + + let res = Node { + key: self.key, + value: v_clone(&self.value), + // Ordinarily, we would use Option::clone rather than inlining + // the case statement here; we write it this way to work around + // this issue: https://github.com/verus-lang/verus/issues/1346 + left: (match &self.left { + Some(node) => Some(Box::new((&**node).clone())), + None => None, + }), + right: (match &self.right { + Some(node) => Some(Box::new((&**node).clone())), + None => None, + }), + }; + + proof { + assert(Node::optional_as_map(res.left).dom() =~= + Node::optional_as_map(self.left).dom()); + assert(Node::optional_as_map(res.right).dom() =~= + Node::optional_as_map(self.right).dom()); + } + + return res; + } +} + +// ANCHOR: clone_signature +impl Clone for TreeMap { + fn clone(&self) -> (res: Self) + ensures self@.dom() =~= res@.dom(), + forall |key| #[trigger] res@.dom().contains(key) ==> + call_ensures(V::clone, (&self@[key],), res@[key]) +// ANCHOR_END: clone_signature + { + proof { + use_type_invariant(self); + } + + TreeMap { + // This calls Option>::Clone + root: self.root.clone(), + } + } +} +// ANCHOR_END: clone_full_impl + +impl TotalOrdered for u64 { + open spec fn le(self, other: Self) -> bool { self <= other } + + proof fn reflexive(x: Self) { } + proof fn transitive(x: Self, y: Self, z: Self) { } + proof fn antisymmetric(x: Self, y: Self) { } + proof fn total(x: Self, y: Self) { } + + fn compare(&self, other: &Self) -> (c: Cmp) { + if *self == *other { + Cmp::Equal + } else if *self < *other { + Cmp::Less + } else { + Cmp::Greater + } + } +} + +fn test() { + let mut tree_map = TreeMap::::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) { + let mut tree_map = tree_map; + tree_map.insert(25, true); + tree_map.insert(100, true); +} + +// ANCHOR: clone_u32 +fn test_clone_u32(tree_map: TreeMap) { + let tree_map2 = tree_map.clone(); + assert(tree_map2@ =~= tree_map@); +} +// ANCHOR_END: clone_u32 + +// ANCHOR: clone_int_wrapper +struct IntWrapper { + pub int_value: u32, +} + +impl Clone for IntWrapper { + fn clone(&self) -> (s: Self) + ensures s == *self + { + IntWrapper { int_value: self.int_value } + } +} + +fn test_clone_int_wrapper(tree_map: TreeMap) { + let tree_map2 = tree_map.clone(); + assert(tree_map2@ =~= tree_map@); +} +// ANCHOR_END: clone_int_wrapper + +// ANCHOR: clone_weird_int +struct WeirdInt { + pub int_value: u32, + pub other: u32, +} + +impl Clone for WeirdInt { + fn clone(&self) -> (s: Self) + ensures s.int_value == self.int_value + { + WeirdInt { int_value: self.int_value, other: 0 } + } +} + +fn test_clone_weird_int(tree_map: TreeMap) { + let tree_map2 = tree_map.clone(); + + // assert(tree_map2@ =~= tree_map@); // this would fail + + assert(tree_map2@.dom() == tree_map@.dom()); + assert(forall |k| tree_map@.dom().contains(k) ==> + tree_map2@[k].int_value == tree_map@[k].int_value); +} +// ANCHOR_END: clone_weird_int + + +} +// ANCHOR_END: all + +fn main() { } + diff --git a/source/rust_verify/example/guide/bst_map_type_invariant.rs b/source/rust_verify/example/guide/bst_map_type_invariant.rs new file mode 100644 index 000000000..d499e5c2c --- /dev/null +++ b/source/rust_verify/example/guide/bst_map_type_invariant.rs @@ -0,0 +1,324 @@ +// ANCHOR: all +use vstd::prelude::*; + +verus!{ + +struct Node { + key: u64, + value: V, + left: Option>>, + right: Option>>, +} + +pub struct TreeMap { + root: Option>>, +} + +impl Node { + spec fn optional_as_map(node_opt: Option>>) -> Map + decreases node_opt, + { + match node_opt { + None => Map::empty(), + Some(node) => node.as_map(), + } + } + + spec fn as_map(self) -> Map + decreases self, + { + Node::::optional_as_map(self.left) + .union_prefer_right(Node::::optional_as_map(self.right)) + .insert(self.key, self.value) + } +} + +impl TreeMap { + pub closed spec fn as_map(self) -> Map { + Node::::optional_as_map(self.root) + } +} + +impl View for TreeMap { + type V = Map; + + open spec fn view(&self) -> Map { + self.as_map() + } +} + +impl Node { + spec fn well_formed(self) -> bool + decreases self + { + &&& (forall |elem| Node::::optional_as_map(self.left).dom().contains(elem) ==> elem < self.key) + &&& (forall |elem| Node::::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, + }) + } +} + +// ANCHOR: well_formed_with_attr +impl TreeMap { + #[verifier::type_invariant] + spec fn well_formed(self) -> bool { + match self.root { + Some(node) => node.well_formed(), + None => true, // empty tree always well-formed + } + } +} +// ANCHOR_END: well_formed_with_attr + +// ANCHOR: new +impl TreeMap { +// ANCHOR: new_signature + pub fn new() -> (s: Self) + ensures + s@ == Map::::empty() +// ANCHOR_END: new_signature + { + TreeMap:: { root: None } + } +} +// ANCHOR_END: new + +impl Node { + fn insert_into_optional(node: &mut Option>>, key: u64, value: V) + requires + old(node).is_some() ==> old(node).unwrap().well_formed(), + ensures + node.is_some() ==> node.unwrap().well_formed(), + Node::::optional_as_map(*node) =~= Node::::optional_as_map(*old(node)).insert(key, value) + { + if node.is_none() { + *node = Some(Box::new(Node:: { + 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::::optional_as_map(self.left).dom().contains(key)); + assert(!Node::::optional_as_map(self.right).dom().contains(key)); + } else if key < self.key { + Self::insert_into_optional(&mut self.left, key, value); + + assert(!Node::::optional_as_map(self.right).dom().contains(key)); + } else { + Self::insert_into_optional(&mut self.right, key, value); + + assert(!Node::::optional_as_map(self.left).dom().contains(key)); + } + } +} + +// ANCHOR: insert +impl TreeMap { +// ANCHOR: insert_signature + pub fn insert(&mut self, key: u64, value: V) + ensures + self@ == old(self)@.insert(key, value) +// ANCHOR_END: insert_signature + { + proof { use_type_invariant(&*self); } + let mut root = None; + std::mem::swap(&mut root, &mut self.root); + Node::::insert_into_optional(&mut root, key, value); + self.root = root; + } +} +// ANCHOR_END: insert + +impl Node { + fn delete_from_optional(node: &mut Option>>, key: u64) + requires + old(node).is_some() ==> old(node).unwrap().well_formed(), + ensures + node.is_some() ==> node.unwrap().well_formed(), + Node::::optional_as_map(*node) =~= Node::::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::::optional_as_map(boxed_node.left).dom().contains(key)); + assert(!Node::::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::::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::::optional_as_map(boxed_node.right).dom().contains(key)); + Node::::delete_from_optional(&mut boxed_node.left, key); + *node = Some(boxed_node); + } else { + assert(!Node::::optional_as_map(boxed_node.left).dom().contains(key)); + Node::::delete_from_optional(&mut boxed_node.right, key); + *node = Some(boxed_node); + } + } + } + + fn delete_rightmost(node: &mut Option>>) -> (popped: (u64, V)) + requires + old(node).is_some(), + old(node).unwrap().well_formed(), + ensures + node.is_some() ==> node.unwrap().well_formed(), + Node::::optional_as_map(*node) =~= Node::::optional_as_map(*old(node)).remove(popped.0), + Node::::optional_as_map(*old(node)).dom().contains(popped.0), + Node::::optional_as_map(*old(node))[popped.0] == popped.1, + forall |elem| Node::::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::::optional_as_map(boxed_node.right) =~= Map::empty()); + assert(!Node::::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::::delete_rightmost(&mut boxed_node.right); + assert(!Node::::optional_as_map(boxed_node.left).dom().contains(popped_key)); + *node = Some(boxed_node); + return (popped_key, popped_value); + } + } +} + +// ANCHOR: delete +impl TreeMap { +// ANCHOR: delete_signature + pub fn delete(&mut self, key: u64) + ensures + self@ == old(self)@.remove(key) +// ANCHOR_END: delete_signature + { + proof { use_type_invariant(&*self); } + let mut root = None; + std::mem::swap(&mut root, &mut self.root); + Node::::delete_from_optional(&mut root, key); + self.root = root; + } +} +// ANCHOR_END: delete + +impl Node { + fn get_from_optional(node: &Option>>, 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::::optional_as_map(self.right).dom().contains(key)); } + Self::get_from_optional(&self.left, key) + } else { + proof { assert(!Node::::optional_as_map(self.left).dom().contains(key)); } + Self::get_from_optional(&self.right, key) + } + } +} + +// ANCHOR: get +impl TreeMap { +// ANCHOR: get_signature + pub fn get(&self, key: u64) -> Option<&V> + returns (if self@.dom().contains(key) { Some(&self@[key]) } else { None }) +// ANCHOR_END: get_signature + { + proof { use_type_invariant(&*self); } + Node::::get_from_optional(&self.root, key) + } +} +// ANCHOR_END: get + +// ANCHOR: example_use +fn test() { + let mut tree_map = TreeMap::::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) { + let mut tree_map = tree_map; + tree_map.insert(25, true); + tree_map.insert(100, true); +} +// ANCHOR_END: example_use + + +} +// ANCHOR_END: all + +fn main() { } +