Skip to main content

arctic/sequential/
iter.rs

1use core::marker::PhantomData;
2use core::ops::ControlFlow;
3use core::ops::Deref;
4
5use crate::raw;
6use crate::raw::Edge;
7use crate::raw::Key;
8use crate::raw::iter::Order;
9use crate::sequential::Value;
10
11/// Immutable reference to a subtree rooted at a key prefix,
12/// optionally bounded by a key range.
13pub struct Shard<'g, 'k, K: Key, V, R> {
14    inner: raw::Shard<'g, 'k, K, R>,
15    _value: PhantomData<&'g V>,
16}
17
18impl<'g, 'k, K, V, R> Shard<'g, 'k, K, V, R>
19where
20    K: Key,
21    V: Value,
22    R: raw::iter::Range<K::Read<'k>>,
23{
24    #[inline]
25    pub(crate) unsafe fn new(prefix: raw::Shard<'g, 'k, K, R>) -> Self {
26        Self {
27            inner: prefix,
28            _value: PhantomData,
29        }
30    }
31
32    /// Get an iterator over keys and immutable references to values in `order`.
33    #[inline]
34    pub fn entries(&self, order: Order) -> EntryIter<'g, 'k, K, V, R> {
35        EntryIter {
36            inner: self.inner.entries(Some(order)),
37            _value: PhantomData,
38        }
39    }
40
41    /// Get an iterator over immutable references to values in `order`.
42    #[inline]
43    pub fn values(&self, order: Order) -> ValueIter<'g, 'k, K, V, R> {
44        ValueIter {
45            inner: self.inner.values(Some(order)),
46            _value: PhantomData,
47        }
48    }
49}
50
51/// Mutable reference to a subtree rooted at a key prefix,
52/// optionally bounded by a key range.
53pub struct ShardMut<'g, 'k, K: Key, V, R>(Shard<'g, 'k, K, V, R>);
54
55impl<'g, 'k, K, V, R> ShardMut<'g, 'k, K, V, R>
56where
57    K: Key,
58    V: Value,
59    R: raw::iter::Range<K::Read<'k>>,
60{
61    #[inline]
62    pub(crate) unsafe fn new(prefix: Shard<'g, 'k, K, V, R>) -> Self {
63        Self(prefix)
64    }
65
66    /// Get an iterator over keys and mutable references to values in `O` order.
67    #[inline]
68    pub fn entries_mut(&mut self, order: Order) -> EntryIterMut<'g, 'k, K, V, R> {
69        EntryIterMut {
70            inner: self.0.inner.entries(Some(order)),
71            _value: PhantomData,
72        }
73    }
74
75    /// Get an iterator over mutable references to values in `O` order.
76    #[inline]
77    pub fn values_mut(&mut self, order: Order) -> ValueIterMut<'g, 'k, K, V, R> {
78        ValueIterMut {
79            inner: self.0.inner.values(Some(order)),
80            _value: PhantomData,
81        }
82    }
83}
84
85impl<'g, 'k, K: Key, V: Value, R: raw::iter::Range<K::Read<'k>>> Deref
86    for ShardMut<'g, 'k, K, V, R>
87{
88    type Target = Shard<'g, 'k, K, V, R>;
89    fn deref(&self) -> &Self::Target {
90        &self.0
91    }
92}
93
94/// Iterator over keys and immutable references to values.
95pub struct EntryIter<'g, 'k, K: Key, V, R: raw::iter::Range<K::Read<'k>>> {
96    inner: raw::iter::EntryIter<'g, 'k, K, R>,
97    _value: PhantomData<&'g V>,
98}
99
100impl<'g, 'k, K, V, R> EntryIter<'g, 'k, K, V, R>
101where
102    K: Key,
103    V: Value,
104    R: raw::iter::Range<K::Read<'k>>,
105{
106    /// Lending equivalent to [`Iterator::next`] that borrows the current key from this [`EntryIter`].
107    #[inline]
108    pub fn lend(&mut self) -> Option<(K::Insert<'_>, &'g V)> {
109        self.inner.lend().map(|(key, _, edge)| {
110            (key, unsafe {
111                Edge::as_value_unchecked(edge).cast::<V>().as_ref()
112            })
113        })
114    }
115
116    /// Internal iteration over keys and immutable references to values.
117    #[inline]
118    pub fn try_fold<F, B, C>(self, init: C, mut apply: F) -> ControlFlow<B, C>
119    where
120        F: FnMut(C, (K::Insert<'_>, &'g V)) -> ControlFlow<B, C>,
121    {
122        self.inner.try_fold(init, |acc, (key, _, edge)| {
123            apply(
124                acc,
125                (key, unsafe {
126                    Edge::as_value_unchecked(edge).cast::<V>().as_ref()
127                }),
128            )
129        })
130    }
131}
132
133impl<'g, 'k, K, V, R> Iterator for EntryIter<'g, 'k, K, V, R>
134where
135    K: Key,
136    V: Value,
137    R: raw::iter::Range<K::Read<'k>>,
138{
139    type Item = (K, &'g V);
140
141    #[inline]
142    fn next(&mut self) -> Option<Self::Item> {
143        self.lend()
144            .map(|(key, value)| (K::insert_to_key(key), value))
145    }
146}
147
148/// Iterator over keys and mutable references to values.
149pub struct EntryIterMut<'g, 'k, K: Key, V, R: raw::iter::Range<K::Read<'k>>> {
150    inner: raw::iter::EntryIter<'g, 'k, K, R>,
151    _value: PhantomData<&'g mut V>,
152}
153
154impl<'g, 'k, K, V, R> EntryIterMut<'g, 'k, K, V, R>
155where
156    K: Key,
157    V: Value,
158    R: raw::iter::Range<K::Read<'k>>,
159{
160    /// Lending equivalent to [`Iterator::next`] that borrows the current key from this [`EntryIter`].
161    #[inline]
162    pub fn lend(&mut self) -> Option<(K::Insert<'_>, &'g mut V)> {
163        self.inner.lend().map(|(key, _, edge)| {
164            (key, unsafe {
165                Edge::as_value_unchecked(edge).cast::<V>().as_mut()
166            })
167        })
168    }
169
170    /// Internal iteration over keys and mutable references to values.
171    #[inline]
172    pub fn try_fold<F, B, C>(self, init: C, mut apply: F) -> ControlFlow<B, C>
173    where
174        F: FnMut(C, (K::Insert<'_>, &'g mut V)) -> ControlFlow<B, C>,
175    {
176        self.inner.try_fold(init, |acc, (key, _, edge)| {
177            apply(
178                acc,
179                (key, unsafe {
180                    Edge::as_value_unchecked(edge).cast::<V>().as_mut()
181                }),
182            )
183        })
184    }
185}
186
187impl<'g, 'k, K, V, R> Iterator for EntryIterMut<'g, 'k, K, V, R>
188where
189    K: Key,
190    V: Value,
191    R: raw::iter::Range<K::Read<'k>>,
192{
193    type Item = (K, &'g mut V);
194
195    #[inline]
196    fn next(&mut self) -> Option<Self::Item> {
197        self.lend()
198            .map(|(key, value)| (K::insert_to_key(key), value))
199    }
200}
201
202/// Iterator over references to values.
203pub struct ValueIter<'g, 'k, K: Key, V, R: raw::iter::Range<K::Read<'k>>> {
204    inner: raw::iter::ValueIter<'g, 'k, K, R>,
205    _value: PhantomData<&'g V>,
206}
207
208impl<'g, 'k, K, V, R> ValueIter<'g, 'k, K, V, R>
209where
210    K: Key,
211    V: Value,
212    R: raw::iter::Range<K::Read<'k>>,
213{
214    /// Internal iteration over immutable references to values.
215    #[inline]
216    pub fn try_fold<F, B, C>(self, init: C, mut apply: F) -> ControlFlow<B, C>
217    where
218        F: FnMut(C, &'g V) -> ControlFlow<B, C>,
219    {
220        self.inner.try_fold(init, |acc, (_, edge)| {
221            apply(acc, unsafe {
222                Edge::as_value_unchecked(edge).cast::<V>().as_ref()
223            })
224        })
225    }
226}
227
228impl<'g, 'k, K, V, R> Iterator for ValueIter<'g, 'k, K, V, R>
229where
230    K: Key,
231    V: Value,
232    R: crate::raw::iter::Range<K::Read<'k>>,
233{
234    type Item = &'g V;
235
236    #[inline]
237    fn next(&mut self) -> Option<Self::Item> {
238        self.inner
239            .lend()
240            .map(|(_, edge)| unsafe { Edge::as_value_unchecked(edge).cast::<V>().as_ref() })
241    }
242}
243
244/// Iterator over mutable references to values.
245pub struct ValueIterMut<'g, 'k, K: Key, V, R: raw::iter::Range<K::Read<'k>>> {
246    inner: raw::iter::ValueIter<'g, 'k, K, R>,
247    _value: PhantomData<&'g mut V>,
248}
249
250impl<'g, 'k, K, V, R> ValueIterMut<'g, 'k, K, V, R>
251where
252    K: Key,
253    V: Value,
254    R: raw::iter::Range<K::Read<'k>>,
255{
256    /// Internal iteration over mutable references to values.
257    #[inline]
258    pub fn try_fold<F, B, C>(self, init: C, mut apply: F) -> ControlFlow<B, C>
259    where
260        F: FnMut(C, &'g mut V) -> ControlFlow<B, C>,
261    {
262        self.inner.try_fold(init, |acc, (_, edge)| {
263            apply(acc, unsafe {
264                Edge::as_value_unchecked(edge).cast::<V>().as_mut()
265            })
266        })
267    }
268}
269
270impl<'g, 'k, K, V, R> Iterator for ValueIterMut<'g, 'k, K, V, R>
271where
272    K: Key,
273    V: Value,
274    R: crate::raw::iter::Range<K::Read<'k>>,
275{
276    type Item = &'g mut V;
277
278    #[inline]
279    fn next(&mut self) -> Option<Self::Item> {
280        self.inner
281            .lend()
282            .map(|(_, edge)| unsafe { Edge::as_value_unchecked(edge).cast::<V>().as_mut() })
283    }
284}
285
286#[cfg(test)]
287mod tests {
288    use core::convert::Infallible;
289    use core::ops::ControlFlow;
290
291    use crate::Order;
292    use crate::sequential::Map;
293
294    #[test]
295    fn indirect_values_mut() {
296        let mut map = Map::<u64, _>::default();
297
298        for i in 0..1024 {
299            map.upsert(i, Box::new(i)).unwrap_err();
300        }
301
302        map.all_mut()
303            .values_mut(Order::Ascend)
304            .try_fold((), |(), value| {
305                **value += 1;
306                ControlFlow::<Infallible>::Continue(())
307            });
308
309        map.all()
310            .entries(Order::Descend)
311            .try_fold((), |(), (key, value)| {
312                assert_eq!(key + 1, **value);
313                ControlFlow::<Infallible>::Continue(())
314            });
315    }
316
317    #[test]
318    fn direct_values_mut() {
319        let mut map = Map::<u64, _>::default();
320
321        for i in 0..1024 {
322            map.upsert(i, i).unwrap_err();
323        }
324
325        map.all_mut()
326            .values_mut(Order::Ascend)
327            .try_fold((), |(), value| {
328                *value += 1;
329                ControlFlow::<Infallible>::Continue(())
330            });
331
332        map.all()
333            .entries(Order::Descend)
334            .try_fold((), |(), (key, value)| {
335                assert_eq!(key + 1, *value);
336                ControlFlow::<Infallible>::Continue(())
337            });
338    }
339}