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