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.
9pub struct Iter<'a, K, V> {
10    metadata: &'a [u8],
11    slots: &'a [MaybeUninit<(K, V)>],
12    pos: usize,
13}
14
15impl<'a, K, V> Iter<'a, K, V> {
16    /// Construct an iterator from a map's raw bucket slices.
17    ///
18    /// `metadata` may be longer than `slots` because the map keeps a
19    /// trailing `GROUP_WIDTH - 1` byte mirror for SIMD-safe wraparound
20    /// loads; only the first `slots.len()` metadata bytes correspond to
21    /// real slots, and that's what we iterate over.
22    pub(crate) fn new(metadata: &'a [u8], slots: &'a [MaybeUninit<(K, V)>]) -> Self {
23        let real_len = slots.len();
24        let metadata = &metadata[..real_len];
25        Self { metadata, slots, pos: 0 }
26    }
27
28    /// Construct an iterator that starts at bucket `start` (clamped to
29    /// `slots.len()`). Powers reservoir / random-start sampling for the
30    /// `kevy-store` eviction sampler — chain two of these (start..end, then
31    /// 0..start) to walk the table in a ring beginning at any position.
32    pub(crate) fn with_start(
33        metadata: &'a [u8],
34        slots: &'a [MaybeUninit<(K, V)>],
35        start: usize,
36    ) -> Self {
37        let real_len = slots.len();
38        let metadata = &metadata[..real_len];
39        Self { metadata, slots, pos: start.min(real_len) }
40    }
41}
42
43impl<'a, K, V> Iterator for Iter<'a, K, V> {
44    type Item = (&'a K, &'a V);
45    fn next(&mut self) -> Option<Self::Item> {
46        while self.pos < self.metadata.len() {
47            let i = self.pos;
48            self.pos += 1;
49            if self.metadata[i] & 0x80 == 0 {
50                // SAFETY: full slot. The borrow's lifetime is tied to
51                // self.slots: &'a [MaybeUninit<(K, V)>].
52                let kv = unsafe { self.slots[i].assume_init_ref() };
53                return Some((&kv.0, &kv.1));
54            }
55        }
56        None
57    }
58}
59
60impl<'a, K, V> IntoIterator for &'a KevyMap<K, V> {
61    type Item = (&'a K, &'a V);
62    type IntoIter = Iter<'a, K, V>;
63    fn into_iter(self) -> Self::IntoIter {
64        self.iter()
65    }
66}
67
68/// `(&K, &mut V)` iterator over all live entries of a [`KevyMap`]; order
69/// unspecified. Keys stay shared — mutating a key would corrupt its bucket.
70pub struct IterMut<'a, K, V> {
71    metadata: &'a [u8],
72    slots: &'a mut [MaybeUninit<(K, V)>],
73    pos: usize,
74}
75
76impl<'a, K, V> IterMut<'a, K, V> {
77    /// Construct from a map's raw bucket slices (same mirror-tail trim as
78    /// [`Iter::new`]).
79    pub(crate) fn new(metadata: &'a [u8], slots: &'a mut [MaybeUninit<(K, V)>]) -> Self {
80        let metadata = &metadata[..slots.len()];
81        Self { metadata, slots, pos: 0 }
82    }
83}
84
85impl<'a, K, V> Iterator for IterMut<'a, K, V> {
86    type Item = (&'a K, &'a mut V);
87    fn next(&mut self) -> Option<Self::Item> {
88        while self.pos < self.metadata.len() {
89            let i = self.pos;
90            self.pos += 1;
91            if self.metadata[i] & 0x80 == 0 {
92                // SAFETY: full slot ⇒ initialised, and `pos` only advances, so
93                // each index is yielded at most once — the returned `&mut V`s
94                // are disjoint and all live within the `'a` borrow of `slots`.
95                let kv = unsafe { &mut *self.slots.as_mut_ptr().add(i).cast::<(K, V)>() };
96                return Some((&kv.0, &mut kv.1));
97            }
98        }
99        None
100    }
101}
102
103impl<'a, K, V> IntoIterator for &'a mut KevyMap<K, V> {
104    type Item = (&'a K, &'a mut V);
105    type IntoIter = IterMut<'a, K, V>;
106    fn into_iter(self) -> Self::IntoIter {
107        self.iter_mut()
108    }
109}
110
111/// `&K` iterator over all live entries of a [`KevyMap`].
112pub struct Keys<'a, K, V>(Iter<'a, K, V>);
113
114impl<'a, K, V> Keys<'a, K, V> {
115    pub(crate) fn new(inner: Iter<'a, K, V>) -> Self {
116        Self(inner)
117    }
118}
119
120impl<'a, K, V> Iterator for Keys<'a, K, V> {
121    type Item = &'a K;
122    fn next(&mut self) -> Option<Self::Item> {
123        self.0.next().map(|(k, _)| k)
124    }
125}
126
127/// `&V` iterator over all live entries of a [`KevyMap`].
128pub struct Values<'a, K, V>(Iter<'a, K, V>);
129
130impl<'a, K, V> Values<'a, K, V> {
131    pub(crate) fn new(inner: Iter<'a, K, V>) -> Self {
132        Self(inner)
133    }
134}
135
136impl<'a, K, V> Iterator for Values<'a, K, V> {
137    type Item = &'a V;
138    fn next(&mut self) -> Option<Self::Item> {
139        self.0.next().map(|(_, v)| v)
140    }
141}