Skip to main content

planarian_schema/
pathmap.rs

1//! `PathMap`: a map keyed by option path, stored as a vector sorted by key.
2//!
3//! It stands in for `BTreeMap<String, V>` where the `policy` proofs read the
4//! keys: vstd only specifies `BTreeMap` for key types with an ordering spec,
5//! which `String` lacks, while a `Vec` is specified unconditionally. The
6//! methods keep the entries sorted and unique; the proofs do not depend on
7//! that, only on which keys are present. Serialized as a JSON object, as the
8//! `BTreeMap` was.
9
10use serde::{Deserialize, Deserializer, Serialize, Serializer};
11use std::collections::BTreeMap;
12use vstd::prelude::*;
13
14verus! {
15
16#[derive(Clone, Debug, PartialEq, Eq)]
17pub struct PathMap<V> {
18    /// Sorted by key, keys unique. Read freely; change only through the methods.
19    pub entries: Vec<(String, V)>,
20}
21
22} // verus!
23
24impl<V> Default for PathMap<V> {
25    fn default() -> Self {
26        PathMap { entries: Vec::new() }
27    }
28}
29
30impl<V> PathMap<V> {
31    pub fn new() -> Self {
32        Self::default()
33    }
34
35    fn find(&self, key: &str) -> Result<usize, usize> {
36        self.entries.binary_search_by(|(k, _)| k.as_str().cmp(key))
37    }
38
39    pub fn get(&self, key: &str) -> Option<&V> {
40        self.find(key).ok().map(|i| &self.entries[i].1)
41    }
42
43    pub fn contains_key(&self, key: &str) -> bool {
44        self.find(key).is_ok()
45    }
46
47    pub fn insert(&mut self, key: String, value: V) -> Option<V> {
48        match self.find(&key) {
49            Ok(i) => Some(std::mem::replace(&mut self.entries[i].1, value)),
50            Err(i) => {
51                self.entries.insert(i, (key, value));
52                None
53            }
54        }
55    }
56
57    pub fn remove(&mut self, key: &str) -> Option<V> {
58        self.find(key).ok().map(|i| self.entries.remove(i).1)
59    }
60
61    pub fn retain(&mut self, mut keep: impl FnMut(&String, &mut V) -> bool) {
62        self.entries.retain_mut(|(k, v)| keep(k, v));
63    }
64
65    pub fn len(&self) -> usize {
66        self.entries.len()
67    }
68
69    pub fn is_empty(&self) -> bool {
70        self.entries.is_empty()
71    }
72
73    pub fn iter(&self) -> impl Iterator<Item = (&String, &V)> {
74        self.entries.iter().map(|(k, v)| (k, v))
75    }
76
77    pub fn keys(&self) -> impl Iterator<Item = &String> {
78        self.entries.iter().map(|(k, _)| k)
79    }
80
81    pub fn values(&self) -> impl Iterator<Item = &V> {
82        self.entries.iter().map(|(_, v)| v)
83    }
84}
85
86impl<V> FromIterator<(String, V)> for PathMap<V> {
87    fn from_iter<I: IntoIterator<Item = (String, V)>>(iter: I) -> Self {
88        let mut m = PathMap::new();
89        for (k, v) in iter {
90            m.insert(k, v);
91        }
92        m
93    }
94}
95
96impl<V, const N: usize> From<[(String, V); N]> for PathMap<V> {
97    fn from(items: [(String, V); N]) -> Self {
98        items.into_iter().collect()
99    }
100}
101
102impl<V> From<BTreeMap<String, V>> for PathMap<V> {
103    fn from(m: BTreeMap<String, V>) -> Self {
104        // Already sorted and unique.
105        PathMap { entries: m.into_iter().collect() }
106    }
107}
108
109impl<'a, V> IntoIterator for &'a PathMap<V> {
110    type Item = (&'a String, &'a V);
111    type IntoIter = std::iter::Map<std::slice::Iter<'a, (String, V)>, fn(&'a (String, V)) -> (&'a String, &'a V)>;
112
113    fn into_iter(self) -> Self::IntoIter {
114        self.entries.iter().map(|(k, v)| (k, v))
115    }
116}
117
118impl<V> IntoIterator for PathMap<V> {
119    type Item = (String, V);
120    type IntoIter = std::vec::IntoIter<(String, V)>;
121
122    fn into_iter(self) -> Self::IntoIter {
123        self.entries.into_iter()
124    }
125}
126
127impl<V: Serialize> Serialize for PathMap<V> {
128    fn serialize<S: Serializer>(&self, s: S) -> Result<S::Ok, S::Error> {
129        s.collect_map(self.iter())
130    }
131}
132
133impl<'de, V: Deserialize<'de>> Deserialize<'de> for PathMap<V> {
134    fn deserialize<D: Deserializer<'de>>(d: D) -> Result<Self, D::Error> {
135        Ok(BTreeMap::<String, V>::deserialize(d)?.into())
136    }
137}
138
139#[cfg(test)]
140mod tests {
141    use super::*;
142
143    #[test]
144    fn behaves_like_a_btreemap() {
145        let mut m: PathMap<i32> = [("b".to_string(), 2), ("a".to_string(), 1)].into();
146        assert_eq!(m.keys().map(String::as_str).collect::<Vec<_>>(), ["a", "b"]);
147        assert_eq!(m.insert("a".into(), 3), Some(1));
148        assert_eq!(m.insert("c".into(), 4), None);
149        assert_eq!(m.get("a"), Some(&3));
150        m.retain(|k, _| k != "b");
151        assert_eq!(m.remove("c"), Some(4));
152        assert_eq!(m.len(), 1);
153        let json = serde_json::to_string(&m).unwrap();
154        assert_eq!(json, r#"{"a":3}"#);
155        assert_eq!(serde_json::from_str::<PathMap<i32>>(&json).unwrap(), m);
156    }
157}