unc_sdk/store/vec/
iter.rs

1use borsh::{BorshDeserialize, BorshSerialize};
2use core::{iter::FusedIterator, ops::Range};
3
4use super::{Vector, ERR_INDEX_OUT_OF_BOUNDS};
5use crate::env;
6
7/// An iterator over references to each element in the stored vector.
8#[derive(Debug)]
9pub struct Iter<'a, T>
10where
11    T: BorshSerialize + BorshDeserialize,
12{
13    /// Underlying vector to iterate through
14    vec: &'a Vector<T>,
15    /// Range of indices to iterate.
16    range: Range<u32>,
17}
18
19impl<'a, T> Iter<'a, T>
20where
21    T: BorshSerialize + BorshDeserialize,
22{
23    pub(super) fn new(vec: &'a Vector<T>) -> Self {
24        Self { vec, range: Range { start: 0, end: vec.len() } }
25    }
26
27    /// Returns number of elements left to iterate.
28    fn remaining(&self) -> usize {
29        self.range.len()
30    }
31}
32
33impl<'a, T> Iterator for Iter<'a, T>
34where
35    T: BorshSerialize + BorshDeserialize,
36{
37    type Item = &'a T;
38
39    fn next(&mut self) -> Option<Self::Item> {
40        <Self as Iterator>::nth(self, 0)
41    }
42
43    fn size_hint(&self) -> (usize, Option<usize>) {
44        let remaining = self.remaining();
45        (remaining, Some(remaining))
46    }
47
48    fn count(self) -> usize {
49        self.remaining()
50    }
51
52    fn nth(&mut self, n: usize) -> Option<Self::Item> {
53        let idx = self.range.nth(n)?;
54        Some(self.vec.get(idx).unwrap_or_else(|| env::panic_str(ERR_INDEX_OUT_OF_BOUNDS)))
55    }
56}
57
58impl<'a, T> ExactSizeIterator for Iter<'a, T> where T: BorshSerialize + BorshDeserialize {}
59impl<'a, T> FusedIterator for Iter<'a, T> where T: BorshSerialize + BorshDeserialize {}
60
61impl<'a, T> DoubleEndedIterator for Iter<'a, T>
62where
63    T: BorshSerialize + BorshDeserialize,
64{
65    fn next_back(&mut self) -> Option<Self::Item> {
66        <Self as DoubleEndedIterator>::nth_back(self, 0)
67    }
68
69    fn nth_back(&mut self, n: usize) -> Option<Self::Item> {
70        let idx = self.range.nth_back(n)?;
71        Some(self.vec.get(idx).unwrap_or_else(|| env::panic_str(ERR_INDEX_OUT_OF_BOUNDS)))
72    }
73}
74
75/// An iterator over exclusive references to each element of a stored vector.
76#[derive(Debug)]
77pub struct IterMut<'a, T>
78where
79    T: BorshSerialize + BorshDeserialize,
80{
81    /// Mutable reference to vector used to iterate through.
82    vec: &'a mut Vector<T>,
83    /// Range of indices to iterate.
84    range: Range<u32>,
85}
86
87impl<'a, T> IterMut<'a, T>
88where
89    T: BorshSerialize + BorshDeserialize,
90{
91    /// Creates a new iterator for the given storage vector.
92    pub(crate) fn new(vec: &'a mut Vector<T>) -> Self {
93        let end = vec.len();
94        Self { vec, range: Range { start: 0, end } }
95    }
96
97    /// Returns the amount of remaining elements to yield by the iterator.
98    fn remaining(&self) -> usize {
99        self.range.len()
100    }
101}
102
103impl<'a, T> IterMut<'a, T>
104where
105    T: BorshSerialize + BorshDeserialize,
106{
107    fn get_mut<'b>(&'b mut self, at: u32) -> Option<&'a mut T> {
108        self.vec.get_mut(at).map(|value| {
109            //* SAFETY: The lifetime can be swapped here because we can assert that the iterator
110            //*         will only give out one mutable reference for every individual item
111            //*         during the iteration, and there is no overlap. This must be checked
112            //*         that no element in this iterator is ever revisited during iteration.
113            unsafe { &mut *(value as *mut T) }
114        })
115    }
116}
117
118impl<'a, T> Iterator for IterMut<'a, T>
119where
120    T: BorshSerialize + BorshDeserialize,
121{
122    type Item = &'a mut T;
123
124    fn next(&mut self) -> Option<Self::Item> {
125        <Self as Iterator>::nth(self, 0)
126    }
127
128    fn size_hint(&self) -> (usize, Option<usize>) {
129        let remaining = self.remaining();
130        (remaining, Some(remaining))
131    }
132
133    fn count(self) -> usize {
134        self.remaining()
135    }
136
137    fn nth(&mut self, n: usize) -> Option<Self::Item> {
138        let idx = self.range.nth(n)?;
139        Some(self.get_mut(idx).unwrap_or_else(|| env::panic_str(ERR_INDEX_OUT_OF_BOUNDS)))
140    }
141}
142
143impl<'a, T> ExactSizeIterator for IterMut<'a, T> where T: BorshSerialize + BorshDeserialize {}
144impl<'a, T> FusedIterator for IterMut<'a, T> where T: BorshSerialize + BorshDeserialize {}
145
146impl<'a, T> DoubleEndedIterator for IterMut<'a, T>
147where
148    T: BorshSerialize + BorshDeserialize,
149{
150    fn next_back(&mut self) -> Option<Self::Item> {
151        <Self as DoubleEndedIterator>::nth_back(self, 0)
152    }
153
154    fn nth_back(&mut self, n: usize) -> Option<Self::Item> {
155        let idx = self.range.nth_back(n)?;
156        Some(self.get_mut(idx).unwrap_or_else(|| env::panic_str(ERR_INDEX_OUT_OF_BOUNDS)))
157    }
158}
159
160/// A draining iterator for [`Vector<T>`].
161#[derive(Debug)]
162pub struct Drain<'a, T>
163where
164    T: BorshSerialize + BorshDeserialize,
165{
166    /// Mutable reference to vector used to iterate through.
167    vec: &'a mut Vector<T>,
168    /// Range of indices to iterate.
169    range: Range<u32>,
170    /// Range of elements to delete.
171    delete_range: Range<u32>,
172}
173
174impl<'a, T> Drain<'a, T>
175where
176    T: BorshSerialize + BorshDeserialize,
177{
178    /// Creates a new iterator for the given storage vector.
179    pub(crate) fn new(vec: &'a mut Vector<T>, range: Range<u32>) -> Self {
180        Self { vec, delete_range: range.clone(), range }
181    }
182
183    /// Returns the amount of remaining elements to yield by the iterator.
184    pub(crate) fn remaining(&self) -> usize {
185        self.range.len()
186    }
187    fn remove(&mut self, index: u32) -> T {
188        self.vec
189            .values
190            .get_mut_inner(index)
191            .replace(None)
192            // Element must exist within bounds of vector
193            .unwrap_or_else(|| env::abort())
194    }
195}
196
197impl<'a, T> Drop for Drain<'a, T>
198where
199    T: BorshSerialize + BorshDeserialize,
200{
201    fn drop(&mut self) {
202        let delete_indices = (self.delete_range.start..self.range.start)
203            .chain(self.range.end..self.delete_range.end);
204
205        // Delete any non-deleted elements from iterator (not loading from storage)
206        for i in delete_indices {
207            self.vec.values.set(i, None);
208        }
209
210        // Shift values after delete into slots deleted.
211        let shift_len = self.delete_range.len() as u32;
212        for i in self.delete_range.end..self.vec.len() {
213            self.vec.swap(i, i - shift_len);
214        }
215
216        // Adjust length of vector.
217        self.vec.len -= self.delete_range.len() as u32;
218    }
219}
220
221impl<'a, T> Iterator for Drain<'a, T>
222where
223    T: BorshSerialize + BorshDeserialize,
224{
225    type Item = T;
226
227    fn next(&mut self) -> Option<Self::Item> {
228        // Load and replace value at next index
229        let delete_idx = self.range.next()?;
230        let prev = self.remove(delete_idx);
231
232        Some(prev)
233    }
234
235    fn nth(&mut self, n: usize) -> Option<Self::Item> {
236        for _ in 0..n {
237            let next = self.range.next()?;
238            // Delete all values in advance, values will be shifted over on drop.
239            // This avoids having to load and deserialize any elements skipped over.
240            self.vec.values.set(next, None);
241        }
242        self.next()
243    }
244
245    fn size_hint(&self) -> (usize, Option<usize>) {
246        let remaining = self.remaining();
247        (remaining, Some(remaining))
248    }
249
250    fn count(self) -> usize {
251        self.remaining()
252    }
253}
254
255impl<'a, T> ExactSizeIterator for Drain<'a, T> where T: BorshSerialize + BorshDeserialize {}
256impl<'a, T> FusedIterator for Drain<'a, T> where T: BorshSerialize + BorshDeserialize {}
257
258impl<'a, T> DoubleEndedIterator for Drain<'a, T>
259where
260    T: BorshSerialize + BorshDeserialize,
261{
262    fn next_back(&mut self) -> Option<Self::Item> {
263        let delete_idx = self.range.next_back()?;
264        let prev = self.remove(delete_idx);
265
266        Some(prev)
267    }
268
269    fn nth_back(&mut self, n: usize) -> Option<Self::Item> {
270        // Only delete and don't load any values before n
271        for _ in 0..n {
272            let next = self.range.next_back()?;
273            // Delete all values in advance, values will be shifted over on drop.
274            // This avoids having to load and deserialize any elements skipped over.
275            self.vec.values.set(next, None);
276        }
277        self.next_back()
278    }
279}
unc_sdk/store/vec/iter.rs

unc_sdk/store/vec/
iter.rs
