Skip to main content

kovan_map/hashmap/
iter.rs

1//! Walking the map: the borrowed walks over a table snapshot and the owned walk that consumes
2//! the map.
3//!
4//! A borrowed walk keeps the table that was current when it was created, to its end: a resize
5//! freezes that table (every link final), so it keeps the entries it held then. It takes each
6//! bucket in one pass, collecting the nodes that are not deleted, before it yields any of them
7//! (a pass that fails a validation starts that bucket over from its head, with nothing yielded
8//! from it yet), so an entry present for the whole walk is yielded exactly once: a single pass
9//! meets each live node once, a replace puts the new node right after the old one (which the
10//! pass either met live, or passes deleted), and a pass never goes back. A key removed and
11//! inserted again during the walk can be yielded once for each of its lives. Modelled in
12//! `proofs/tla/chained/ChainedMap.tla` (actions T0 to T9).
13
14extern crate alloc;
15
16use super::node::{Node, is_marked, ptr};
17use super::table::TableRef;
18use super::walk::still_links;
19use super::{HashMap, MIN_CAPACITY};
20use alloc::vec::Vec;
21use core::hash::{BuildHasher, Hash};
22use core::sync::atomic::Ordering;
23use kovan::{Atomic, drop_node, free_node, pin};
24
25// Construction only: none of `iter`/`keys`/`values` hashes or clones a value (the walk that
26// does lives in `Iterator for Iter` below), so this block needs only the struct's own `'static`
27// bound, as std's `HashMap::iter`/`keys`/`values` need none of `Hash`, `Eq`, `Clone` or
28// `BuildHasher` either.
29impl<K: 'static, V: 'static, S> HashMap<K, V, S> {
30    /// Returns an iterator over the map entries, `(K, V)` clones.
31    ///
32    /// The iterator walks the table that is current when it is created, to its end, even when a
33    /// resize replaces it meanwhile. An entry present from the iterator's creation to its end is
34    /// yielded exactly once, with a value it held during the walk; an entry inserted, removed or
35    /// updated concurrently may or may not be reflected, and a key removed and inserted again
36    /// meanwhile may be yielded once per life. The iterator holds one kovan guard for its
37    /// lifetime (nodes it passed are not freed while it lives).
38    pub fn iter(&self) -> Iter<'_, K, V, S> {
39        let guard = pin();
40        let table = self.table_ref(&guard);
41        Iter {
42            _map: self,
43            table,
44            bucket_idx: 0,
45            batch: Batch::new(),
46            pos: 0,
47            guard,
48        }
49    }
50
51    /// Returns an iterator over the map keys.
52    /// Yields K clones.
53    pub fn keys(&self) -> Keys<'_, K, V, S> {
54        Keys { iter: self.iter() }
55    }
56
57    /// Returns an iterator over the map values (clones `V`).
58    pub fn values(&self) -> Values<'_, K, V, S> {
59        Values { iter: self.iter() }
60    }
61}
62
63/// How many of a bucket's nodes a walk keeps without allocating. A bucket holds under one node
64/// on average (the map grows past three quarters of its buckets), so a longer chain is rare;
65/// one past this spills into a vector the walk reuses for every later bucket.
66const BATCH_INLINE: usize = 8;
67
68/// The nodes one bucket pass collected, in chain order.
69struct Batch<K, V> {
70    inline: [*const Node<K, V>; BATCH_INLINE],
71    spill: Vec<*const Node<K, V>>,
72    len: usize,
73}
74
75impl<K, V> Batch<K, V> {
76    fn new() -> Self {
77        Self {
78            inline: [core::ptr::null(); BATCH_INLINE],
79            spill: Vec::new(),
80            len: 0,
81        }
82    }
83
84    #[inline]
85    fn clear(&mut self) {
86        self.len = 0;
87        self.spill.clear();
88    }
89
90    #[inline]
91    fn push(&mut self, node: *const Node<K, V>) {
92        if self.len < BATCH_INLINE {
93            self.inline[self.len] = node;
94        } else {
95            self.spill.push(node);
96        }
97        self.len += 1;
98    }
99
100    #[inline]
101    fn get(&self, i: usize) -> Option<*const Node<K, V>> {
102        match i {
103            _ if i >= self.len => None,
104            _ if i < BATCH_INLINE => Some(self.inline[i]),
105            _ => Some(self.spill[i - BATCH_INLINE]),
106        }
107    }
108}
109
110/// Iterator over HashMap entries ([`HashMap::iter`]).
111///
112/// Field ordering matters for drop safety.
113/// Rust drops struct fields in declaration order.
114/// The `guard` must be dropped *after* `batch`/`table` so that the epoch
115/// pin covering the snapshot is not released before we're done with the raw
116/// pointers.
117pub struct Iter<'a, K: 'static, V: 'static, S> {
118    _map: &'a HashMap<K, V, S>,
119    table: TableRef<K, V>,
120    bucket_idx: usize,
121    batch: Batch<K, V>,
122    pos: usize,
123    guard: kovan::Guard,
124}
125
126impl<K, V, S> Iter<'_, K, V, S> {
127    /// Collect bucket `b`'s nodes that are not deleted, in one validated pass.
128    fn collect(&mut self, b: usize) {
129        let guard = &self.guard;
130        'bucket: loop {
131            self.batch.clear();
132            let mut prev: &Atomic<Node<K, V>> = self.table.bucket(b);
133            let mut cur = ptr(prev.load(Ordering::Acquire, guard).as_raw());
134            // The first deleted node of the run the pass is in (see `walk::still_links`).
135            let mut run: *mut Node<K, V> = core::ptr::null_mut();
136            while !cur.is_null() {
137                // SAFETY: loaded under the iterator's guard while reachable (see
138                // `hashmap::walk`); the guard lives as long as the iterator.
139                let node = unsafe { &*cur };
140                let next = node.next.load(Ordering::Acquire, guard).as_raw();
141                if !is_marked(next) {
142                    self.batch.push(cur);
143                    prev = &node.next;
144                    run = core::ptr::null_mut();
145                } else {
146                    if run.is_null() {
147                        run = cur;
148                    }
149                    if !still_links(prev, run, guard) {
150                        continue 'bucket;
151                    }
152                }
153                cur = ptr(next);
154            }
155            self.pos = 0;
156            return;
157        }
158    }
159}
160
161impl<'a, K, V, S> Iterator for Iter<'a, K, V, S>
162where
163    K: Clone,
164    V: Clone,
165{
166    type Item = (K, V);
167
168    fn next(&mut self) -> Option<Self::Item> {
169        loop {
170            if let Some(node) = self.batch.get(self.pos) {
171                self.pos += 1;
172                // SAFETY: collected under the iterator's guard, which keeps it from being freed.
173                let node = unsafe { &*node };
174                return Some((node.key.clone(), node.value.clone()));
175            }
176            if self.bucket_idx >= self.table.capacity() {
177                return None;
178            }
179            let b = self.bucket_idx;
180            self.bucket_idx += 1;
181            self.collect(b);
182        }
183    }
184}
185
186/// Iterator over HashMap keys.
187pub struct Keys<'a, K: 'static, V: 'static, S> {
188    iter: Iter<'a, K, V, S>,
189}
190
191impl<'a, K, V, S> Iterator for Keys<'a, K, V, S>
192where
193    K: Clone,
194    V: Clone,
195{
196    type Item = K;
197
198    fn next(&mut self) -> Option<Self::Item> {
199        self.iter.next().map(|(k, _)| k)
200    }
201}
202
203// Bounded by exactly what `Iterator for Iter` needs: a concurrent walk yields owned clones, so
204// `K: Clone` and `V: Clone` are unavoidable here, unlike std's unconstrained `IntoIterator for
205// &HashMap`, which yields borrowed `(&K, &V)` and hashes nothing at this bound-checked level
206// either.
207impl<'a, K, V, S> IntoIterator for &'a HashMap<K, V, S>
208where
209    K: Clone + 'static,
210    V: Clone + 'static,
211{
212    type Item = (K, V);
213    type IntoIter = Iter<'a, K, V, S>;
214
215    fn into_iter(self) -> Self::IntoIter {
216        self.iter()
217    }
218}
219
220/// Iterator over HashMap values (clones `V`).
221pub struct Values<'a, K: 'static, V: 'static, S> {
222    iter: Iter<'a, K, V, S>,
223}
224
225impl<'a, K, V, S> Iterator for Values<'a, K, V, S>
226where
227    K: Clone,
228    V: Clone,
229{
230    type Item = V;
231
232    #[inline]
233    fn next(&mut self) -> Option<V> {
234        self.iter.next().map(|(_, v)| v)
235    }
236}
237
238/// Owned iterator yielding `(K, V)` by value - moves out of the nodes, no
239/// clone. Consuming the map gives exclusive access, so no guard protection of
240/// the yielded values is needed. A deleted node still in a chain (never
241/// unlinked, so never retired) is freed here with its key and value.
242pub struct IntoIter<K: 'static, V: 'static> {
243    table: TableRef<K, V>,
244    bucket_idx: usize,
245    current: *mut Node<K, V>,
246    guard: kovan::Guard,
247}
248
249impl<K, V> Iterator for IntoIter<K, V> {
250    type Item = (K, V);
251
252    fn next(&mut self) -> Option<(K, V)> {
253        loop {
254            if !self.current.is_null() {
255                let node = self.current;
256                let next = unsafe { (*node).next.load(Ordering::Acquire, &self.guard).as_raw() };
257                self.current = ptr(next);
258                if is_marked(next) {
259                    // SAFETY: exclusive access; the table owns a deleted node still in a chain.
260                    unsafe { drop_node(node) };
261                    continue;
262                }
263                // Move K and V out, then free the shell without running drop.
264                let k = unsafe { core::ptr::read(&(*node).key) };
265                let v = unsafe { core::ptr::read(&(*node).value) };
266                unsafe { free_node(node) };
267                return Some((k, v));
268            }
269            if self.bucket_idx >= self.table.capacity() {
270                return None;
271            }
272            let bucket = self.table.bucket(self.bucket_idx);
273            self.bucket_idx += 1;
274            self.current = ptr(bucket.load(Ordering::Acquire, &self.guard).as_raw());
275        }
276    }
277}
278
279impl<K, V> Drop for IntoIter<K, V> {
280    fn drop(&mut self) {
281        while self.next().is_some() {} // drop remaining live K/V + free shells
282        unsafe { self.table.free_array_only() };
283    }
284}
285
286impl<K, V, S> IntoIterator for HashMap<K, V, S>
287where
288    K: 'static,
289    V: 'static,
290{
291    type Item = (K, V);
292    type IntoIter = IntoIter<K, V>;
293
294    fn into_iter(self) -> IntoIter<K, V> {
295        let mut me = core::mem::ManuallyDrop::new(self);
296        let guard = pin();
297        let table = TableRef::<K, V>::from_raw(me.table.load(Ordering::Relaxed, &guard).as_raw());
298        // Suppress HashMap::drop (we own the table now); drop the hasher and the
299        // latch (a no-op but under shuttle) - the rest is a pointer, a usize and a marker.
300        unsafe {
301            core::ptr::drop_in_place(&mut me.hasher);
302            core::ptr::drop_in_place(&mut me.resizing);
303        }
304        IntoIter {
305            table,
306            bucket_idx: 0,
307            current: core::ptr::null_mut(),
308            guard,
309        }
310    }
311}
312
313// `K: Send, V: Send`: an entry an insert replaces is retired, and its destructor may run on
314// another thread.
315impl<K, V, S> core::iter::FromIterator<(K, V)> for HashMap<K, V, S>
316where
317    K: Hash + Eq + Clone + Send + 'static,
318    V: Clone + Send + 'static,
319    S: BuildHasher + Default,
320{
321    fn from_iter<I: IntoIterator<Item = (K, V)>>(iter: I) -> Self {
322        let map = Self::with_capacity_and_hasher(MIN_CAPACITY, S::default());
323        for (k, v) in iter {
324            map.insert(k, v);
325        }
326        map
327    }
328}