ic_stable_memory/collections/hash_set/
mod.rs

1use crate::collections::hash_map::SHashMap;
2use crate::collections::hash_set::iter::SHashSetIter;
3use crate::encoding::AsFixedSizeBytes;
4use crate::primitive::StableType;
5use crate::OutOfMemory;
6use std::borrow::Borrow;
7use std::fmt::{Debug, Formatter};
8use std::hash::Hash;
9
10#[doc(hidden)]
11pub mod iter;
12
13/// Hashmap-based hashset
14///
15/// This is just a wrapper around [SHashMap]`<T, ()>`, read it's documentation to get info on the internals.
16pub struct SHashSet<T: StableType + AsFixedSizeBytes + Hash + Eq> {
17    map: SHashMap<T, ()>,
18}
19
20impl<T: StableType + AsFixedSizeBytes + Hash + Eq> SHashSet<T> {
21    /// See [SHashMap::new]
22    #[inline]
23    pub fn new() -> Self {
24        Self {
25            map: SHashMap::new(),
26        }
27    }
28
29    /// See [SHashMap::new_with_capacity]
30    #[inline]
31    pub fn new_with_capacity(capacity: usize) -> Result<Self, OutOfMemory> {
32        Ok(Self {
33            map: SHashMap::new_with_capacity(capacity)?,
34        })
35    }
36
37    /// See [SHashMap::insert]
38    #[inline]
39    pub fn insert(&mut self, value: T) -> Result<bool, T> {
40        self.map
41            .insert(value, ())
42            .map(|it| it.is_some())
43            .map_err(|(k, _)| k)
44    }
45
46    /// See [SHashMap::remove]
47    #[inline]
48    pub fn remove<Q>(&mut self, value: &Q) -> bool
49    where
50        T: Borrow<Q>,
51        Q: Hash + Eq + ?Sized,
52    {
53        self.map.remove(value).is_some()
54    }
55
56    /// See [SHashMap::contains_key]
57    #[inline]
58    pub fn contains<Q>(&self, value: &Q) -> bool
59    where
60        T: Borrow<Q>,
61        Q: Hash + Eq + ?Sized,
62    {
63        self.map.contains_key(value)
64    }
65
66    /// See [SHashMap::len]
67    #[inline]
68    pub fn len(&self) -> usize {
69        self.map.len()
70    }
71
72    /// See [SHashMap::capacity]
73    #[inline]
74    pub fn capacity(&self) -> usize {
75        self.map.capacity()
76    }
77
78    /// See [SHashMap::is_empty]
79    #[inline]
80    pub fn is_empty(&self) -> bool {
81        self.map.is_empty()
82    }
83
84    /// See [SHashMap::is_full]
85    #[inline]
86    pub fn is_full(&self) -> bool {
87        self.map.is_full()
88    }
89
90    /// See [SHashMap::iter]
91    #[inline]
92    pub fn iter(&self) -> SHashSetIter<T> {
93        SHashSetIter::new(self)
94    }
95
96    /// See [SHashMap::clear]
97    #[inline]
98    pub fn clear(&mut self) {
99        self.map.clear();
100    }
101}
102
103impl<T: StableType + AsFixedSizeBytes + Hash + Eq> Default for SHashSet<T> {
104    #[inline]
105    fn default() -> Self {
106        SHashSet::new()
107    }
108}
109
110impl<T: StableType + AsFixedSizeBytes + Hash + Eq> AsFixedSizeBytes for SHashSet<T> {
111    const SIZE: usize = SHashMap::<T, ()>::SIZE;
112    type Buf = <SHashMap<T, ()> as AsFixedSizeBytes>::Buf;
113
114    #[inline]
115    fn as_fixed_size_bytes(&self, buf: &mut [u8]) {
116        self.map.as_fixed_size_bytes(buf)
117    }
118
119    #[inline]
120    fn from_fixed_size_bytes(arr: &[u8]) -> Self {
121        let map = SHashMap::<T, ()>::from_fixed_size_bytes(arr);
122        Self { map }
123    }
124}
125
126impl<T: StableType + AsFixedSizeBytes + Hash + Eq> StableType for SHashSet<T> {
127    #[inline]
128    unsafe fn stable_drop_flag_off(&mut self) {
129        self.map.stable_drop_flag_off();
130    }
131
132    #[inline]
133    unsafe fn stable_drop_flag_on(&mut self) {
134        self.map.stable_drop_flag_on();
135    }
136}
137
138impl<T: StableType + AsFixedSizeBytes + Hash + Eq + Debug> Debug for SHashSet<T> {
139    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
140        f.write_str("(")?;
141        for (idx, elem) in self.iter().enumerate() {
142            elem.fmt(f)?;
143
144            if idx < self.len() - 1 {
145                f.write_str(", ")?;
146            }
147        }
148        f.write_str(")")
149    }
150}
151
152#[cfg(test)]
153mod tests {
154    use crate::collections::hash_set::SHashSet;
155    use crate::encoding::{AsFixedSizeBytes, Buffer};
156    use crate::utils::test::generate_random_string;
157    use crate::utils::DebuglessUnwrap;
158    use crate::{
159        _debug_validate_allocator, get_allocated_size, init_allocator, retrieve_custom_data,
160        stable, stable_memory_init, stable_memory_post_upgrade, stable_memory_pre_upgrade,
161        store_custom_data, SBox,
162    };
163    use rand::rngs::ThreadRng;
164    use rand::{thread_rng, Rng};
165    use std::collections::HashSet;
166
167    #[test]
168    fn basic_flow_works_fine() {
169        stable::clear();
170        stable_memory_init();
171
172        {
173            let mut set = SHashSet::default();
174
175            assert!(set.is_empty());
176
177            assert!(!set.insert(10).unwrap());
178            assert!(!set.insert(20).unwrap());
179            assert!(set.insert(10).unwrap());
180
181            assert!(set.contains(&10));
182            assert!(!set.contains(&100));
183
184            assert_eq!(set.len(), 2);
185
186            assert!(!set.remove(&100));
187            assert!(set.remove(&10));
188
189            SHashSet::<u64>::new_with_capacity(10).debugless_unwrap();
190        }
191
192        _debug_validate_allocator();
193        assert_eq!(get_allocated_size(), 0);
194    }
195
196    #[test]
197    fn iter_works_fine() {
198        stable::clear();
199        stable_memory_init();
200
201        {
202            let mut set = SHashSet::default();
203
204            for i in 0..100 {
205                set.insert(i);
206            }
207
208            let mut c = 0;
209            for _ in set.iter() {
210                c += 1;
211            }
212
213            assert_eq!(c, 100);
214        }
215
216        _debug_validate_allocator();
217        assert_eq!(get_allocated_size(), 0);
218    }
219
220    #[test]
221    fn serialization_works_fine() {
222        stable::clear();
223        stable_memory_init();
224
225        {
226            let set = SHashSet::<u32>::default();
227
228            let len = set.len();
229            let cap = set.capacity();
230
231            let buf = set.as_new_fixed_size_bytes();
232            let set1 = SHashSet::<u32>::from_fixed_size_bytes(buf._deref());
233
234            assert_eq!(len, set1.len());
235            assert_eq!(cap, set1.capacity());
236        }
237
238        _debug_validate_allocator();
239        assert_eq!(get_allocated_size(), 0);
240    }
241
242    #[derive(Debug)]
243    enum Action {
244        Insert,
245        Remove,
246        Clear,
247        CanisterUpgrade,
248    }
249
250    struct Fuzzer {
251        set: Option<SHashSet<SBox<String>>>,
252        example: HashSet<String>,
253        keys: Vec<String>,
254        rng: ThreadRng,
255        log: Vec<Action>,
256    }
257
258    impl Fuzzer {
259        fn new() -> Fuzzer {
260            Fuzzer {
261                set: Some(SHashSet::new()),
262                example: HashSet::new(),
263                keys: Vec::new(),
264                rng: thread_rng(),
265                log: Vec::new(),
266            }
267        }
268
269        fn set(&mut self) -> &mut SHashSet<SBox<String>> {
270            self.set.as_mut().unwrap()
271        }
272
273        fn next(&mut self) {
274            let action = self.rng.gen_range(0..101);
275
276            match action {
277                // INSERT ~60%
278                0..=59 => {
279                    let key = generate_random_string(&mut self.rng);
280                    if let Ok(data) = SBox::new(key.clone()) {
281                        if self.set().insert(data).is_err() {
282                            return;
283                        }
284                        self.example.insert(key.clone());
285
286                        match self.keys.binary_search(&key) {
287                            Ok(idx) => {}
288                            Err(idx) => {
289                                self.keys.insert(idx, key);
290                            }
291                        };
292
293                        self.log.push(Action::Insert);
294                    }
295                }
296                // REMOVE
297                60..=89 => {
298                    let len = self.set().len();
299
300                    if len == 0 {
301                        return self.next();
302                    }
303
304                    let idx = self.rng.gen_range(0..len);
305                    let key = self.keys.remove(idx);
306
307                    self.set().remove(&key);
308                    self.example.remove(&key);
309
310                    self.log.push(Action::Remove);
311                }
312                90..=91 => {
313                    self.set().clear();
314                    self.example.clear();
315                    self.keys.clear();
316
317                    self.log.push(Action::Clear);
318                }
319                // CANISTER UPGRADE
320                _ => match SBox::new(self.set.take().unwrap()) {
321                    Ok(data) => {
322                        store_custom_data(1, data);
323
324                        if stable_memory_pre_upgrade().is_ok() {
325                            stable_memory_post_upgrade();
326                        }
327
328                        self.set = retrieve_custom_data::<SHashSet<SBox<String>>>(1)
329                            .map(|it| it.into_inner());
330
331                        self.log.push(Action::CanisterUpgrade);
332                    }
333                    Err(set) => {
334                        self.set = Some(set);
335                    }
336                },
337            }
338
339            _debug_validate_allocator();
340            assert_eq!(self.set().len() as usize, self.example.len());
341
342            for key in self.keys.clone() {
343                assert!(self.set().contains(&key));
344                assert!(self.example.contains(&key));
345            }
346        }
347    }
348
349    #[test]
350    fn fuzzer_works_fine() {
351        stable::clear();
352        init_allocator(0);
353
354        {
355            let mut fuzzer = Fuzzer::new();
356
357            for _ in 0..10_000 {
358                fuzzer.next();
359            }
360        }
361
362        assert_eq!(get_allocated_size(), 0);
363    }
364
365    #[test]
366    fn fuzzer_works_fine_limited_memory() {
367        stable::clear();
368        init_allocator(10);
369
370        {
371            let mut fuzzer = Fuzzer::new();
372
373            for _ in 0..10_000 {
374                fuzzer.next();
375            }
376        }
377
378        assert_eq!(get_allocated_size(), 0);
379    }
380}
ic_stable_memory/collections/hash_set/mod.rs

ic_stable_memory/collections/hash_set/
mod.rs
