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#[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
86pub 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
140pub 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
194pub 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
241pub 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#[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}