Skip to main content

kevy_map/
iter.rs

1//! Borrowing iterators over a [`KevyMap`] — `(&K, &V)`, `&K`-only, and
2//! `&V`-only flavours.
3
4use core::mem::MaybeUninit;
5
6use crate::map::KevyMap;
7
8/// `(&K, &V)` iterator over all live entries of a [`KevyMap`]; order unspecified.
9///
10/// ```
11/// use kevy_map::KevyMap;
12/// let mut m: KevyMap<u64, &str> = KevyMap::new();
13/// m.insert(1, "a");
14/// m.insert(2, "b");
15/// let mut pairs: Vec<(u64, &str)> = m.iter().map(|(k, v)| (*k, *v)).collect();
16/// pairs.sort();
17/// assert_eq!(pairs, [(1, "a"), (2, "b")]);
18/// ```
19#[derive(Debug)]
20pub struct Iter<'a, K, V> {
21    metadata: &'a [u8],
22    slots: &'a [MaybeUninit<(K, V)>],
23    pos: usize,
24}
25
26impl<'a, K, V> Iter<'a, K, V> {
27    /// Construct an iterator from a map's raw bucket slices.
28    ///
29    /// `metadata` may be longer than `slots` because the map keeps a
30    /// trailing `GROUP_WIDTH - 1` byte mirror for SIMD-safe wraparound
31    /// loads; only the first `slots.len()` metadata bytes correspond to
32    /// real slots, and that's what we iterate over.
33    pub(crate) fn new(metadata: &'a [u8], slots: &'a [MaybeUninit<(K, V)>]) -> Self {
34        let real_len = slots.len();
35        let metadata = &metadata[..real_len];
36        Self { metadata, slots, pos: 0 }
37    }
38
39    /// Construct an iterator that starts at bucket `start` (clamped to
40    /// `slots.len()`). Powers reservoir / random-start sampling for the
41    /// `kevy-store` eviction sampler — chain two of these (start..end, then
42    /// 0..start) to walk the table in a ring beginning at any position.
43    pub(crate) fn with_start(
44        metadata: &'a [u8],
45        slots: &'a [MaybeUninit<(K, V)>],
46        start: usize,
47    ) -> Self {
48        let real_len = slots.len();
49        let metadata = &metadata[..real_len];
50        Self { metadata, slots, pos: start.min(real_len) }
51    }
52}
53
54impl<'a, K, V> Iterator for Iter<'a, K, V> {
55    type Item = (&'a K, &'a V);
56    fn next(&mut self) -> Option<Self::Item> {
57        while self.pos < self.metadata.len() {
58            let i = self.pos;
59            self.pos += 1;
60            if self.metadata[i] & 0x80 == 0 {
61                // SAFETY: full slot. The borrow's lifetime is tied to
62                // self.slots: &'a [MaybeUninit<(K, V)>].
63                let kv = unsafe { self.slots[i].assume_init_ref() };
64                return Some((&kv.0, &kv.1));
65            }
66        }
67        None
68    }
69}
70
71impl<'a, K, V> IntoIterator for &'a KevyMap<K, V> {
72    type Item = (&'a K, &'a V);
73    type IntoIter = Iter<'a, K, V>;
74    fn into_iter(self) -> Self::IntoIter {
75        self.iter()
76    }
77}
78
79/// `(&K, &mut V)` iterator over all live entries of a [`KevyMap`]; order
80/// unspecified. Keys stay shared — mutating a key would corrupt its bucket.
81///
82/// ```
83/// use kevy_map::KevyMap;
84/// let mut m: KevyMap<u64, u64> = KevyMap::new();
85/// m.insert(1, 10);
86/// m.insert(2, 20);
87/// for (_, v) in m.iter_mut() {
88///     *v += 1;
89/// }
90/// assert_eq!((m.get(&1), m.get(&2)), (Some(&11), Some(&21)));
91/// ```
92#[derive(Debug)]
93pub struct IterMut<'a, K, V> {
94    metadata: &'a [u8],
95    slots: &'a mut [MaybeUninit<(K, V)>],
96    pos: usize,
97}
98
99impl<'a, K, V> IterMut<'a, K, V> {
100    /// Construct from a map's raw bucket slices (same mirror-tail trim as
101    /// [`Iter::new`]).
102    pub(crate) fn new(metadata: &'a [u8], slots: &'a mut [MaybeUninit<(K, V)>]) -> Self {
103        let metadata = &metadata[..slots.len()];
104        Self { metadata, slots, pos: 0 }
105    }
106}
107
108impl<'a, K, V> Iterator for IterMut<'a, K, V> {
109    type Item = (&'a K, &'a mut V);
110    fn next(&mut self) -> Option<Self::Item> {
111        while self.pos < self.metadata.len() {
112            let i = self.pos;
113            self.pos += 1;
114            if self.metadata[i] & 0x80 == 0 {
115                // SAFETY: full slot ⇒ initialised, and `pos` only advances, so
116                // each index is yielded at most once — the returned `&mut V`s
117                // are disjoint and all live within the `'a` borrow of `slots`.
118                let kv = unsafe { &mut *self.slots.as_mut_ptr().add(i).cast::<(K, V)>() };
119                return Some((&kv.0, &mut kv.1));
120            }
121        }
122        None
123    }
124}
125
126impl<'a, K, V> IntoIterator for &'a mut KevyMap<K, V> {
127    type Item = (&'a K, &'a mut V);
128    type IntoIter = IterMut<'a, K, V>;
129    fn into_iter(self) -> Self::IntoIter {
130        self.iter_mut()
131    }
132}
133
134/// `&K` iterator over all live entries of a [`KevyMap`].
135///
136/// ```
137/// use kevy_map::KevyMap;
138/// let mut m: KevyMap<u64, ()> = KevyMap::new();
139/// m.insert(3, ());
140/// m.insert(1, ());
141/// let mut keys: Vec<u64> = m.keys().copied().collect();
142/// keys.sort();
143/// assert_eq!(keys, [1, 3]);
144/// ```
145#[derive(Debug)]
146pub struct Keys<'a, K, V>(Iter<'a, K, V>);
147
148impl<'a, K, V> Keys<'a, K, V> {
149    pub(crate) fn new(inner: Iter<'a, K, V>) -> Self {
150        Self(inner)
151    }
152}
153
154impl<'a, K, V> Iterator for Keys<'a, K, V> {
155    type Item = &'a K;
156    fn next(&mut self) -> Option<Self::Item> {
157        self.0.next().map(|(k, _)| k)
158    }
159}
160
161/// `&V` iterator over all live entries of a [`KevyMap`].
162///
163/// ```
164/// use kevy_map::KevyMap;
165/// let mut m: KevyMap<u64, u64> = KevyMap::new();
166/// m.insert(1, 5);
167/// m.insert(2, 7);
168/// assert_eq!(m.values().sum::<u64>(), 12);
169/// ```
170#[derive(Debug)]
171pub struct Values<'a, K, V>(Iter<'a, K, V>);
172
173impl<'a, K, V> Values<'a, K, V> {
174    pub(crate) fn new(inner: Iter<'a, K, V>) -> Self {
175        Self(inner)
176    }
177}
178
179impl<'a, K, V> Iterator for Values<'a, K, V> {
180    type Item = &'a V;
181    fn next(&mut self) -> Option<Self::Item> {
182        self.0.next().map(|(_, v)| v)
183    }
184}