unc_sdk/store/unordered_set/
iter.rs

1use super::UnorderedSet;
2use crate::store::free_list::FreeListIndex;
3use crate::store::key::ToKey;
4use crate::store::{free_list, LookupMap};
5use borsh::{BorshDeserialize, BorshSerialize};
6use std::iter::{Chain, FusedIterator};
7
8impl<'a, T, H> IntoIterator for &'a UnorderedSet<T, H>
9where
10    T: BorshSerialize + Ord + BorshDeserialize + Clone,
11    H: ToKey,
12{
13    type Item = &'a T;
14    type IntoIter = Iter<'a, T>;
15
16    fn into_iter(self) -> Self::IntoIter {
17        self.iter()
18    }
19}
20
21/// An iterator over elements of a [`UnorderedSet`].
22///
23/// This `struct` is created by the [`iter`] method on [`UnorderedSet`].
24/// See its documentation for more.
25///
26/// [`iter`]: UnorderedSet::iter
27pub struct Iter<'a, T>
28where
29    T: BorshSerialize + Ord + BorshDeserialize,
30{
31    elements: free_list::Iter<'a, T>,
32}
33
34impl<'a, T> Iter<'a, T>
35where
36    T: BorshSerialize + Ord + BorshDeserialize,
37{
38    pub(super) fn new<H>(set: &'a UnorderedSet<T, H>) -> Self
39    where
40        H: ToKey,
41    {
42        Self { elements: set.elements.iter() }
43    }
44}
45
46impl<'a, T> Iterator for Iter<'a, T>
47where
48    T: BorshSerialize + Ord + BorshDeserialize,
49{
50    type Item = &'a T;
51
52    fn next(&mut self) -> Option<Self::Item> {
53        <Self as Iterator>::nth(self, 0)
54    }
55
56    fn size_hint(&self) -> (usize, Option<usize>) {
57        self.elements.size_hint()
58    }
59
60    fn count(self) -> usize {
61        self.elements.count()
62    }
63
64    fn nth(&mut self, n: usize) -> Option<Self::Item> {
65        self.elements.nth(n)
66    }
67}
68
69impl<'a, T> ExactSizeIterator for Iter<'a, T> where T: BorshSerialize + Ord + BorshDeserialize {}
70impl<'a, T> FusedIterator for Iter<'a, T> where T: BorshSerialize + Ord + BorshDeserialize {}
71
72impl<'a, T> DoubleEndedIterator for Iter<'a, T>
73where
74    T: BorshSerialize + Ord + BorshDeserialize,
75{
76    fn next_back(&mut self) -> Option<Self::Item> {
77        <Self as DoubleEndedIterator>::nth_back(self, 0)
78    }
79
80    fn nth_back(&mut self, n: usize) -> Option<Self::Item> {
81        self.elements.nth_back(n)
82    }
83}
84
85/// A lazy iterator producing elements in the difference of `UnorderedSet`s.
86///
87/// This `struct` is created by the [`difference`] method on [`UnorderedSet`].
88/// See its documentation for more.
89///
90/// [`difference`]: UnorderedSet::difference
91pub struct Difference<'a, T, H>
92where
93    T: BorshSerialize + Ord + BorshDeserialize,
94    H: ToKey,
95{
96    elements: free_list::Iter<'a, T>,
97
98    other: &'a UnorderedSet<T, H>,
99}
100
101impl<'a, T, H> Difference<'a, T, H>
102where
103    T: BorshSerialize + Ord + BorshDeserialize,
104    H: ToKey,
105{
106    pub(super) fn new(set: &'a UnorderedSet<T, H>, other: &'a UnorderedSet<T, H>) -> Self {
107        Self { elements: set.elements.iter(), other }
108    }
109}
110
111impl<'a, T, H> Iterator for Difference<'a, T, H>
112where
113    T: BorshSerialize + Ord + BorshDeserialize + Clone,
114    H: ToKey,
115{
116    type Item = &'a T;
117
118    fn next(&mut self) -> Option<Self::Item> {
119        loop {
120            let elt = self.elements.next()?;
121            if !self.other.contains(elt) {
122                return Some(elt);
123            }
124        }
125    }
126
127    fn size_hint(&self) -> (usize, Option<usize>) {
128        (0, self.elements.size_hint().1)
129    }
130}
131
132impl<'a, T, H> FusedIterator for Difference<'a, T, H>
133where
134    T: BorshSerialize + Ord + BorshDeserialize + Clone,
135    H: ToKey,
136{
137}
138
139/// A lazy iterator producing elements in the intersection of `UnorderedSet`s.
140///
141/// This `struct` is created by the [`intersection`] method on [`UnorderedSet`].
142/// See its documentation for more.
143///
144/// [`intersection`]: UnorderedSet::intersection
145pub struct Intersection<'a, T, H>
146where
147    T: BorshSerialize + Ord + BorshDeserialize,
148    H: ToKey,
149{
150    elements: free_list::Iter<'a, T>,
151
152    other: &'a UnorderedSet<T, H>,
153}
154
155impl<'a, T, H> Intersection<'a, T, H>
156where
157    T: BorshSerialize + Ord + BorshDeserialize,
158    H: ToKey,
159{
160    pub(super) fn new(set: &'a UnorderedSet<T, H>, other: &'a UnorderedSet<T, H>) -> Self {
161        Self { elements: set.elements.iter(), other }
162    }
163}
164
165impl<'a, T, H> Iterator for Intersection<'a, T, H>
166where
167    T: BorshSerialize + Ord + BorshDeserialize + Clone,
168    H: ToKey,
169{
170    type Item = &'a T;
171
172    fn next(&mut self) -> Option<Self::Item> {
173        loop {
174            let elt = self.elements.next()?;
175            if self.other.contains(elt) {
176                return Some(elt);
177            }
178        }
179    }
180
181    fn size_hint(&self) -> (usize, Option<usize>) {
182        (0, self.elements.size_hint().1)
183    }
184}
185
186impl<'a, T, H> FusedIterator for Intersection<'a, T, H>
187where
188    T: BorshSerialize + Ord + BorshDeserialize + Clone,
189    H: ToKey,
190{
191}
192
193/// A lazy iterator producing elements in the symmetrical difference of [`UnorderedSet`]s.
194///
195/// This `struct` is created by the [`symmetric_difference`] method on [`UnorderedSet`].
196/// See its documentation for more.
197///
198/// [`symmetric_difference`]: UnorderedSet::symmetric_difference
199pub struct SymmetricDifference<'a, T, H>
200where
201    T: BorshSerialize + Ord + BorshDeserialize,
202    H: ToKey,
203{
204    iter: Chain<Difference<'a, T, H>, Difference<'a, T, H>>,
205}
206
207impl<'a, T, H> SymmetricDifference<'a, T, H>
208where
209    T: BorshSerialize + Ord + BorshDeserialize + Clone,
210    H: ToKey,
211{
212    pub(super) fn new(set: &'a UnorderedSet<T, H>, other: &'a UnorderedSet<T, H>) -> Self {
213        Self { iter: set.difference(other).chain(other.difference(set)) }
214    }
215}
216
217impl<'a, T, H> Iterator for SymmetricDifference<'a, T, H>
218where
219    T: BorshSerialize + Ord + BorshDeserialize + Clone,
220    H: ToKey,
221{
222    type Item = &'a T;
223
224    fn next(&mut self) -> Option<Self::Item> {
225        self.iter.next()
226    }
227
228    fn size_hint(&self) -> (usize, Option<usize>) {
229        self.iter.size_hint()
230    }
231}
232
233impl<'a, T, H> FusedIterator for SymmetricDifference<'a, T, H>
234where
235    T: BorshSerialize + Ord + BorshDeserialize + Clone,
236    H: ToKey,
237{
238}
239
240/// A lazy iterator producing elements in the union of `UnorderedSet`s.
241///
242/// This `struct` is created by the [`union`] method on [`UnorderedSet`].
243/// See its documentation for more.
244///
245/// [`union`]: UnorderedSet::union
246pub struct Union<'a, T, H>
247where
248    T: BorshSerialize + Ord + BorshDeserialize,
249    H: ToKey,
250{
251    iter: Chain<Iter<'a, T>, Difference<'a, T, H>>,
252}
253
254impl<'a, T, H> Union<'a, T, H>
255where
256    T: BorshSerialize + Ord + BorshDeserialize + Clone,
257    H: ToKey,
258{
259    pub(super) fn new(set: &'a UnorderedSet<T, H>, other: &'a UnorderedSet<T, H>) -> Self {
260        Self { iter: set.iter().chain(other.difference(set)) }
261    }
262}
263
264impl<'a, T, H> Iterator for Union<'a, T, H>
265where
266    T: BorshSerialize + Ord + BorshDeserialize + Clone,
267    H: ToKey,
268{
269    type Item = &'a T;
270
271    fn next(&mut self) -> Option<Self::Item> {
272        self.iter.next()
273    }
274
275    fn size_hint(&self) -> (usize, Option<usize>) {
276        self.iter.size_hint()
277    }
278}
279
280impl<'a, T, H> FusedIterator for Union<'a, T, H>
281where
282    T: BorshSerialize + Ord + BorshDeserialize + Clone,
283    H: ToKey,
284{
285}
286
287/// A draining iterator for [`UnorderedSet`].
288///
289/// This `struct` is created by the [`drain`] method on [`UnorderedSet`].
290/// See its documentation for more.
291///
292/// [`drain`]: UnorderedSet::drain
293#[derive(Debug)]
294pub struct Drain<'a, T, H>
295where
296    T: BorshSerialize + BorshDeserialize + Ord,
297    H: ToKey,
298{
299    elements: free_list::Drain<'a, T>,
300
301    index: &'a mut LookupMap<T, FreeListIndex, H>,
302}
303
304impl<'a, T, H> Drain<'a, T, H>
305where
306    T: BorshSerialize + BorshDeserialize + Ord,
307    H: ToKey,
308{
309    pub(crate) fn new(set: &'a mut UnorderedSet<T, H>) -> Self {
310        Self { elements: set.elements.drain(), index: &mut set.index }
311    }
312
313    fn remaining(&self) -> usize {
314        self.elements.remaining()
315    }
316}
317
318impl<'a, T, H> Iterator for Drain<'a, T, H>
319where
320    T: BorshSerialize + BorshDeserialize + Ord + Clone,
321    H: ToKey,
322{
323    type Item = T;
324
325    fn next(&mut self) -> Option<Self::Item> {
326        let key = self.elements.next()?;
327        self.index.remove(&key);
328        Some(key)
329    }
330
331    fn size_hint(&self) -> (usize, Option<usize>) {
332        let remaining = self.remaining();
333        (remaining, Some(remaining))
334    }
335
336    fn count(self) -> usize {
337        self.remaining()
338    }
339}
340
341impl<'a, T, H> ExactSizeIterator for Drain<'a, T, H>
342where
343    T: BorshSerialize + Ord + BorshDeserialize + Clone,
344    H: ToKey,
345{
346}
347
348impl<'a, T, H> FusedIterator for Drain<'a, T, H>
349where
350    T: BorshSerialize + Ord + BorshDeserialize + Clone,
351    H: ToKey,
352{
353}
354
355impl<'a, T, H> DoubleEndedIterator for Drain<'a, T, H>
356where
357    T: BorshSerialize + Ord + BorshDeserialize + Clone,
358    H: ToKey,
359{
360    fn next_back(&mut self) -> Option<Self::Item> {
361        self.elements.next_back()
362    }
363}