stable_map/
drain.rs

1#[cfg(test)]
2mod tests;
3
4use {
5    crate::{
6        linear_storage::LinearStorage,
7        pos_vec::pos::{InUse, Pos},
8    },
9    core::{
10        fmt::{Debug, Formatter},
11        iter::FusedIterator,
12    },
13    hashbrown::hash_map,
14};
15
16/// A draining iterator over the entries of a `StableMap` in arbitrary
17/// order. The iterator element type is `(K, V)`.
18///
19/// This `struct` is created by the [`drain`] method on [`StableMap`]. See its
20/// documentation for more.
21///
22/// [`drain`]: crate::StableMap::drain
23/// [`StableMap`]: crate::StableMap
24///
25/// # Examples
26///
27/// ```
28/// use stable_map::StableMap;
29///
30/// let mut map: StableMap<_, _> = [(1, "a"), (2, "b"), (3, "c")].into();
31///
32/// let mut drain_iter = map.drain();
33/// let mut vec = vec![drain_iter.next(), drain_iter.next(), drain_iter.next()];
34///
35/// // The `Drain` iterator produces items in arbitrary order, so the
36/// // items must be sorted to test them against a sorted array.
37/// vec.sort_unstable();
38/// assert_eq!(vec, [Some((1, "a")), Some((2, "b")), Some((3, "c"))]);
39///
40/// // It is fused iterator
41/// assert_eq!(drain_iter.next(), None);
42/// assert_eq!(drain_iter.next(), None);
43/// ```
44pub struct Drain<'a, K, V> {
45    pub(crate) drain: hash_map::Drain<'a, K, Pos<InUse>>,
46    pub(crate) entries: &'a mut LinearStorage<V>,
47}
48
49impl<K, V> Drop for Drain<'_, K, V> {
50    fn drop(&mut self) {
51        self.entries.clear();
52        // SAFETY(invariants):
53        // - Dropping hash_map::Drain clears key_to_pos.
54    }
55}
56
57impl<K, V> Iterator for Drain<'_, K, V> {
58    type Item = (K, V);
59
60    fn next(&mut self) -> Option<Self::Item> {
61        let (k, pos) = self.drain.next()?;
62        let value = unsafe {
63            // SAFETY: By the invariants, pos is valid.
64            self.entries.take_unchecked(pos)
65        };
66        Some((k, value))
67    }
68
69    fn size_hint(&self) -> (usize, Option<usize>) {
70        self.drain.size_hint()
71    }
72}
73
74impl<K, V> Debug for Drain<'_, K, V> {
75    fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result {
76        f.debug_struct("Drain").finish_non_exhaustive()
77    }
78}
79
80impl<K, V> ExactSizeIterator for Drain<'_, K, V> {
81    fn len(&self) -> usize {
82        self.drain.len()
83    }
84}
85
86impl<K, V> FusedIterator for Drain<'_, K, V> {}
87
88// SAFETY:
89// - This impl is required because Pos<InUse>, Pos<Stored> allow for conflicting access
90//   but this API prevents this.
91unsafe impl<K, V> Send for Drain<'_, K, V>
92where
93    K: Send,
94    V: Send,
95{
96}
97
98// SAFETY:
99// - This impl is required because Pos<InUse>, Pos<Stored> allow for conflicting access
100//   but this API prevents this.
101unsafe impl<K, V> Sync for Drain<'_, K, V>
102where
103    K: Sync,
104    V: Sync,
105{
106}
stable_map/drain.rs

stable_map/
drain.rs
