near_sdk/store/iterable_set/
mod.rs

1// This suppresses the depreciation warnings for uses of IterableSet in this module
2#![allow(deprecated)]
3
4mod impls;
5mod iter;
6
7pub use self::iter::{Difference, Drain, Intersection, Iter, SymmetricDifference, Union};
8use super::{ERR_INCONSISTENT_STATE, LookupMap};
9use crate::store::Vector;
10use crate::store::key::{Sha256, ToKey};
11use crate::{IntoStorageKey, env};
12use borsh::{BorshDeserialize, BorshSerialize};
13use near_sdk_macros::near;
14use std::borrow::Borrow;
15use std::fmt;
16
17type VecIndex = u32;
18
19/// A lazily loaded storage set that stores its content directly on the storage trie.
20/// This structure is similar to [`near_sdk::store::LookupSet`](crate::store::LookupSet), except
21/// that it keeps track of the elements so that [`IterableSet`] can be iterable among other things.
22///
23/// As with the [`LookupSet`] type, an `IterableSet` requires that the elements
24/// implement the [`BorshSerialize`] and [`Ord`] traits. This can frequently be achieved by
25/// using `#[derive(BorshSerialize, Ord)]`. Some functions also require elements to implement the
26/// [`BorshDeserialize`] trait.
27///
28/// This set stores the values under a hash of the set's `prefix` and [`BorshSerialize`] of the
29/// element using the set's [`ToKey`] implementation.
30///
31/// The default hash function for [`IterableSet`] is [`Sha256`] which uses a syscall
32/// (or host function) built into the NEAR runtime to hash the element. To use a custom function,
33/// use [`with_hasher`]. Alternative builtin hash functions can be found at
34/// [`near_sdk::store::key`](crate::store::key).
35///
36/// # Examples
37///
38/// ```
39/// use near_sdk::store::IterableSet;
40///
41/// // Initializes a set, the generic types can be inferred to `IterableSet<String, Sha256>`
42/// // The `b"a"` parameter is a prefix for the storage keys of this data structure.
43/// let mut set = IterableSet::new(b"a");
44///
45/// set.insert("test".to_string());
46/// assert!(set.contains("test"));
47/// assert!(set.remove("test"));
48/// ```
49///
50/// [`IterableSet`] also implements various binary operations, which allow
51/// for iterating various combinations of two sets.
52///
53/// ```
54/// use near_sdk::store::IterableSet;
55/// use std::collections::HashSet;
56///
57/// let mut set1 = IterableSet::new(b"m");
58/// set1.insert(1);
59/// set1.insert(2);
60/// set1.insert(3);
61///
62/// let mut set2 = IterableSet::new(b"n");
63/// set2.insert(2);
64/// set2.insert(3);
65/// set2.insert(4);
66///
67/// assert_eq!(
68///     set1.union(&set2).collect::<HashSet<_>>(),
69///     [1, 2, 3, 4].iter().collect()
70/// );
71/// assert_eq!(
72///     set1.intersection(&set2).collect::<HashSet<_>>(),
73///     [2, 3].iter().collect()
74/// );
75/// assert_eq!(
76///     set1.difference(&set2).collect::<HashSet<_>>(),
77///     [1].iter().collect()
78/// );
79/// assert_eq!(
80///     set1.symmetric_difference(&set2).collect::<HashSet<_>>(),
81///     [1, 4].iter().collect()
82/// );
83/// ```
84///
85/// [`with_hasher`]: Self::with_hasher
86/// [`LookupSet`]: crate::store::LookupSet
87#[near(inside_nearsdk)]
88pub struct IterableSet<T, H = Sha256>
89where
90    T: BorshSerialize + Ord,
91    H: ToKey,
92{
93    // ser/de is independent of `T` ser/de, `BorshSerialize`/`BorshDeserialize`/`BorshSchema` bounds removed
94    #[cfg_attr(not(feature = "abi"), borsh(bound(serialize = "", deserialize = "")))]
95    #[cfg_attr(
96        feature = "abi",
97        borsh(bound(serialize = "", deserialize = ""), schema(params = ""))
98    )]
99    elements: Vector<T>,
100    // ser/de is independent of `T`,`H` ser/de, `BorshSerialize`/`BorshDeserialize`/`BorshSchema` bounds removed
101    #[cfg_attr(not(feature = "abi"), borsh(bound(serialize = "", deserialize = "")))]
102    #[cfg_attr(
103        feature = "abi",
104        borsh(bound(serialize = "", deserialize = ""), schema(params = ""))
105    )]
106    index: LookupMap<T, VecIndex, H>,
107}
108
109impl<T, H> Drop for IterableSet<T, H>
110where
111    T: BorshSerialize + Ord,
112    H: ToKey,
113{
114    fn drop(&mut self) {
115        self.flush()
116    }
117}
118
119impl<T, H> fmt::Debug for IterableSet<T, H>
120where
121    T: BorshSerialize + Ord + BorshDeserialize + fmt::Debug,
122    H: ToKey,
123{
124    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
125        f.debug_struct("IterableSet")
126            .field("elements", &self.elements)
127            .field("index", &self.index)
128            .finish()
129    }
130}
131
132impl<T> IterableSet<T, Sha256>
133where
134    T: BorshSerialize + Ord,
135{
136    /// Create a new iterable set. Use `prefix` as a unique prefix for keys.
137    ///
138    /// This prefix can be anything that implements [`IntoStorageKey`]. The prefix is used when
139    /// storing and looking up values in storage to ensure no collisions with other collections.
140    ///
141    /// # Examples
142    ///
143    /// ```
144    /// use near_sdk::store::IterableSet;
145    ///
146    /// let mut map: IterableSet<String> = IterableSet::new(b"b");
147    /// ```
148    #[inline]
149    pub fn new<S>(prefix: S) -> Self
150    where
151        S: IntoStorageKey,
152    {
153        Self::with_hasher(prefix)
154    }
155}
156
157impl<T, H> IterableSet<T, H>
158where
159    T: BorshSerialize + Ord,
160    H: ToKey,
161{
162    /// Initialize a [`IterableSet`] with a custom hash function.
163    ///
164    /// # Example
165    /// ```
166    /// use near_sdk::store::key::Keccak256;
167    /// use near_sdk::store::IterableSet;
168    ///
169    /// let map = IterableSet::<String, Keccak256>::with_hasher(b"m");
170    /// ```
171    pub fn with_hasher<S>(prefix: S) -> Self
172    where
173        S: IntoStorageKey,
174    {
175        let mut vec_key = prefix.into_storage_key();
176        let map_key = [vec_key.as_slice(), b"m"].concat();
177        vec_key.push(b'v');
178        Self { elements: Vector::new(vec_key), index: LookupMap::with_hasher(map_key) }
179    }
180
181    /// Returns the number of elements in the set.
182    pub fn len(&self) -> u32 {
183        self.elements.len()
184    }
185
186    /// Returns true if the set contains no elements.
187    pub fn is_empty(&self) -> bool {
188        self.elements.is_empty()
189    }
190
191    /// Clears the set, removing all values.
192    pub fn clear(&mut self)
193    where
194        T: BorshDeserialize + Clone,
195    {
196        for e in self.elements.drain(..) {
197            self.index.set(e, None);
198        }
199    }
200
201    /// Visits the values representing the difference, i.e., the values that are in `self` but not
202    /// in `other`.
203    ///
204    /// # Examples
205    ///
206    /// ```
207    /// use near_sdk::store::IterableSet;
208    ///
209    /// let mut set1 = IterableSet::new(b"m");
210    /// set1.insert("a".to_string());
211    /// set1.insert("b".to_string());
212    /// set1.insert("c".to_string());
213    ///
214    /// let mut set2 = IterableSet::new(b"n");
215    /// set2.insert("b".to_string());
216    /// set2.insert("c".to_string());
217    /// set2.insert("d".to_string());
218    ///
219    /// // Can be seen as `set1 - set2`.
220    /// for x in set1.difference(&set2) {
221    ///     println!("{}", x); // Prints "a"
222    /// }
223    /// ```
224    pub fn difference<'a>(&'a self, other: &'a IterableSet<T, H>) -> Difference<'a, T, H>
225    where
226        T: BorshDeserialize,
227    {
228        Difference::new(self, other)
229    }
230
231    /// Visits the values representing the symmetric difference, i.e., the values that are in
232    /// `self` or in `other` but not in both.
233    ///
234    /// # Examples
235    ///
236    /// ```
237    /// use near_sdk::store::IterableSet;
238    ///
239    /// let mut set1 = IterableSet::new(b"m");
240    /// set1.insert("a".to_string());
241    /// set1.insert("b".to_string());
242    /// set1.insert("c".to_string());
243    ///
244    /// let mut set2 = IterableSet::new(b"n");
245    /// set2.insert("b".to_string());
246    /// set2.insert("c".to_string());
247    /// set2.insert("d".to_string());
248    ///
249    /// // Prints "a", "d" in arbitrary order.
250    /// for x in set1.symmetric_difference(&set2) {
251    ///     println!("{}", x);
252    /// }
253    /// ```
254    pub fn symmetric_difference<'a>(
255        &'a self,
256        other: &'a IterableSet<T, H>,
257    ) -> SymmetricDifference<'a, T, H>
258    where
259        T: BorshDeserialize + Clone,
260    {
261        SymmetricDifference::new(self, other)
262    }
263
264    /// Visits the values representing the intersection, i.e., the values that are both in `self`
265    /// and `other`.
266    ///
267    /// # Examples
268    ///
269    /// ```
270    /// use near_sdk::store::IterableSet;
271    ///
272    /// let mut set1 = IterableSet::new(b"m");
273    /// set1.insert("a".to_string());
274    /// set1.insert("b".to_string());
275    /// set1.insert("c".to_string());
276    ///
277    /// let mut set2 = IterableSet::new(b"n");
278    /// set2.insert("b".to_string());
279    /// set2.insert("c".to_string());
280    /// set2.insert("d".to_string());
281    ///
282    /// // Prints "b", "c" in arbitrary order.
283    /// for x in set1.intersection(&set2) {
284    ///     println!("{}", x);
285    /// }
286    /// ```
287    pub fn intersection<'a>(&'a self, other: &'a IterableSet<T, H>) -> Intersection<'a, T, H>
288    where
289        T: BorshDeserialize,
290    {
291        Intersection::new(self, other)
292    }
293
294    /// Visits the values representing the union, i.e., all the values in `self` or `other`, without
295    /// duplicates.
296    ///
297    /// # Examples
298    ///
299    /// ```
300    /// use near_sdk::store::IterableSet;
301    ///
302    /// let mut set1 = IterableSet::new(b"m");
303    /// set1.insert("a".to_string());
304    /// set1.insert("b".to_string());
305    /// set1.insert("c".to_string());
306    ///
307    /// let mut set2 = IterableSet::new(b"n");
308    /// set2.insert("b".to_string());
309    /// set2.insert("c".to_string());
310    /// set2.insert("d".to_string());
311    ///
312    /// // Prints "a", "b", "c", "d" in arbitrary order.
313    /// for x in set1.union(&set2) {
314    ///     println!("{}", x);
315    /// }
316    /// ```
317    pub fn union<'a>(&'a self, other: &'a IterableSet<T, H>) -> Union<'a, T, H>
318    where
319        T: BorshDeserialize + Clone,
320    {
321        Union::new(self, other)
322    }
323
324    /// Returns `true` if `self` has no elements in common with `other`. This is equivalent to
325    /// checking for an empty intersection.
326    ///
327    /// # Examples
328    ///
329    /// ```
330    /// use near_sdk::store::IterableSet;
331    ///
332    /// let mut set1 = IterableSet::new(b"m");
333    /// set1.insert("a".to_string());
334    /// set1.insert("b".to_string());
335    /// set1.insert("c".to_string());
336    ///
337    /// let mut set2 = IterableSet::new(b"n");
338    ///
339    /// assert_eq!(set1.is_disjoint(&set2), true);
340    /// set2.insert("d".to_string());
341    /// assert_eq!(set1.is_disjoint(&set2), true);
342    /// set2.insert("a".to_string());
343    /// assert_eq!(set1.is_disjoint(&set2), false);
344    /// ```
345    pub fn is_disjoint(&self, other: &IterableSet<T, H>) -> bool
346    where
347        T: BorshDeserialize + Clone,
348    {
349        if self.len() <= other.len() {
350            self.iter().all(|v| !other.contains(v))
351        } else {
352            other.iter().all(|v| !self.contains(v))
353        }
354    }
355
356    /// Returns `true` if the set is a subset of another, i.e., `other` contains at least all the
357    /// values in `self`.
358    ///
359    /// # Examples
360    ///
361    /// ```
362    /// use near_sdk::store::IterableSet;
363    ///
364    /// let mut sup = IterableSet::new(b"m");
365    /// sup.insert("a".to_string());
366    /// sup.insert("b".to_string());
367    /// sup.insert("c".to_string());
368    ///
369    /// let mut set = IterableSet::new(b"n");
370    ///
371    /// assert_eq!(set.is_subset(&sup), true);
372    /// set.insert("b".to_string());
373    /// assert_eq!(set.is_subset(&sup), true);
374    /// set.insert("d".to_string());
375    /// assert_eq!(set.is_subset(&sup), false);
376    /// ```
377    pub fn is_subset(&self, other: &IterableSet<T, H>) -> bool
378    where
379        T: BorshDeserialize + Clone,
380    {
381        if self.len() <= other.len() { self.iter().all(|v| other.contains(v)) } else { false }
382    }
383
384    /// Returns `true` if the set is a superset of another, i.e., `self` contains at least all the
385    /// values in `other`.
386    ///
387    /// # Examples
388    ///
389    /// ```
390    /// use near_sdk::store::IterableSet;
391    ///
392    /// let mut sub = IterableSet::new(b"m");
393    /// sub.insert("a".to_string());
394    /// sub.insert("b".to_string());
395    ///
396    /// let mut set = IterableSet::new(b"n");
397    ///
398    /// assert_eq!(set.is_superset(&sub), false);
399    /// set.insert("b".to_string());
400    /// set.insert("d".to_string());
401    /// assert_eq!(set.is_superset(&sub), false);
402    /// set.insert("a".to_string());
403    /// assert_eq!(set.is_superset(&sub), true);
404    /// ```
405    pub fn is_superset(&self, other: &IterableSet<T, H>) -> bool
406    where
407        T: BorshDeserialize + Clone,
408    {
409        other.is_subset(self)
410    }
411
412    /// An iterator visiting all elements in arbitrary order.
413    /// The iterator element type is `&'a T`.
414    ///
415    /// # Examples
416    ///
417    /// ```
418    /// use near_sdk::store::IterableSet;
419    ///
420    /// let mut set = IterableSet::new(b"m");
421    /// set.insert("a".to_string());
422    /// set.insert("b".to_string());
423    /// set.insert("c".to_string());
424    ///
425    /// for val in set.iter() {
426    ///     println!("val: {}", val);
427    /// }
428    /// ```
429    pub fn iter(&self) -> Iter<'_, T>
430    where
431        T: BorshDeserialize,
432    {
433        Iter::new(self)
434    }
435
436    /// Clears the set, returning all elements in an iterator.
437    ///
438    /// This will clear all elements, even if only some of them are yielded.
439    ///
440    /// # Examples
441    ///
442    /// ```
443    /// use near_sdk::store::IterableSet;
444    ///
445    /// let mut a = IterableSet::new(b"m");
446    /// a.insert(1);
447    /// a.insert(2);
448    ///
449    /// for v in a.drain().take(1) {
450    ///     assert!(v == 1 || v == 2);
451    /// }
452    ///
453    /// assert!(a.is_empty());
454    /// ```
455    pub fn drain(&mut self) -> Drain<'_, T, H>
456    where
457        T: BorshDeserialize,
458    {
459        Drain::new(self)
460    }
461
462    /// Returns `true` if the set contains the specified value.
463    ///
464    /// The value may be any borrowed form of the set's value type, but
465    /// [`BorshSerialize`], [`ToOwned<Owned = T>`](ToOwned) and [`Ord`] on the borrowed form *must*
466    /// match those for the value type.
467    pub fn contains<Q: ?Sized>(&self, value: &Q) -> bool
468    where
469        T: Borrow<Q>,
470        Q: BorshSerialize + ToOwned<Owned = T> + Ord,
471    {
472        self.index.contains_key(value)
473    }
474
475    /// Adds a value to the set.
476    ///
477    /// If the set did not have this value present, true is returned.
478    ///
479    /// If the set did have this value present, false is returned.
480    pub fn insert(&mut self, value: T) -> bool
481    where
482        T: Clone + BorshDeserialize,
483    {
484        let entry = self.index.get_mut_inner(&value);
485        if entry.value_mut().is_some() {
486            false
487        } else {
488            self.elements.push(value);
489            let element_index = self.elements.len() - 1;
490            entry.replace(Some(element_index));
491            true
492        }
493    }
494
495    /// Removes a value from the set. Returns whether the value was present in the set.
496    ///
497    /// The value may be any borrowed form of the set's value type, but
498    /// [`BorshSerialize`], [`ToOwned<Owned = K>`](ToOwned) and [`Ord`] on the borrowed form *must*
499    /// match those for the value type.
500    ///
501    /// # Performance
502    ///
503    /// When elements are removed, the underlying vector of keys is rearranged by means of swapping
504    /// an obsolete key with the last element in the list and deleting that. Note that that requires
505    /// updating the `index` map due to the fact that it holds `elements` vector indices.
506    pub fn remove<Q: ?Sized>(&mut self, value: &Q) -> bool
507    where
508        T: Borrow<Q> + BorshDeserialize + Clone,
509        Q: BorshSerialize + ToOwned<Owned = T> + Ord,
510    {
511        match self.index.remove(value) {
512            Some(element_index) => {
513                let last_index = self.elements.len() - 1;
514                let _ = self.elements.swap_remove(element_index);
515
516                match element_index {
517                    // If it's the last/only element - do nothing.
518                    x if x == last_index => {}
519                    // Otherwise update it's index.
520                    _ => {
521                        let element = self
522                            .elements
523                            .get(element_index)
524                            .unwrap_or_else(|| env::panic_str(ERR_INCONSISTENT_STATE));
525                        self.index.set(element.clone(), Some(element_index));
526                    }
527                }
528
529                true
530            }
531            None => false,
532        }
533    }
534
535    /// Flushes the intermediate values of the map before this is called when the structure is
536    /// [`Drop`]ed. This will write all modified values to storage but keep all cached values
537    /// in memory.
538    pub fn flush(&mut self) {
539        self.elements.flush();
540        self.index.flush();
541    }
542}
543
544#[cfg(not(target_arch = "wasm32"))]
545#[cfg(test)]
546mod tests {
547    use crate::store::IterableSet;
548    use crate::test_utils::test_env::setup_free;
549    use arbitrary::{Arbitrary, Unstructured};
550    use borsh::{BorshDeserialize, to_vec};
551    use rand::RngCore;
552    use rand::SeedableRng;
553    use std::collections::HashSet;
554
555    #[test]
556    fn basic_functionality() {
557        let mut set = IterableSet::new(b"b");
558        assert!(set.is_empty());
559        assert!(set.insert("test".to_string()));
560        assert!(set.contains("test"));
561        assert_eq!(set.len(), 1);
562
563        assert!(set.remove("test"));
564        assert_eq!(set.len(), 0);
565    }
566
567    #[test]
568    fn set_iterator() {
569        let mut set = IterableSet::new(b"b");
570
571        set.insert(0u8);
572        set.insert(1);
573        set.insert(2);
574        set.insert(3);
575        set.remove(&1);
576        let iter = set.iter();
577        assert_eq!(iter.len(), 3);
578        assert_eq!(iter.collect::<Vec<_>>(), [(&0), (&3), (&2)]);
579
580        let mut iter = set.iter();
581        assert_eq!(iter.nth(2), Some(&2));
582        // Check fused iterator assumption that each following one will be None
583        assert_eq!(iter.next(), None);
584
585        // Drain
586        assert_eq!(set.drain().collect::<Vec<_>>(), [0, 3, 2]);
587        assert!(set.is_empty());
588    }
589
590    #[test]
591    fn test_drain() {
592        let mut s = IterableSet::new(b"m");
593        s.extend(1..100);
594
595        // Drain the set a few times to make sure that it does have any random residue
596        for _ in 0..20 {
597            assert_eq!(s.len(), 99);
598
599            for _ in s.drain() {}
600
601            #[allow(clippy::never_loop)]
602            for _ in &s {
603                panic!("s should be empty!");
604            }
605
606            assert_eq!(s.len(), 0);
607            assert!(s.is_empty());
608
609            s.extend(1..100);
610        }
611    }
612
613    #[test]
614    fn test_extend() {
615        let mut a = IterableSet::<u64>::new(b"m");
616        a.insert(1);
617
618        a.extend([2, 3, 4]);
619
620        assert_eq!(a.len(), 4);
621        assert!(a.contains(&1));
622        assert!(a.contains(&2));
623        assert!(a.contains(&3));
624        assert!(a.contains(&4));
625    }
626
627    #[test]
628    fn test_difference() {
629        let mut set1 = IterableSet::new(b"m");
630        set1.insert("a".to_string());
631        set1.insert("b".to_string());
632        set1.insert("c".to_string());
633        set1.insert("d".to_string());
634
635        let mut set2 = IterableSet::new(b"n");
636        set2.insert("b".to_string());
637        set2.insert("c".to_string());
638        set2.insert("e".to_string());
639
640        assert_eq!(
641            set1.difference(&set2).collect::<HashSet<_>>(),
642            ["a".to_string(), "d".to_string()].iter().collect::<HashSet<_>>()
643        );
644        assert_eq!(
645            set2.difference(&set1).collect::<HashSet<_>>(),
646            ["e".to_string()].iter().collect::<HashSet<_>>()
647        );
648        assert!(set1.difference(&set2).nth(1).is_some());
649        assert!(set1.difference(&set2).nth(2).is_none());
650    }
651
652    #[test]
653    fn test_difference_empty() {
654        let mut set1 = IterableSet::new(b"m");
655        set1.insert(1);
656        set1.insert(2);
657        set1.insert(3);
658
659        let mut set2 = IterableSet::new(b"n");
660        set2.insert(3);
661        set2.insert(1);
662        set2.insert(2);
663        set2.insert(4);
664
665        assert_eq!(set1.difference(&set2).collect::<HashSet<_>>(), HashSet::new());
666    }
667
668    #[test]
669    fn test_symmetric_difference() {
670        let mut set1 = IterableSet::new(b"m");
671        set1.insert("a".to_string());
672        set1.insert("b".to_string());
673        set1.insert("c".to_string());
674
675        let mut set2 = IterableSet::new(b"n");
676        set2.insert("b".to_string());
677        set2.insert("c".to_string());
678        set2.insert("d".to_string());
679
680        assert_eq!(
681            set1.symmetric_difference(&set2).collect::<HashSet<_>>(),
682            ["a".to_string(), "d".to_string()].iter().collect::<HashSet<_>>()
683        );
684        assert_eq!(
685            set2.symmetric_difference(&set1).collect::<HashSet<_>>(),
686            ["a".to_string(), "d".to_string()].iter().collect::<HashSet<_>>()
687        );
688    }
689
690    #[test]
691    fn test_symmetric_difference_empty() {
692        let mut set1 = IterableSet::new(b"m");
693        set1.insert(1);
694        set1.insert(2);
695        set1.insert(3);
696
697        let mut set2 = IterableSet::new(b"n");
698        set2.insert(3);
699        set2.insert(1);
700        set2.insert(2);
701
702        assert_eq!(set1.symmetric_difference(&set2).collect::<HashSet<_>>(), HashSet::new());
703    }
704
705    #[test]
706    fn test_intersection() {
707        let mut set1 = IterableSet::new(b"m");
708        set1.insert("a".to_string());
709        set1.insert("b".to_string());
710        set1.insert("c".to_string());
711
712        let mut set2 = IterableSet::new(b"n");
713        set2.insert("b".to_string());
714        set2.insert("c".to_string());
715        set2.insert("d".to_string());
716
717        assert_eq!(
718            set1.intersection(&set2).collect::<HashSet<_>>(),
719            ["b".to_string(), "c".to_string()].iter().collect::<HashSet<_>>()
720        );
721        assert_eq!(
722            set2.intersection(&set1).collect::<HashSet<_>>(),
723            ["b".to_string(), "c".to_string()].iter().collect::<HashSet<_>>()
724        );
725        assert!(set1.intersection(&set2).nth(1).is_some());
726        assert!(set1.intersection(&set2).nth(2).is_none());
727    }
728
729    #[test]
730    fn test_intersection_empty() {
731        let mut set1 = IterableSet::new(b"m");
732        set1.insert(1);
733        set1.insert(2);
734        set1.insert(3);
735
736        let mut set2 = IterableSet::new(b"n");
737        set2.insert(4);
738        set2.insert(6);
739        set2.insert(5);
740
741        assert_eq!(set1.intersection(&set2).collect::<HashSet<_>>(), HashSet::new());
742    }
743
744    #[test]
745    fn test_union() {
746        let mut set1 = IterableSet::new(b"m");
747        set1.insert("a".to_string());
748        set1.insert("b".to_string());
749        set1.insert("c".to_string());
750
751        let mut set2 = IterableSet::new(b"n");
752        set2.insert("b".to_string());
753        set2.insert("c".to_string());
754        set2.insert("d".to_string());
755
756        assert_eq!(
757            set1.union(&set2).collect::<HashSet<_>>(),
758            ["a".to_string(), "b".to_string(), "c".to_string(), "d".to_string()]
759                .iter()
760                .collect::<HashSet<_>>()
761        );
762        assert_eq!(
763            set2.union(&set1).collect::<HashSet<_>>(),
764            ["a".to_string(), "b".to_string(), "c".to_string(), "d".to_string()]
765                .iter()
766                .collect::<HashSet<_>>()
767        );
768    }
769
770    #[test]
771    fn test_union_empty() {
772        let set1 = IterableSet::<u64>::new(b"m");
773        let set2 = IterableSet::<u64>::new(b"n");
774
775        assert_eq!(set1.union(&set2).collect::<HashSet<_>>(), HashSet::new());
776    }
777
778    #[test]
779    fn test_subset_and_superset() {
780        let mut a = IterableSet::new(b"m");
781        assert!(a.insert(0));
782        assert!(a.insert(50));
783        assert!(a.insert(110));
784        assert!(a.insert(70));
785
786        let mut b = IterableSet::new(b"n");
787        assert!(b.insert(0));
788        assert!(b.insert(70));
789        assert!(b.insert(190));
790        assert!(b.insert(2500));
791        assert!(b.insert(110));
792        assert!(b.insert(2000));
793
794        assert!(!a.is_subset(&b));
795        assert!(!a.is_superset(&b));
796        assert!(!b.is_subset(&a));
797        assert!(!b.is_superset(&a));
798
799        assert!(b.insert(50));
800
801        assert!(a.is_subset(&b));
802        assert!(!a.is_superset(&b));
803        assert!(!b.is_subset(&a));
804        assert!(b.is_superset(&a));
805    }
806
807    #[test]
808    fn test_disjoint() {
809        let mut xs = IterableSet::new(b"m");
810        let mut ys = IterableSet::new(b"n");
811
812        assert!(xs.is_disjoint(&ys));
813        assert!(ys.is_disjoint(&xs));
814
815        assert!(xs.insert(50));
816        assert!(ys.insert(110));
817        assert!(xs.is_disjoint(&ys));
818        assert!(ys.is_disjoint(&xs));
819
820        assert!(xs.insert(70));
821        assert!(xs.insert(190));
822        assert!(xs.insert(40));
823        assert!(ys.insert(20));
824        assert!(ys.insert(-110));
825        assert!(xs.is_disjoint(&ys));
826        assert!(ys.is_disjoint(&xs));
827
828        assert!(ys.insert(70));
829        assert!(!xs.is_disjoint(&ys));
830        assert!(!ys.is_disjoint(&xs));
831    }
832
833    #[derive(Arbitrary, Debug)]
834    enum Op {
835        Insert(u8),
836        Remove(u8),
837        Flush,
838        Restore,
839        Contains(u8),
840    }
841
842    #[test]
843    fn arbitrary() {
844        setup_free();
845
846        let mut rng = rand_xorshift::XorShiftRng::seed_from_u64(0);
847        let mut buf = vec![0; 4096];
848        for _ in 0..512 {
849            // Clear storage in-between runs
850            crate::mock::with_mocked_blockchain(|b| b.take_storage());
851            rng.fill_bytes(&mut buf);
852
853            let mut us = IterableSet::new(b"l");
854            let mut hs = HashSet::new();
855            let u = Unstructured::new(&buf);
856            if let Ok(ops) = Vec::<Op>::arbitrary_take_rest(u) {
857                for op in ops {
858                    match op {
859                        Op::Insert(v) => {
860                            let r1 = us.insert(v);
861                            let r2 = hs.insert(v);
862                            assert_eq!(r1, r2)
863                        }
864                        Op::Remove(v) => {
865                            let r1 = us.remove(&v);
866                            let r2 = hs.remove(&v);
867                            assert_eq!(r1, r2)
868                        }
869                        Op::Flush => {
870                            us.flush();
871                        }
872                        Op::Restore => {
873                            let serialized = to_vec(&us).unwrap();
874                            us = IterableSet::deserialize(&mut serialized.as_slice()).unwrap();
875                        }
876                        Op::Contains(v) => {
877                            let r1 = us.contains(&v);
878                            let r2 = hs.contains(&v);
879                            assert_eq!(r1, r2)
880                        }
881                    }
882                }
883            }
884        }
885    }
886
887    #[cfg(feature = "abi")]
888    #[test]
889    fn test_borsh_schema() {
890        #[derive(
891            borsh::BorshSerialize, borsh::BorshDeserialize, PartialEq, Eq, PartialOrd, Ord,
892        )]
893        struct NoSchemaStruct;
894
895        assert_eq!(
896            "IterableSet".to_string(),
897            <IterableSet<NoSchemaStruct> as borsh::BorshSchema>::declaration()
898        );
899        let mut defs = Default::default();
900        <IterableSet<NoSchemaStruct> as borsh::BorshSchema>::add_definitions_recursively(&mut defs);
901
902        insta::assert_snapshot!(format!("{:#?}", defs));
903    }
904
905    #[test]
906    fn test_drain_next_back_removes_from_index() {
907        let mut set = IterableSet::new(b"t");
908
909        // Insert elements
910        for i in 0..10 {
911            set.insert(i);
912        }
913
914        // Drain from the back using next_back()
915        let mut drain = set.drain();
916        let last = drain.next_back().unwrap();
917        drop(drain);
918
919        // If the bug exists: contains() will return true (stale index entry)
920        // If bug is fixed: contains() will return false
921        assert!(
922            !set.contains(&last),
923            "Element {} was drained but contains() returns true (stale index)",
924            last
925        );
926
927        // If the bug exists: insert should fail silently or panic
928        // If bug is fixed: insert should succeed
929        assert!(
930            set.insert(last),
931            "Cannot re-insert element {} after draining it (stale index)",
932            last
933        );
934    }
935
936    #[test]
937    fn test_drain_bidirectional() {
938        let mut set = IterableSet::new(b"t");
939        for i in 0..10 {
940            set.insert(i);
941        }
942
943        let mut drain = set.drain();
944        let first = drain.next().unwrap();
945        let last = drain.next_back().unwrap();
946        drop(drain);
947
948        // Both should be gone
949        assert!(!set.contains(&first), "Element {} from next() still in index", first);
950        assert!(!set.contains(&last), "Element {} from next_back() still in index", last);
951
952        // Should be able to re-insert both
953        assert!(set.insert(first));
954        assert!(set.insert(last));
955    }
956
957    #[test]
958    fn test_drain_partial_consumption_clears_index() {
959        let mut set = IterableSet::new(b"t");
960        for i in 1..=3 {
961            set.insert(i);
962        }
963
964        {
965            let mut drain = set.drain();
966            drain.next().unwrap();
967            // Dropped with two elements unconsumed
968        }
969
970        assert_eq!(set.len(), 0);
971        assert!(set.is_empty());
972
973        for i in 1..=3 {
974            assert!(!set.contains(&i), "stale index entry for {} after partial drain", i);
975            assert!(!set.remove(&i), "remove({}) should be a no-op on an empty set", i);
976            assert!(set.insert(i), "re-insert of {} should succeed after drain", i);
977        }
978        assert_eq!(set.len(), 3);
979    }
980
981    #[test]
982    fn test_drain_unconsumed_clears_index() {
983        let mut set = IterableSet::new(b"t");
984        for i in 1..=3 {
985            set.insert(i);
986        }
987
988        set.drain(); // Bare statement: dropped without consuming anything
989
990        assert!(set.is_empty());
991        for i in 1..=3 {
992            assert!(!set.contains(&i), "stale index entry for {} after unconsumed drain", i);
993        }
994    }
995
996    #[test]
997    fn test_drain_mixed_direction_partial_consumption() {
998        let mut set = IterableSet::new(b"t");
999        for i in 1..=5 {
1000            set.insert(i);
1001        }
1002
1003        {
1004            let mut drain = set.drain();
1005            drain.next().unwrap();
1006            drain.next_back().unwrap();
1007            // Dropped with three elements unconsumed
1008        }
1009
1010        assert!(set.is_empty());
1011        for i in 1..=5 {
1012            assert!(!set.contains(&i), "stale index entry for {} after mixed-direction drain", i);
1013            assert!(set.insert(i), "re-insert of {} should succeed after drain", i);
1014        }
1015        assert_eq!(set.len(), 5);
1016    }
1017}
near_sdk/store/iterable_set/mod.rs

near_sdk/store/iterable_set/
mod.rs
