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}