ic_stable_memory/collections/certified_btree_set/
mod.rs

1use crate::collections::certified_btree_map::SCertifiedBTreeMap;
2use crate::collections::certified_btree_set::iter::SCertifiedBTreeSetIter;
3use crate::encoding::AsFixedSizeBytes;
4use crate::primitive::s_ref::SRef;
5use crate::primitive::StableType;
6use crate::utils::certification::HashTree;
7use crate::{AsHashTree, AsHashableBytes};
8use std::borrow::Borrow;
9use std::fmt::{Debug, Formatter};
10
11pub mod iter;
12
13/// Certified B-plus tree based set data structure
14///
15/// This is just a wrapper around [SCertifiedBTreeMap]`<T, ()>`, read its documentation for more info on the internals.
16/// () is encoded as `empty` [utils::certification::HashTree].
17pub struct SCertifiedBTreeSet<T: StableType + AsFixedSizeBytes + Ord + AsHashableBytes> {
18    map: SCertifiedBTreeMap<T, ()>,
19}
20
21impl<T: Ord + StableType + AsFixedSizeBytes + AsHashableBytes> SCertifiedBTreeSet<T> {
22    /// See [SCertifiedBTreeMap::new]
23    #[inline]
24    pub fn new() -> Self {
25        Self {
26            map: SCertifiedBTreeMap::new(),
27        }
28    }
29
30    /// See [SCertifiedBTreeMap::len]
31    #[inline]
32    pub fn len(&self) -> u64 {
33        self.map.len()
34    }
35
36    /// See [SCertifiedBTreeMap::is_empty]
37    #[inline]
38    pub fn is_empty(&self) -> bool {
39        self.map.is_empty()
40    }
41
42    /// See [SCertifiedBTreeMap::insert]
43    #[inline]
44    pub fn insert(&mut self, value: T) -> Result<bool, T> {
45        self.map
46            .insert(value, ())
47            .map(|it| it.is_some())
48            .map_err(|(k, _)| k)
49    }
50
51    /// See [SCertifiedBTreeMap::insert_and_commit]
52    #[inline]
53    pub fn insert_and_commit(&mut self, value: T) -> Result<bool, T> {
54        self.map
55            .insert_and_commit(value, ())
56            .map(|it| it.is_some())
57            .map_err(|(k, _)| k)
58    }
59
60    /// See [SCertifiedBTreeMap::remove]
61    #[inline]
62    pub fn remove<Q>(&mut self, value: &Q) -> bool
63    where
64        T: Borrow<Q>,
65        Q: Ord + ?Sized,
66    {
67        self.map.remove(value).is_some()
68    }
69
70    /// See [SCertifiedBTreeMap::remove_and_commit]
71    #[inline]
72    pub fn remove_and_commit<Q>(&mut self, value: &Q) -> bool
73    where
74        T: Borrow<Q>,
75        Q: Ord + ?Sized,
76    {
77        self.map.remove_and_commit(value).is_some()
78    }
79
80    /// See [SCertifiedBTreeMap::commit]
81    #[inline]
82    pub fn commit(&mut self) {
83        self.map.commit();
84    }
85
86    /// See [SCertifiedBTreeMap::prove_absence]
87    #[inline]
88    pub fn prove_absence<Q>(&self, index: &Q) -> HashTree
89    where
90        T: Borrow<Q>,
91        Q: Ord + ?Sized,
92    {
93        self.map.prove_absence(index)
94    }
95
96    /// See [SCertifiedBTreeMap::prove_range]
97    #[inline]
98    pub fn prove_range<Q>(&self, from: &Q, to: &Q) -> HashTree
99    where
100        T: Borrow<Q>,
101        Q: Ord + ?Sized,
102    {
103        self.map.prove_range(from, to)
104    }
105
106    /// See [SCertifiedBTreeMap::witness]
107    #[inline]
108    pub fn witness<Q>(&self, index: &Q) -> HashTree
109    where
110        T: Borrow<Q>,
111        Q: Ord + ?Sized,
112    {
113        self.map.witness(index)
114    }
115
116    /// See [SCertifiedBTreeMap::clear]
117    #[inline]
118    pub fn clear(&mut self) {
119        self.map.clear();
120    }
121
122    /// See [SCertifiedBTreeMap::contains_key]
123    #[inline]
124    pub fn contains<Q>(&self, value: &Q) -> bool
125    where
126        T: Borrow<Q>,
127        Q: Ord + ?Sized,
128    {
129        self.map.contains_key(value)
130    }
131
132    /// See [SBTreeMap::get]
133    #[inline]
134    pub fn get_random(&self, seed: u32) -> Option<SRef<T>> {
135        self.map.get_random_key(seed)
136    }
137
138    /// See [SCertifiedBTreeMap::iter]
139    #[inline]
140    pub fn iter(&self) -> SCertifiedBTreeSetIter<T> {
141        SCertifiedBTreeSetIter::new(self)
142    }
143}
144
145impl<T: StableType + AsFixedSizeBytes + Ord + AsHashableBytes> AsHashTree
146    for SCertifiedBTreeSet<T>
147{
148    #[inline]
149    fn root_hash(&self) -> crate::utils::certification::Hash {
150        self.map.root_hash()
151    }
152
153    #[inline]
154    fn hash_tree(&self) -> HashTree {
155        self.map.hash_tree()
156    }
157}
158
159impl<T: Ord + StableType + AsFixedSizeBytes + AsHashableBytes> Default for SCertifiedBTreeSet<T> {
160    #[inline]
161    fn default() -> Self {
162        Self::new()
163    }
164}
165
166impl<T: StableType + AsFixedSizeBytes + Ord + AsHashableBytes> AsFixedSizeBytes
167    for SCertifiedBTreeSet<T>
168{
169    const SIZE: usize = SCertifiedBTreeMap::<T, ()>::SIZE;
170    type Buf = <SCertifiedBTreeMap<T, ()> as AsFixedSizeBytes>::Buf;
171
172    #[inline]
173    fn as_fixed_size_bytes(&self, buf: &mut [u8]) {
174        self.map.as_fixed_size_bytes(buf);
175    }
176
177    #[inline]
178    fn from_fixed_size_bytes(arr: &[u8]) -> Self {
179        let map = SCertifiedBTreeMap::<T, ()>::from_fixed_size_bytes(&arr);
180        Self { map }
181    }
182}
183
184impl<T: StableType + AsFixedSizeBytes + Ord + AsHashableBytes> StableType
185    for SCertifiedBTreeSet<T>
186{
187    #[inline]
188    unsafe fn stable_drop_flag_on(&mut self) {
189        self.map.stable_drop_flag_on();
190    }
191
192    #[inline]
193    unsafe fn stable_drop_flag_off(&mut self) {
194        self.map.stable_drop_flag_off()
195    }
196}
197
198impl<T: StableType + AsFixedSizeBytes + Ord + Debug + AsHashableBytes> Debug
199    for SCertifiedBTreeSet<T>
200{
201    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
202        f.write_str("(")?;
203        for (idx, elem) in self.iter().enumerate() {
204            elem.fmt(f)?;
205
206            if idx < (self.len() - 1) as usize {
207                f.write_str(", ")?;
208            }
209        }
210        f.write_str(")")
211    }
212}
213
214#[cfg(test)]
215mod tests {
216    use crate::collections::certified_btree_set::SCertifiedBTreeSet;
217    use crate::encoding::{AsFixedSizeBytes, Buffer};
218    use crate::{_debug_validate_allocator, get_allocated_size, stable, stable_memory_init};
219
220    #[test]
221    fn serialization_works_fine() {
222        stable::clear();
223        stable_memory_init();
224
225        {
226            let set = SCertifiedBTreeSet::<usize>::new();
227
228            let buf = set.as_new_fixed_size_bytes();
229            SCertifiedBTreeSet::<usize>::from_fixed_size_bytes(buf._deref());
230        }
231
232        _debug_validate_allocator();
233        assert_eq!(get_allocated_size(), 0);
234    }
235
236    #[test]
237    fn iter_works_fine() {
238        stable::clear();
239        stable_memory_init();
240
241        {
242            let mut set = SCertifiedBTreeSet::<usize>::default();
243            for i in 0..100 {
244                set.insert(i);
245            }
246
247            for (idx, mut i) in set.iter().enumerate() {
248                assert_eq!(idx, *i);
249            }
250        }
251
252        _debug_validate_allocator();
253        assert_eq!(get_allocated_size(), 0);
254    }
255}
ic_stable_memory/collections/certified_btree_set/mod.rs

ic_stable_memory/collections/certified_btree_set/
mod.rs
