iris_ztd/
zmap.rs

1use core::borrow::Borrow;
2
3use crate::Zeroable;
4use crate::{Digest, Hashable, Noun, NounDecode, NounEncode};
5use alloc::boxed::Box;
6use alloc::fmt::Debug;
7use alloc::vec;
8use alloc::vec::Vec;
9
10#[derive(Debug, Clone, PartialEq, Eq)]
11pub struct ZMap<K, V> {
12    root: Zeroable<Box<Node<K, V>>>,
13}
14
15#[derive(Debug, Clone, PartialEq, Eq)]
16struct Node<K, V> {
17    key: K,
18    value: V,
19    left: Zeroable<Box<Node<K, V>>>,
20    right: Zeroable<Box<Node<K, V>>>,
21}
22
23impl<K, V> Default for ZMap<K, V> {
24    fn default() -> Self {
25        Self::new()
26    }
27}
28
29impl<K, V> ZMap<K, V> {
30    pub fn new() -> Self {
31        ZMap {
32            root: Zeroable(None),
33        }
34    }
35}
36
37impl<K: NounEncode, V: NounEncode> ZMap<K, V> {
38    pub fn insert(&mut self, key: K, value: V) -> bool {
39        let (new_root, inserted) = Self::put(self.root.take(), key, value);
40        self.root = Zeroable(Some(new_root));
41        inserted
42    }
43
44    pub fn get<Q: NounEncode + ?Sized>(&self, key: &Q) -> Option<&V>
45    where
46        K: Borrow<Q>,
47    {
48        Self::get_inner(self.root.0.as_ref()?, key)
49    }
50
51    fn get_inner<'a, Q: NounEncode + ?Sized>(n: &'a Node<K, V>, key: &Q) -> Option<&'a V>
52    where
53        K: Borrow<Q>,
54    {
55        if Self::tip_eq(&key, &n.key) {
56            return Some(&n.value);
57        }
58        let go_left = Self::gor_tip(&key, &n.key);
59        if go_left {
60            Self::get_inner(n.left.as_ref()?, key)
61        } else {
62            Self::get_inner(n.right.as_ref()?, key)
63        }
64    }
65
66    fn put(node: Option<Box<Node<K, V>>>, key: K, value: V) -> (Box<Node<K, V>>, bool) {
67        match node {
68            None => (
69                Box::new(Node {
70                    key,
71                    value,
72                    left: Zeroable(None),
73                    right: Zeroable(None),
74                }),
75                true,
76            ),
77            Some(mut n) => {
78                if Self::tip_eq(&key, &n.key) {
79                    return (n, false);
80                }
81                let go_left = Self::gor_tip(&key, &n.key);
82                if go_left {
83                    let (new_left, inserted) = Self::put(n.left.take(), key, value);
84                    n.left = Zeroable(Some(new_left));
85                    if !Self::mor_tip(&n.key, &n.left.as_ref().unwrap().key) {
86                        // Rotate right
87                        let mut new_root = n.left.take().unwrap();
88                        n.left = Zeroable(new_root.right.take());
89                        new_root.right = Zeroable(Some(n));
90                        (new_root, inserted)
91                    } else {
92                        (n, inserted)
93                    }
94                } else {
95                    let (new_right, inserted) = Self::put(n.right.take(), key, value);
96                    n.right = Zeroable(Some(new_right));
97                    if !Self::mor_tip(&n.key, &n.right.as_ref().unwrap().key) {
98                        // Rotate left
99                        let mut new_root = n.right.take().unwrap();
100                        n.right = Zeroable(new_root.left.take());
101                        new_root.left = Zeroable(Some(n));
102                        (new_root, inserted)
103                    } else {
104                        (n, inserted)
105                    }
106                }
107            }
108        }
109    }
110
111    fn tip_eq<Q: NounEncode + ?Sized>(a: &Q, b: &K) -> bool {
112        a.to_noun().hash() == b.to_noun().hash()
113    }
114
115    fn gor_tip<Q: NounEncode + ?Sized>(a: &Q, b: &K) -> bool {
116        a.to_noun().hash().to_bytes() < b.to_noun().hash().to_bytes()
117    }
118
119    fn mor_tip<Q: NounEncode + ?Sized>(a: &Q, b: &K) -> bool {
120        Self::double_tip(a).to_bytes() < Self::double_tip(b).to_bytes()
121    }
122
123    fn double_tip<Q: NounEncode + ?Sized>(a: &Q) -> Digest {
124        (a.to_noun().hash(), a.to_noun().hash()).hash()
125    }
126}
127
128impl<K: NounEncode, V: NounEncode> core::iter::FromIterator<(K, V)> for ZMap<K, V> {
129    fn from_iter<I: IntoIterator<Item = (K, V)>>(iter: I) -> Self {
130        let mut set = ZMap::new();
131        for (k, v) in iter {
132            set.insert(k, v);
133        }
134        set
135    }
136}
137
138impl<K: NounEncode + Hashable, V: NounEncode + Hashable> Hashable for ZMap<K, V> {
139    fn hash(&self) -> Digest {
140        fn hash_node<K: NounEncode + Hashable, V: NounEncode + Hashable>(
141            node: &Zeroable<Box<Node<K, V>>>,
142        ) -> Digest {
143            match &node.0 {
144                None => 0.hash(),
145                Some(n) => {
146                    let left_hash = hash_node(&n.left);
147                    let right_hash = hash_node(&n.right);
148                    ((&n.key, &n.value), (left_hash, right_hash)).hash()
149                }
150            }
151        }
152        hash_node(&self.root)
153    }
154}
155
156impl<K: NounEncode + Hashable, V: NounEncode> NounEncode for ZMap<K, V> {
157    fn to_noun(&self) -> Noun {
158        fn visit<K: NounEncode + Hashable, V: NounEncode>(
159            node: &Zeroable<Box<Node<K, V>>>,
160        ) -> Noun {
161            match &node.0 {
162                None => 0.to_noun(),
163                Some(n) => {
164                    let left_hash = visit(&n.left);
165                    let right_hash = visit(&n.right);
166                    ((&n.key, &n.value), (left_hash, right_hash)).to_noun()
167                }
168            }
169        }
170        visit(&self.root)
171    }
172}
173
174impl<K: NounDecode, V: NounDecode> NounDecode for Node<K, V> {
175    fn from_noun(noun: &Noun) -> Option<Self> {
176        let ((key, value), left, right) = NounDecode::from_noun(noun)?;
177        Some(Self {
178            key,
179            value,
180            left,
181            right,
182        })
183    }
184}
185
186impl<K: NounDecode, V: NounDecode> NounDecode for ZMap<K, V> {
187    fn from_noun(noun: &Noun) -> Option<Self> {
188        let root: Zeroable<Box<Node<K, V>>> = NounDecode::from_noun(noun)?;
189        Some(Self { root })
190    }
191}
192
193pub struct ZMapIntoIterator<K, V> {
194    stack: Vec<Box<Node<K, V>>>,
195}
196
197impl<K, V> Iterator for ZMapIntoIterator<K, V> {
198    type Item = (K, V);
199
200    fn next(&mut self) -> Option<Self::Item> {
201        let cur = self.stack.pop()?;
202        if let Some(n) = cur.left.0 {
203            self.stack.push(n);
204        }
205        if let Some(n) = cur.right.0 {
206            self.stack.push(n);
207        }
208        Some((cur.key, cur.value))
209    }
210}
211
212impl<K, V> IntoIterator for ZMap<K, V> {
213    type Item = (K, V);
214    type IntoIter = ZMapIntoIterator<K, V>;
215
216    fn into_iter(self) -> Self::IntoIter {
217        let mut stack = vec![];
218        if let Some(n) = self.root.0 {
219            stack.push(n);
220        }
221        ZMapIntoIterator { stack }
222    }
223}
224
225impl<K, V> From<ZMap<K, V>> for Vec<(K, V)> {
226    fn from(map: ZMap<K, V>) -> Self {
227        map.into_iter().collect()
228    }
229}
230
231#[cfg(test)]
232mod tests {
233    use super::*;
234    use alloc::string::{String, ToString};
235
236    #[test]
237    fn test_zmap_encode_decode() {
238        let mut zm = ZMap::<String, u64>::new();
239        zm.insert("ver".to_string(), 10);
240        zm.insert("ve2".to_string(), 11);
241        let zm_noun = zm.to_noun();
242        let zm_decode = ZMap::<String, u64>::from_noun(&zm_noun).unwrap();
243        assert_eq!(Vec::from(zm), Vec::from(zm_decode));
244    }
245}
iris_ztd/zmap.rs

iris_ztd/
zmap.rs
