Skip to main content

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