dashmap/
iter.rs

1use super::mapref::multiple::{RefMulti, RefMutMulti};
2use super::util;
3use crate::t::Map;
4use crate::util::SharedValue;
5use crate::{DashMap, HashMap};
6use core::hash::{BuildHasher, Hash};
7use core::mem;
8use parking_lot::{RwLockReadGuard, RwLockWriteGuard};
9use std::collections::hash_map;
10use std::collections::hash_map::RandomState;
11use std::sync::Arc;
12
13/// Iterator over a DashMap yielding key value pairs.
14///
15/// # Examples
16///
17/// ```
18/// use dashmap::DashMap;
19///
20/// let map = DashMap::new();
21/// map.insert("hello", "world");
22/// map.insert("alex", "steve");
23/// let pairs: Vec<(&'static str, &'static str)> = map.into_iter().collect();
24/// assert_eq!(pairs.len(), 2);
25/// ```
26pub struct OwningIter<K, V, S = RandomState> {
27    map: DashMap<K, V, S>,
28    shard_i: usize,
29    current: Option<GuardOwningIter<K, V>>,
30}
31
32impl<K: Eq + Hash, V, S: BuildHasher + Clone> OwningIter<K, V, S> {
33    pub(crate) fn new(map: DashMap<K, V, S>) -> Self {
34        Self {
35            map,
36            shard_i: 0,
37            current: None,
38        }
39    }
40}
41
42type GuardOwningIter<K, V> = hash_map::IntoIter<K, SharedValue<V>>;
43
44impl<K: Eq + Hash, V, S: BuildHasher + Clone> Iterator for OwningIter<K, V, S> {
45    type Item = (K, V);
46
47    fn next(&mut self) -> Option<Self::Item> {
48        loop {
49            if let Some(current) = self.current.as_mut() {
50                if let Some((k, v)) = current.next() {
51                    return Some((k, v.into_inner()));
52                }
53            }
54
55            if self.shard_i == self.map._shard_count() {
56                return None;
57            }
58
59            //let guard = unsafe { self.map._yield_read_shard(self.shard_i) };
60            let mut shard_wl = unsafe { self.map._yield_write_shard(self.shard_i) };
61
62            let hasher = self.map._hasher();
63
64            let map = mem::replace(&mut *shard_wl, HashMap::with_hasher(hasher));
65
66            drop(shard_wl);
67
68            let iter = map.into_iter();
69
70            //unsafe { ptr::write(&mut self.current, Some((arcee, iter))); }
71            self.current = Some(iter);
72
73            self.shard_i += 1;
74        }
75    }
76}
77
78unsafe impl<K, V, S> Send for OwningIter<K, V, S>
79where
80    K: Eq + Hash + Send,
81    V: Send,
82    S: BuildHasher + Clone + Send,
83{
84}
85
86unsafe impl<K, V, S> Sync for OwningIter<K, V, S>
87where
88    K: Eq + Hash + Sync,
89    V: Sync,
90    S: BuildHasher + Clone + Sync,
91{
92}
93
94type GuardIter<'a, K, V, S> = (
95    Arc<RwLockReadGuard<'a, HashMap<K, V, S>>>,
96    hash_map::Iter<'a, K, SharedValue<V>>,
97);
98
99type GuardIterMut<'a, K, V, S> = (
100    Arc<RwLockWriteGuard<'a, HashMap<K, V, S>>>,
101    hash_map::IterMut<'a, K, SharedValue<V>>,
102);
103
104/// Iterator over a DashMap yielding immutable references.
105///
106/// # Examples
107///
108/// ```
109/// use dashmap::DashMap;
110///
111/// let map = DashMap::new();
112/// map.insert("hello", "world");
113/// assert_eq!(map.iter().count(), 1);
114/// ```
115pub struct Iter<'a, K, V, S = RandomState, M = DashMap<K, V, S>> {
116    map: &'a M,
117    shard_i: usize,
118    current: Option<GuardIter<'a, K, V, S>>,
119}
120
121impl<'a, 'i, K: Clone + Hash + Eq, V: Clone, S: Clone + BuildHasher> Clone for Iter<'i, K, V, S> {
122    fn clone(&self) -> Self {
123        Iter::new(self.map)
124    }
125}
126
127unsafe impl<'a, 'i, K, V, S, M> Send for Iter<'i, K, V, S, M>
128where
129    K: 'a + Eq + Hash + Send,
130    V: 'a + Send,
131    S: 'a + BuildHasher + Clone,
132    M: Map<'a, K, V, S>,
133{
134}
135
136unsafe impl<'a, 'i, K, V, S, M> Sync for Iter<'i, K, V, S, M>
137where
138    K: 'a + Eq + Hash + Sync,
139    V: 'a + Sync,
140    S: 'a + BuildHasher + Clone,
141    M: Map<'a, K, V, S>,
142{
143}
144
145impl<'a, K: Eq + Hash, V, S: 'a + BuildHasher + Clone, M: Map<'a, K, V, S>> Iter<'a, K, V, S, M> {
146    pub(crate) fn new(map: &'a M) -> Self {
147        Self {
148            map,
149            shard_i: 0,
150            current: None,
151        }
152    }
153}
154
155impl<'a, K: Eq + Hash, V, S: 'a + BuildHasher + Clone, M: Map<'a, K, V, S>> Iterator
156    for Iter<'a, K, V, S, M>
157{
158    type Item = RefMulti<'a, K, V, S>;
159
160    fn next(&mut self) -> Option<Self::Item> {
161        loop {
162            if let Some(current) = self.current.as_mut() {
163                if let Some((k, v)) = current.1.next() {
164                    let guard = current.0.clone();
165
166                    return unsafe { Some(RefMulti::new(guard, k, v.get())) };
167                }
168            }
169
170            if self.shard_i == self.map._shard_count() {
171                return None;
172            }
173
174            let guard = unsafe { self.map._yield_read_shard(self.shard_i) };
175
176            let sref: &HashMap<K, V, S> = unsafe { util::change_lifetime_const(&*guard) };
177
178            let iter = sref.iter();
179
180            self.current = Some((Arc::new(guard), iter));
181
182            self.shard_i += 1;
183        }
184    }
185}
186
187/// Iterator over a DashMap yielding mutable references.
188///
189/// # Examples
190///
191/// ```
192/// use dashmap::DashMap;
193///
194/// let map = DashMap::new();
195/// map.insert("Johnny", 21);
196/// map.iter_mut().for_each(|mut r| *r += 1);
197/// assert_eq!(*map.get("Johnny").unwrap(), 22);
198/// ```
199pub struct IterMut<'a, K, V, S = RandomState, M = DashMap<K, V, S>> {
200    map: &'a M,
201    shard_i: usize,
202    current: Option<GuardIterMut<'a, K, V, S>>,
203}
204
205unsafe impl<'a, 'i, K, V, S, M> Send for IterMut<'i, K, V, S, M>
206where
207    K: 'a + Eq + Hash + Send,
208    V: 'a + Send,
209    S: 'a + BuildHasher + Clone,
210    M: Map<'a, K, V, S>,
211{
212}
213
214unsafe impl<'a, 'i, K, V, S, M> Sync for IterMut<'i, K, V, S, M>
215where
216    K: 'a + Eq + Hash + Sync,
217    V: 'a + Sync,
218    S: 'a + BuildHasher + Clone,
219    M: Map<'a, K, V, S>,
220{
221}
222
223impl<'a, K: Eq + Hash, V, S: 'a + BuildHasher + Clone, M: Map<'a, K, V, S>>
224    IterMut<'a, K, V, S, M>
225{
226    pub(crate) fn new(map: &'a M) -> Self {
227        Self {
228            map,
229            shard_i: 0,
230            current: None,
231        }
232    }
233}
234
235impl<'a, K: Eq + Hash, V, S: 'a + BuildHasher + Clone, M: Map<'a, K, V, S>> Iterator
236    for IterMut<'a, K, V, S, M>
237{
238    type Item = RefMutMulti<'a, K, V, S>;
239
240    fn next(&mut self) -> Option<Self::Item> {
241        loop {
242            if let Some(current) = self.current.as_mut() {
243                if let Some((k, v)) = current.1.next() {
244                    let guard = current.0.clone();
245
246                    unsafe {
247                        let k = util::change_lifetime_const(k);
248
249                        let v = &mut *v.as_ptr();
250
251                        return Some(RefMutMulti::new(guard, k, v));
252                    }
253                }
254            }
255
256            if self.shard_i == self.map._shard_count() {
257                return None;
258            }
259
260            let mut guard = unsafe { self.map._yield_write_shard(self.shard_i) };
261
262            let sref: &mut HashMap<K, V, S> = unsafe { util::change_lifetime_mut(&mut *guard) };
263
264            let iter = sref.iter_mut();
265
266            self.current = Some((Arc::new(guard), iter));
267
268            self.shard_i += 1;
269        }
270    }
271}
272
273#[cfg(test)]
274mod tests {
275    use crate::DashMap;
276
277    #[test]
278    fn iter_mut_manual_count() {
279        let map = DashMap::new();
280
281        map.insert("Johnny", 21);
282
283        assert_eq!(map.len(), 1);
284
285        let mut c = 0;
286
287        for shard in map.shards() {
288            c += shard.write().iter_mut().count();
289        }
290
291        assert_eq!(c, 1);
292    }
293
294    #[test]
295    fn iter_mut_count() {
296        let map = DashMap::new();
297
298        map.insert("Johnny", 21);
299
300        assert_eq!(map.len(), 1);
301
302        assert_eq!(map.iter_mut().count(), 1);
303    }
304
305    #[test]
306    fn iter_count() {
307        let map = DashMap::new();
308
309        map.insert("Johnny", 21);
310
311        assert_eq!(map.len(), 1);
312
313        assert_eq!(map.iter().count(), 1);
314    }
315}