iris_ztd/
zset.rs

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

iris_ztd/
zset.rs
