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