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
21pub 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
85pub 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
139pub 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
193pub 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
240pub 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#[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}