dashmap/
iter.rs

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