Skip to main content

vstd/
map.rs

1#![allow(unused_imports)]
2
3use super::pervasive::*;
4use super::prelude::*;
5use super::set::*;
6
7use verus as verus_; // skip verusfmt due to unhandled return-value-pattern
8verus_! {
9
10/// `Map<K, V>` is an abstract map type for specifications.
11/// To use a "map" in compiled code, use an `exec` type like HashMap (TODO)
12/// that has a `Map<K, V>` as its specification type.
13///
14/// An object `map: Map<K, V>` has a _domain_, a set of keys given by [`map.dom()`](Map::dom),
15/// and a mapping for keys in the domain to values, given by [`map[key]`](Map::index).
16/// Alternatively, a map can be thought of as a set of `(K, V)` pairs where each key
17/// appears in at most entry.
18///
19/// In general, a map might be infinite.
20/// To work specifically with finite maps, see the [`self.finite()`](Set::finite) predicate.
21///
22/// Maps can be constructed in a few different ways:
23///  * [`Map::empty()`] constructs an empty map.
24///  * [`Map::new`] and [`Map::total`] construct a map given functions that specify its domain and the mapping
25///     from keys to values (a _map comprehension_).
26///  * The [`map!`] macro, to construct small maps of a fixed size.
27///  * By manipulating an existing map with [`Map::insert`] or [`Map::remove`].
28///
29/// To prove that two maps are equal, it is usually easiest to use the extensionality operator `=~=`.
30#[verifier::ext_equal]
31#[verifier::reject_recursive_types(K)]
32#[verifier::accept_recursive_types(V)]
33pub tracked struct Map<K, V> {
34    mapping: spec_fn(K) -> Option<V>,
35}
36
37impl<K, V> Map<K, V> {
38    /// An empty map.
39    pub closed spec fn empty() -> Map<K, V> {
40        Map { mapping: |k| None }
41    }
42
43    /// Gives a `Map<K, V>` whose domain contains every key, and maps each key
44    /// to the value given by `fv`.
45    pub open spec fn total(fv: spec_fn(K) -> V) -> Map<K, V> {
46        Set::full().mk_map(fv)
47    }
48
49    /// Gives a `Map<K, V>` whose domain is given by the boolean predicate on keys `fk`,
50    /// and maps each key to the value given by `fv`.
51    pub open spec fn new(fk: spec_fn(K) -> bool, fv: spec_fn(K) -> V) -> Map<K, V> {
52        Set::new(fk).mk_map(fv)
53    }
54
55    /// The domain of the map as a set.
56    pub closed spec fn dom(self) -> Set<K> {
57        Set::new(|k| (self.mapping)(k) is Some)
58    }
59
60    /// Gets the value that the given key `key` maps to.
61    /// For keys not in the domain, the result is meaningless and arbitrary.
62    pub closed spec fn index(self, key: K) -> V
63        recommends
64            self.dom().contains(key),
65    {
66        (self.mapping)(key)->Some_0
67    }
68
69    /// `[]` operator, synonymous with `index`
70    #[verifier::inline]
71    pub open spec fn spec_index(self, key: K) -> V
72        recommends
73            self.dom().contains(key),
74    {
75        self.index(key)
76    }
77
78    /// Inserts the given (key, value) pair into the map.
79    ///
80    /// If the key is already present from the map, then its existing value is overwritten
81    /// by the new value.
82    pub closed spec fn insert(self, key: K, value: V) -> Map<K, V> {
83        Map {
84            mapping: |k|
85                if k == key {
86                    Some(value)
87                } else {
88                    (self.mapping)(k)
89                },
90        }
91    }
92
93    /// Removes the given key and its associated value from the map.
94    ///
95    /// If the key is already absent from the map, then the map is left unchanged.
96    pub closed spec fn remove(self, key: K) -> Map<K, V> {
97        Map {
98            mapping: |k|
99                if k == key {
100                    None
101                } else {
102                    (self.mapping)(k)
103                },
104        }
105    }
106
107    /// Returns the number of key-value pairs in the map
108    pub open spec fn len(self) -> nat {
109        self.dom().len()
110    }
111
112    /// Create an empty tracked map.
113    ///
114    /// This allows us to create a map, which we know is empty, that is _tracked_.
115    pub axiom fn tracked_empty() -> (tracked out_v: Self)
116        ensures
117            out_v == Map::<K, V>::empty(),
118    ;
119
120    /// Inserts the given `(key, tracked value)` pair into the map.
121    ///
122    /// If the key is already present from the map, then its existing value is overwritten
123    /// by the new value.
124    pub axiom fn tracked_insert(tracked &mut self, key: K, tracked value: V)
125        ensures
126            *final(self) == Map::insert(*old(self), key, value),
127    ;
128
129    /// Removes the given key and its associated _tracked_ value from the map.
130    ///
131    /// The key must exist in the map
132    pub axiom fn tracked_remove(tracked &mut self, key: K) -> (tracked v: V)
133        requires
134            old(self).dom().contains(key),
135        ensures
136            *final(self) == Map::remove(*old(self), key),
137            v == old(self)[key],
138    ;
139
140    /// Index into a tracked map, getting a tracked borrow of the value
141    pub axiom fn tracked_borrow(tracked &self, key: K) -> (tracked v: &V)
142        requires
143            self.dom().contains(key),
144        ensures
145            *v == self.index(key),
146    ;
147
148    /// Index into a tracked map, getting a tracked mutable borrow of the value
149    pub axiom fn tracked_borrow_mut(tracked &mut self, key: K) -> (tracked v: &mut V)
150        requires
151            self.dom().contains(key),
152        ensures
153            *v == old(self).index(key),
154            *final(self) == old(self).insert(key, *final(v))
155    ;
156
157    /// Split a mutable borrow of a map into two.
158    pub axiom fn tracked_borrow_mut_split(tracked &mut self, keys: Set<K>)
159        -> (tracked (m1, m2): (&mut Self, &mut Self))
160        requires
161            keys <= self.dom(),
162        ensures
163            *m1 == old(self).restrict(keys),
164            *m2 == old(self).remove_keys(keys),
165            *final(self) == final(m1).union_prefer_right(*final(m2)),
166    ;
167
168    /// Change the keys of a map, by reverse lookup in a different map.
169    ///
170    /// For each `(old_key, new_key)` pair in `key_map`, the new map will have `(new_key, old_map[old_key])`.
171    /// Note the new map may be smaller than the old map if the `key_map` omits mappings for some of the old keys.
172    pub axiom fn tracked_map_keys<J>(
173        tracked old_map: Map<K, V>,
174        key_map: Map<J, K>,
175    ) -> (tracked new_map: Map<J, V>)
176        requires
177            forall|j| #![auto] key_map.contains_key(j) ==> old_map.contains_key(key_map[j]),
178            forall|j1, j2|
179                #![auto]
180                j1 != j2 && key_map.contains_key(j1) && key_map.contains_key(j2) ==> key_map[j1]
181                    != key_map[j2],
182        ensures
183            new_map.dom() == key_map.dom(),
184            forall|j|
185                key_map.contains_key(j) ==> new_map.contains_key(j) && #[trigger] new_map[j]
186                    == old_map[key_map[j]],
187    ;
188
189    /// Extract a set of keys (and their corresponding values) out of the map.
190    ///
191    /// This allows us to split a map based on a subset of the domain.
192    pub axiom fn tracked_remove_keys(tracked &mut self, keys: Set<K>) -> (tracked out_map: Map<
193        K,
194        V,
195    >)
196        requires
197            keys.subset_of(old(self).dom()),
198        ensures
199            *final(self) == old(self).remove_keys(keys),
200            out_map == old(self).restrict(keys),
201    ;
202
203    /// Merge a map into a tracked map.
204    ///
205    /// The new (key, value) pairs take precendece.
206    pub axiom fn tracked_union_prefer_right(tracked &mut self, right: Self)
207        ensures
208            *final(self) == old(self).union_prefer_right(right),
209    ;
210}
211
212// Trusted axioms
213/* REVIEW: this is simpler than the two separate axioms below -- would this be ok?
214pub broadcast axiom fn axiom_map_index_decreases<K, V>(m: Map<K, V>, key: K)
215    requires
216        m.dom().contains(key),
217    ensures
218        #[trigger](decreases_to!(m => m[key]));
219*/
220
221pub broadcast axiom fn axiom_map_index_decreases_finite<K, V>(m: Map<K, V>, key: K)
222    requires
223        m.dom().finite(),
224        m.dom().contains(key),
225    ensures
226        #[trigger] (decreases_to!(m => m[key])),
227;
228
229// REVIEW: this is currently a special case that is hard-wired into the verifier
230// It implements a version of https://github.com/FStarLang/FStar/pull/2954 .
231pub broadcast axiom fn axiom_map_index_decreases_infinite<K, V>(m: Map<K, V>, key: K)
232    requires
233        m.dom().contains(key),
234    ensures
235        #[trigger] is_smaller_than_recursive_function_field(m[key], m),
236;
237
238/// The domain of the empty map is the empty set
239pub broadcast proof fn axiom_map_empty<K, V>()
240    ensures
241        #[trigger] Map::<K, V>::empty().dom() == Set::<K>::empty(),
242{
243    broadcast use super::set::group_set_axioms;
244
245    assert(Set::new(|k: K| (|k| None::<V>)(k) is Some) == Set::<K>::empty());
246}
247
248/// The domain of a map after inserting a key-value pair is equivalent to inserting the key into
249/// the original map's domain set.
250pub broadcast proof fn axiom_map_insert_domain<K, V>(m: Map<K, V>, key: K, value: V)
251    ensures
252        #[trigger] m.insert(key, value).dom() == m.dom().insert(key),
253{
254    broadcast use super::set::group_set_axioms;
255
256    assert(m.insert(key, value).dom() =~= m.dom().insert(key));
257}
258
259/// Inserting `value` at `key` in `m` results in a map that maps `key` to `value`
260pub broadcast proof fn axiom_map_insert_same<K, V>(m: Map<K, V>, key: K, value: V)
261    ensures
262        #[trigger] m.insert(key, value)[key] == value,
263{
264}
265
266/// Inserting `value` at `key2` does not change the value mapped to by any other keys in `m`
267pub broadcast proof fn axiom_map_insert_different<K, V>(m: Map<K, V>, key1: K, key2: K, value: V)
268    requires
269        key1 != key2,
270    ensures
271        #[trigger] m.insert(key2, value)[key1] == m[key1],
272{
273}
274
275/// The domain of a map after removing a key-value pair is equivalent to removing the key from
276/// the original map's domain set.
277pub broadcast proof fn axiom_map_remove_domain<K, V>(m: Map<K, V>, key: K)
278    ensures
279        #[trigger] m.remove(key).dom() == m.dom().remove(key),
280{
281    broadcast use super::set::group_set_axioms;
282
283    assert(m.remove(key).dom() =~= m.dom().remove(key));
284}
285
286/// Removing a key-value pair from a map does not change the value mapped to by
287/// any other keys in the map.
288pub broadcast proof fn axiom_map_remove_different<K, V>(m: Map<K, V>, key1: K, key2: K)
289    requires
290        key1 != key2,
291    ensures
292        #[trigger] m.remove(key2)[key1] == m[key1],
293{
294}
295
296/// Two maps are equivalent if their domains are equivalent and every key in their domains map to the same value.
297pub broadcast proof fn axiom_map_ext_equal<K, V>(m1: Map<K, V>, m2: Map<K, V>)
298    ensures
299        #[trigger] (m1 =~= m2) <==> {
300            &&& m1.dom() =~= m2.dom()
301            &&& forall|k: K| #![auto] m1.dom().contains(k) ==> m1[k] == m2[k]
302        },
303{
304    broadcast use super::set::group_set_axioms;
305
306    if m1 =~= m2 {
307        assert(m1.dom() =~= m2.dom());
308        assert(forall|k: K| #![auto] m1.dom().contains(k) ==> m1[k] == m2[k]);
309    }
310    if ({
311        &&& m1.dom() =~= m2.dom()
312        &&& forall|k: K| #![auto] m1.dom().contains(k) ==> m1[k] == m2[k]
313    }) {
314        if m1.mapping != m2.mapping {
315            assert(exists|k| #[trigger] (m1.mapping)(k) != (m2.mapping)(k));
316            let k = choose|k| #[trigger] (m1.mapping)(k) != (m2.mapping)(k);
317            if m1.dom().contains(k) {
318                assert(m1[k] == m2[k]);
319            }
320            assert(false);
321        }
322        assert(m1 =~= m2);
323    }
324}
325
326pub broadcast proof fn axiom_map_ext_equal_deep<K, V>(m1: Map<K, V>, m2: Map<K, V>)
327    ensures
328        #[trigger] (m1 =~~= m2) <==> {
329            &&& m1.dom() =~~= m2.dom()
330            &&& forall|k: K| #![auto] m1.dom().contains(k) ==> m1[k] =~~= m2[k]
331        },
332{
333    axiom_map_ext_equal(m1, m2);
334}
335
336pub broadcast group group_map_axioms {
337    axiom_map_index_decreases_finite,
338    axiom_map_index_decreases_infinite,
339    axiom_map_empty,
340    axiom_map_insert_domain,
341    axiom_map_insert_same,
342    axiom_map_insert_different,
343    axiom_map_remove_domain,
344    axiom_map_remove_different,
345    axiom_map_ext_equal,
346    axiom_map_ext_equal_deep,
347}
348
349// Macros
350#[doc(hidden)]
351#[macro_export]
352macro_rules! map_internal {
353    [$($key:expr => $value:expr),* $(,)?] => {
354        $crate::vstd::map::Map::empty()
355            $(.insert($key, $value))*
356    }
357}
358
359/// Create a map using syntax like `map![key1 => val1, key2 => val, ...]`.
360///
361/// This is equivalent to `Map::empty().insert(key1, val1).insert(key2, val2)...`.
362///
363/// Note that this does _not_ require all keys to be distinct. In the case that two
364/// or more keys are equal, the resulting map uses the value of the rightmost entry.
365#[macro_export]
366macro_rules! map {
367    [$($tail:tt)*] => {
368        $crate::vstd::prelude::verus_proof_macro_exprs!($crate::vstd::map::map_internal!($($tail)*))
369    };
370}
371
372#[doc(hidden)]
373#[verifier::inline]
374pub open spec fn check_argument_is_map<K, V>(m: Map<K, V>) -> Map<K, V> {
375    m
376}
377
378#[doc(hidden)]
379pub use map_internal;
380pub use map;
381
382/// Prove two maps `map1` and `map2` are equal by proving that their values are equal at each key.
383///
384/// More precisely, `assert_maps_equal!` requires that for each key `k`:
385///  * `map1` contains `k` in its domain if and only if `map2` does (`map1.dom().contains(k) <==> map2.dom().contains(k)`)
386///  * If they contain `k` in their domains, then their values are equal (`map1.dom().contains(k) && map2.dom().contains(k) ==> map1[k] == map2[k]`)
387///
388/// The property that equality follows from these facts is often called _extensionality_.
389///
390/// `assert_maps_equal!` can handle many trivial-looking
391/// identities without any additional help:
392///
393/// ```rust
394/// proof fn insert_remove(m: Map<int, int>, k: int, v: int)
395///     requires !m.dom().contains(k)
396///     ensures m.insert(k, v).remove(k) == m
397/// {
398///     let m2 = m.insert(k, v).remove(k);
399///     assert_maps_equal!(m == m2);
400///     assert(m == m2);
401/// }
402/// ```
403///
404/// For more complex cases, a proof may be required for each key:
405///
406/// ```rust
407/// proof fn bitvector_maps() {
408///     let m1 = Map::<u64, u64>::new(
409///         |key: u64| key & 31 == key,
410///         |key: u64| key | 5);
411///
412///     let m2 = Map::<u64, u64>::new(
413///         |key: u64| key < 32,
414///         |key: u64| 5 | key);
415///
416///     assert_maps_equal!(m1 == m2, key => {
417///         // Show that the domains of m1 and m2 are the same by showing their predicates
418///         // are equivalent.
419///         assert_bit_vector((key & 31 == key) <==> (key < 32));
420///
421///         // Show that the values are the same by showing that these expressions
422///         // are equivalent.
423///         assert_bit_vector(key | 5 == 5 | key);
424///     });
425/// }
426/// ```
427#[macro_export]
428macro_rules! assert_maps_equal {
429    [$($tail:tt)*] => {
430        $crate::vstd::prelude::verus_proof_macro_exprs!($crate::vstd::map::assert_maps_equal_internal!($($tail)*))
431    };
432}
433
434#[macro_export]
435#[doc(hidden)]
436macro_rules! assert_maps_equal_internal {
437    (::verus_builtin::spec_eq($m1:expr, $m2:expr)) => {
438        assert_maps_equal_internal!($m1, $m2)
439    };
440    (::verus_builtin::spec_eq($m1:expr, $m2:expr), $k:ident $( : $t:ty )? => $bblock:block) => {
441        assert_maps_equal_internal!($m1, $m2, $k $( : $t )? => $bblock)
442    };
443    ($m1:expr, $m2:expr $(,)?) => {
444        assert_maps_equal_internal!($m1, $m2, key => { })
445    };
446    ($m1:expr, $m2:expr, $k:ident $( : $t:ty )? => $bblock:block) => {
447        #[verifier::spec] let m1 = $crate::vstd::map::check_argument_is_map($m1);
448        #[verifier::spec] let m2 = $crate::vstd::map::check_argument_is_map($m2);
449        $crate::vstd::prelude::assert_by($crate::vstd::prelude::equal(m1, m2), {
450            $crate::vstd::prelude::assert_forall_by(|$k $( : $t )?| {
451                // TODO better error message here: show the individual conjunct that fails,
452                // and maybe give an error message in english as well
453                $crate::vstd::prelude::ensures([
454                    $crate::vstd::prelude::imply(#[verifier::trigger] m1.dom().contains($k), m2.dom().contains($k))
455                    && $crate::vstd::prelude::imply(m2.dom().contains($k), m1.dom().contains($k))
456                    && $crate::vstd::prelude::imply(m1.dom().contains($k) && m2.dom().contains($k),
457                        $crate::vstd::prelude::equal(m1.index($k), m2.index($k)))
458                ]);
459                { $bblock }
460            });
461            $crate::vstd::prelude::assert_($crate::vstd::prelude::ext_equal(m1, m2));
462        });
463    }
464}
465
466#[doc(hidden)]
467pub use assert_maps_equal_internal;
468pub use assert_maps_equal;
469
470} // verus!
471
472verus_! { // skip verusfmt, issue with 'final'
473
474impl<K, V> Map<K, V> {
475    pub proof fn tracked_map_keys_in_place(tracked &mut self, key_map: Map<K, K>)
476        requires
477            forall|j|
478                #![auto]
479                key_map.dom().contains(j) ==> old(self).dom().contains(key_map.index(j)),
480            forall|j1, j2|
481                #![auto]
482                j1 != j2 && key_map.dom().contains(j1) && key_map.dom().contains(j2)
483                    ==> key_map.index(j1) != key_map.index(j2),
484        ensures
485            forall|j| #[trigger] final(self).dom().contains(j) == key_map.dom().contains(j),
486            forall|j|
487                key_map.dom().contains(j) ==> final(self).dom().contains(j) && #[trigger] final(self).index(j)
488                    == old(self).index(key_map.index(j)),
489    {
490        let tracked mut tmp = Self::tracked_empty();
491        super::modes::tracked_swap(&mut tmp, self);
492        let tracked mut tmp = Self::tracked_map_keys(tmp, key_map);
493        super::modes::tracked_swap(&mut tmp, self);
494    }
495}
496
497} // verus!