Skip to main content

near_sdk/store/iterable_set/
iter.rs

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